"""The PTY helper's output credit is opt-in, so mixed versions keep working.

A running daemon keeps its boot-time worker while albedo_pty.py is read from
disk at each terminal start, so after an update an old worker (no credit in
its spec, never grants) can start a new helper. E2E always pairs the current
worker and helper, so only this test sees that pairing.
"""

import json
import os
import select
import struct
import subprocess
import sys
import time
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
HELPER = ROOT / "priv/python/albedo_pty.py"
FLOOD = b"yes PTY_PROTOCOL | head -n 50000; echo DONE_$((1+1))\r"


class Helper:
    def __init__(self, test, **spec):
        self.process = subprocess.Popen(
            [sys.executable, "-u", str(HELPER)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        assert self.process.stdin is not None and self.process.stdout is not None
        self.stdin = self.process.stdin
        self.stdout = self.process.stdout
        test.addCleanup(self.stop)
        self.send(json.dumps({"cwd": "/", "shell": "/bin/sh", **spec}).encode())
        test.assertEqual(self.receive(5)[0], 0)

    def send(self, payload):
        self.stdin.write(struct.pack("!I", len(payload)) + payload)
        self.stdin.flush()

    def read(self, size, deadline):
        data = b""
        while len(data) < size:
            left = deadline - time.monotonic()
            if left <= 0 or not select.select([self.stdout], [], [], left)[0]:
                return None
            chunk = os.read(self.stdout.fileno(), size - len(data))
            if not chunk:
                return None
            data += chunk
        return data

    def receive(self, seconds):
        """The next frame, or None when none arrives in time."""
        deadline = time.monotonic() + seconds
        header = self.read(4, deadline)
        return header and self.read(struct.unpack("!I", header)[0], deadline)

    def output_frames(self, seconds):
        frames = []
        while (frame := self.receive(seconds)) is not None:
            if frame[0] == 1:
                frames.append(frame[1:])
        return frames

    def stop(self):
        self.process.kill()
        self.process.wait(5)
        self.stdin.close()
        self.stdout.close()


class PtyProtocolTest(unittest.TestCase):
    def test_a_spec_without_credit_never_holds_output_back(self):
        helper = Helper(self)
        helper.send(b"\x01" + FLOOD)
        output = b""
        deadline = time.monotonic() + 20
        while b"DONE_2" not in output and time.monotonic() < deadline:
            frame = helper.receive(5)
            self.assertIsNotNone(frame, "output stopped without credit")
            if frame[0] == 1:
                output += frame[1:]
        self.assertIn(b"DONE_2", output)

    def test_a_spec_with_credit_sends_only_what_was_granted(self):
        helper = Helper(self, credit=2)
        helper.send(b"\x01" + FLOOD)
        self.assertEqual(len(helper.output_frames(1)), 2)
        helper.send(b"\x05" + struct.pack("!H", 3))
        self.assertEqual(len(helper.output_frames(1)), 3)


if __name__ == "__main__":
    unittest.main()