#!/usr/bin/env python3
"""Dokki Terminal's PTY helper for Linux hosts (the twin of `DokkiHelper pty-spawn`).

    dokki-pty pty-spawn --cwd <abs dir> --cols N --rows N [--origin cli] -- <absolute argv0> [args...]

Runs one allowlisted program (claude, codex, opencode, agy, ssh) in a new
pseudo-terminal, directly with execve (never through a shell: when it exits
the terminal ends, there is no prompt to fall back to), and speaks the frame
protocol of src/term/frame.js on stdin/stdout:

    type u8 | length u32 big-endian | payload (at most 64 KiB)
    1 data, 2 resize (cols u16, rows u16), 3 exit (code i32),
    4 signal (u8: 1 HUP, 2 INT, 15 TERM), 5 ready (pid u32, always first)

Trust boundary: the host resolves the tool to an absolute path; this script
checks everything again and refuses with exit 2 before anything is spawned:
- the program: an absolute clean path whose name is in the closed tool list;
  it and its resolved target live outside temporary folders and under no
  world-writable folder; it is exec'd by its resolved path;
- the folder, resolved: home itself (ssh only) or a folder the launcher would offer
  (src/dev/launcher.js isAllowedProject: no ~/.ssh, ~/Library, .git, ...);
- arguments: an agent's preset flags from its closed list (src/term/shell.js
  AGENT_FLAGS, none included) or exactly its resume pair ("Continue here":
  `--resume <uuid>`, `resume <uuid>`, `-s <ses_id>`), exactly
  `[-p N] -- <target>` for ssh; sizes in range;
- `--origin cli` (`dokki claude` typed in a terminal on this computer): only for
  claude, whose arguments are then the user's own, passed through as typed
  (still at most MAX_ARGS of at most MAX_ARG_LEN, no NUL); presets never set it.
  The caller owns the PTY's stdin and claude's list already has its bypass flag,
  so a closed list does not narrow what an interactive claude can do; the
  program, its place and the folder are still checked;
- a shell (a "shell" preset): bash, zsh, fish, sh, dash, ksh, tcsh or tmux
  given from SHELL_DIRS, resolving to one of those names, with exactly `-l`
  (tmux: `attach-session` or `new-session`), in home or an allowed folder.
The child gets a minimal environment (no host secrets, no DOKKI_*), no
inherited descriptor past stdio, default signals. Stdin is always read (so a
signal or EOF is never stuck behind a busy child or a slow host); stdout is
queued and the terminal is not read while the host is behind. Stdin EOF hangs
the child's group up; one still there 2 s later is killed, and so is anything
left in the group after the program exits. Exit 0 after the exit frame.

python3 only (the stdlib's pty, selectors, fcntl); unit tests in
native/linux/tests/test_dokki_pty.py, end to end in scripts/verify-pty.mjs.
"""

import errno
import fcntl
import os
import re
import selectors
import signal
import struct
import sys
import termios
import time

T_DATA, T_RESIZE, T_EXIT, T_SIGNAL, T_READY = 1, 2, 3, 4, 5
KNOWN = {T_DATA, T_RESIZE, T_EXIT, T_SIGNAL, T_READY}
FIXED = {T_RESIZE: 4, T_EXIT: 4, T_SIGNAL: 1, T_READY: 4}
HEADER = 5
MAX_PAYLOAD = 64 * 1024
MAX_COLS, MAX_ROWS = 1000, 500
SIGNALS = {1: signal.SIGHUP, 2: signal.SIGINT, 15: signal.SIGTERM}

TOOLS = frozenset({"claude", "codex", "opencode", "agy", "ssh"})
TEST_FIXTURE = "dokki-pty-fixture"
TEST_ENV = "DOKKI_PTY_TEST"
MAX_ARGS = 64
MAX_ARG_LEN = 4096
SSH_TARGET = re.compile(r"^(?:[A-Za-z0-9._-]+@)?(?:[A-Za-z0-9._-]+|\[[0-9A-Fa-f:.]+\])$")
SIZE = re.compile(r"^[1-9][0-9]{0,3}$")
PORT = re.compile(r"^[1-9][0-9]{0,4}$")
HANGUP_GRACE = 2.0
OUTPUT_DEADLINE = 5.0
WATCHDOG = 8  # seconds after the hang-up or the program's exit: the helper is gone, whatever else
OUTPUT_HIGH_WATER = 256 * 1024
MAX_PENDING_INPUT = 1024 * 1024
AGENTS = frozenset({"claude", "codex", "opencode", "agy"})
# An agent takes no arguments (a preset, src/term/presets.js), or exactly its resume pair
# ("Continue here"). The id shapes are the host's (src/term/resume.js RESUME_ID_RE): keep in step.
UUID = re.compile(r"[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}")
OPENCODE_SESSION = re.compile(r"ses_[A-Za-z0-9]{8,64}")
RESUME_SHAPES = {"claude": ("--resume", UUID), "codex": ("resume", UUID), "opencode": ("-s", OPENCODE_SESSION)}
TEMP_ROOTS = ("/tmp", "/private/tmp", "/var/tmp", "/private/var/tmp", "/var/folders", "/private/var/folders", "/dev/shm", "/run/user")
# A tool name that resolves to a shell or an interpreter would run whatever it is given (or
# open a prompt): refused. Agents that resolve to a script (claude -> cli.js) are fine.
INTERPRETER = re.compile(r"(sh|bash|zsh|fish|dash|ksh|mksh|ash|tcsh|csh|nu|busybox|pwsh|powershell|cmd|osascript|python|perl|ruby|env|node|nodejs|deno|bun|php|lua|tclsh)([0-9][0-9.-]*)?(\.exe)?", re.IGNORECASE)
# ssh must be the system's (or Homebrew's) OpenSSH client.
SYSTEM_SSH = ("/usr/bin/ssh", "/bin/ssh", "/opt/homebrew/bin/ssh", "/usr/local/bin/ssh")
SSH_CELLARS = ("/opt/homebrew/Cellar/openssh/", "/usr/local/Cellar/openssh/")
# Residual risk, by design: whoever can write the folder a tool lives in (~/.local/bin is the
# user's own) can still plant a binary there that is neither a shell nor an interpreter. The
# helper runs as that same user; this narrows what a caller can do with it, it is no sandbox.
# Same lists as src/dev/launcher.js and DokkiCore/DevTerminal.swift.
DENIED_TOP = frozenset({"library", ".ssh", ".gnupg", ".aws", ".azure", ".kube", ".docker", ".trash", ".password-store", "appdata", ".config", ".local", ".cache", ".claude", ".codex", ".gemini", ".npm", ".gcloud"})
DENIED_ANYWHERE = frozenset({".ssh", ".gnupg", ".aws", ".password-store", "keychains", ".git"})


class BadArgs(ValueError):
    pass


class FrameError(ValueError):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code


# ---- frames ----------------------------------------------------------------

def encode(ftype, payload=b""):
    if ftype not in KNOWN:
        raise FrameError("frame-unknown-type", "unknown frame type %d" % ftype)
    if len(payload) > MAX_PAYLOAD:
        raise FrameError("frame-too-large", "frame payload too large")
    if ftype in FIXED and len(payload) != FIXED[ftype]:
        raise FrameError("frame-bad-length", "bad payload length")
    return struct.pack(">BI", ftype, len(payload)) + bytes(payload)


def encode_data(data):
    return b"".join(encode(T_DATA, data[i:i + MAX_PAYLOAD]) for i in range(0, len(data), MAX_PAYLOAD)) if data else encode(T_DATA, b"")


def valid_size(cols, rows):
    return 1 <= cols <= MAX_COLS and 1 <= rows <= MAX_ROWS


class Decoder:
    """Streaming: feed() any slice of the stream, get the complete frames as (type, payload).
    Raises FrameError (unknown type, over 64 KiB, bad fixed length, bad value) from the header
    alone and stays failed."""

    def __init__(self):
        self.buf = bytearray()
        self.failed = None

    def feed(self, chunk):
        if self.failed:
            raise self.failed
        self.buf += chunk
        out = []
        try:
            while len(self.buf) >= HEADER:
                ftype, length = struct.unpack_from(">BI", self.buf, 0)
                if ftype not in KNOWN:
                    raise FrameError("frame-unknown-type", "unknown frame type %d" % ftype)
                if length > MAX_PAYLOAD:
                    raise FrameError("frame-too-large", "frame payload too large")
                if ftype in FIXED and length != FIXED[ftype]:
                    raise FrameError("frame-bad-length", "bad payload length")
                if len(self.buf) < HEADER + length:
                    break
                payload = bytes(self.buf[HEADER:HEADER + length])
                del self.buf[:HEADER + length]
                if ftype == T_RESIZE and not valid_size(*struct.unpack(">HH", payload)):
                    raise FrameError("frame-bad-value", "bad size")
                if ftype == T_SIGNAL and payload[0] not in SIGNALS:
                    raise FrameError("frame-bad-value", "bad signal")
                out.append((ftype, payload))
        except FrameError as error:
            self.failed = error
            self.buf = bytearray()
            raise
        return out


# ---- validation -------------------------------------------------------------

def has_control(value):
    return any(ord(c) < 0x20 or 0x7F <= ord(c) <= 0x9F or c in "  " for c in value)


def is_safe_abs_path(path):
    """Absolute, at most 1024 bytes, no control characters, no `.`/`..` segments, no empty name."""
    if not isinstance(path, str) or not path.startswith("/") or len(path) < 2:
        return False
    if len(path.encode("utf-8", "surrogateescape")) > 1024 or has_control(path):
        return False
    parts = path.split("/")[1:]
    if parts and parts[-1] == "":
        parts = parts[:-1]
    return all(p not in ("", ".", "..") for p in parts) and bool(parts)


def is_ssh_target(target):
    return 0 < len(target) <= 253 and not target.startswith("-") and not target.startswith("@") and bool(SSH_TARGET.fullmatch(target))


def is_ssh_args(args):
    """Exactly `-- <target>` or `-p <port> -- <target>`: no option ever reaches ssh from a preset."""
    if len(args) == 2 and args[0] == "--":
        return is_ssh_target(args[1])
    if len(args) == 4 and args[0] == "-p" and args[2] == "--" and PORT.fullmatch(args[1]) and int(args[1]) <= 65535:
        return is_ssh_target(args[3])
    return False


def is_agent_args(tool, args):
    """None, or exactly the tool's resume pair (`claude --resume <uuid>`, `codex resume <uuid>`,
    `opencode -s <ses_id>`); agy has no resume. fullmatch: no trailing newline slips through."""
    if tool not in AGENTS:
        return False
    if is_agent_flags(tool, args):
        return True
    shape = RESUME_SHAPES.get(tool)
    return bool(shape) and len(args) == 2 and args[0] == shape[0] and shape[1].fullmatch(args[1]) is not None


SHELLS = frozenset({"bash", "zsh", "fish", "sh", "dash", "ksh", "tcsh", "tmux"})
SHELL_DIRS = frozenset({"/usr/bin", "/bin", "/usr/local/bin", "/home/linuxbrew/.linuxbrew/bin", "/run/current-system/sw/bin"})
# Where a shell given from SHELL_DIRS may resolve to besides them (Homebrew's and Nix's stores).
SHELL_STORES = ("/home/linuxbrew/.linuxbrew/Cellar/", "/opt/homebrew/Cellar/", "/usr/local/Cellar/", "/nix/store/")
MODEL = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:/@\[\]-]{0,99}")
# Each tool's preset options (src/term/shell.js FLAG_OPTION): a tuple of values, () for a
# switch, None for a model id. Keep in step with DokkiCore/Pty.swift and native/win.
AGENT_OPTIONS = {
    "claude": {"--model": None, "--permission-mode": ("plan", "acceptEdits", "default"), "--continue": (), "--dangerously-skip-permissions": ()},
    "codex": {"--model": None, "--sandbox": ("read-only", "workspace-write"), "--ask-for-approval": ("untrusted", "on-failure", "on-request", "never"),
              "--dangerously-bypass-approvals-and-sandbox": ()},
    "opencode": {"--model": None},
    "agy": {},
}


def is_agent_flags(tool, args):
    """Options from the tool's own list, each at most once, with a valid value where one is taken."""
    options = AGENT_OPTIONS.get(tool)
    if options is None or len(args) > 12:
        return False
    seen, i = set(), 0
    while i < len(args):
        option = args[i]
        if option not in options or option in seen:
            return False
        seen.add(option)
        i += 1
        values = options[option]
        if values == ():
            continue
        if i >= len(args):
            return False
        value = args[i]
        if values is None:
            if not MODEL.fullmatch(value):
                return False
        elif value not in values:
            return False
        i += 1
    # A bypass replaces the permission mode, the sandbox and the approvals (src/term/shell.js).
    bypass = "--dangerously-skip-permissions" in seen or "--dangerously-bypass-approvals-and-sandbox" in seen
    return not (bypass and seen & {"--permission-mode", "--sandbox", "--ask-for-approval"})


def is_shell_args(tool, args):
    """tmux: exactly `attach-session` or `new-session`; any other shell exactly `-l`."""
    if tool == "tmux":
        return args in (["attach-session"], ["new-session"])
    return args == ["-l"]


def allowed_tools(environ):
    return TOOLS | {TEST_FIXTURE} if environ.get(TEST_ENV) == "1" else TOOLS


class HostFS:
    """The file system questions validation asks (a fake one in the unit tests)."""

    def realpath(self, path):
        return os.path.realpath(path)

    def is_dir(self, path):
        return os.path.isdir(path)

    def is_exec_file(self, path):
        return os.path.isfile(path) and os.access(path, os.X_OK)

    def world_writable(self, path):
        try:
            return bool(os.stat(path).st_mode & 0o002)
        except OSError:
            return True

    def home(self, environ):
        return environ.get("HOME") or os.path.expanduser("~")


def under(path, root):
    return path == root or path.startswith(root.rstrip("/") + "/")


def is_allowed_folder(path, home):
    """Inside home (never home itself), not under a denied folder, no `.`/`..` segments."""
    base = home.rstrip("/")
    if len(base) < 2 or not path.lower().startswith(base.lower() + "/"):
        return False
    parts = [p.lower() for p in path[len(base) + 1:].split("/")]
    if not parts[0] or any(p in (".", "..") for p in parts):
        return False
    return parts[0] not in DENIED_TOP and not any(p in DENIED_ANYWHERE for p in parts)


def is_trusted_program(given, real, environ, fs):
    """Neither the path given nor its target in a temporary folder or under a world-writable one."""
    temp = list(TEMP_ROOTS)
    if environ.get("TMPDIR"):
        temp.append(fs.realpath(environ["TMPDIR"]))
    for path in (given, real):
        folder = os.path.dirname(path)
        if any(under(folder, root) for root in temp):
            return False
        while True:
            if fs.world_writable(folder):
                return False
            if folder == "/":
                break
            folder = os.path.dirname(folder)
    return True


def parse_spawn(args, environ, fs=None):
    """`--cwd <dir> --cols N --rows N -- <argv0> [args...]` (flags once each, any order) →
    {cwd (resolved), cols, rows, program (resolved), argv}, or BadArgs."""
    fs = fs or HostFS()
    if "--" not in args:
        raise BadArgs("missing --")
    split = args.index("--")
    flags, argv = args[:split], args[split + 1:]
    if len(flags) not in (6, 8):
        raise BadArgs("need --cwd, --cols, --rows")
    seen = {}
    for key, value in zip(flags[0::2], flags[1::2]):
        if key not in ("--cwd", "--cols", "--rows", "--origin") or key in seen:
            raise BadArgs("bad flag")
        seen[key] = value
    if not all(k in seen for k in ("--cwd", "--cols", "--rows")):
        raise BadArgs("need --cwd, --cols, --rows")
    if "--origin" in seen and seen["--origin"] != "cli":
        raise BadArgs("bad origin")
    from_cli = seen.get("--origin") == "cli"
    cwd, cols, rows = seen["--cwd"], seen["--cols"], seen["--rows"]
    if not SIZE.fullmatch(cols) or not SIZE.fullmatch(rows) or not valid_size(int(cols), int(rows)):
        raise BadArgs("bad size")
    if not argv or not is_safe_abs_path(argv[0]):
        raise BadArgs("argv0 must be an absolute path")
    tool = argv[0].rsplit("/", 1)[-1]
    rest = argv[1:]
    if len(rest) > MAX_ARGS or any("\0" in a or len(a) > MAX_ARG_LEN for a in rest):
        raise BadArgs("bad arguments")
    if from_cli and tool != "claude":
        raise BadArgs("--origin cli is only for claude")
    if tool in SHELLS:
        return parse_shell(tool, argv, rest, cwd, int(cols), int(rows), environ, fs)
    if tool not in allowed_tools(environ):
        raise BadArgs("tool not allowed")
    if tool == "ssh" and not is_ssh_args(rest):
        raise BadArgs("ssh takes only [-p N] -- <target>")
    if tool in AGENTS and not from_cli and not is_agent_args(tool, rest):
        raise BadArgs("argument not allowed for " + tool)
    # The folder once the tool is known: home itself only for ssh; every agent (and the
    # test fixture) needs a folder the launcher would offer.
    if not is_safe_abs_path(cwd):
        raise BadArgs("bad cwd")
    real_cwd, home = fs.realpath(cwd), fs.realpath(fs.home(environ))
    if not fs.is_dir(real_cwd) or not ((tool == "ssh" and real_cwd == home) or is_allowed_folder(real_cwd, home)):
        raise BadArgs("folder not allowed")
    program = fs.realpath(argv[0])
    if not fs.is_exec_file(program):
        raise BadArgs("argv0 is not an executable file")
    if not is_trusted_program(argv[0], program, environ, fs):
        raise BadArgs("argv0 is in a temporary or world-writable folder")
    if INTERPRETER.fullmatch(program.rsplit("/", 1)[-1]):
        raise BadArgs("argv0 resolves to a shell or an interpreter")
    if tool == "ssh" and not (argv[0] in SYSTEM_SSH and (program in SYSTEM_SSH or (program.endswith("/bin/ssh") and program.startswith(SSH_CELLARS)))):
        raise BadArgs("ssh must be the system's")
    return {"cwd": real_cwd, "cols": int(cols), "rows": int(rows), "program": program, "argv": list(argv)}


def parse_shell(tool, argv, rest, cwd, cols, rows, environ, fs):
    """A shell preset: from SHELL_DIRS, its fixed arguments, home or an allowed folder,
    resolving to a shell name outside temporary and world-writable folders."""
    if argv[0].rsplit("/", 1)[0] not in SHELL_DIRS:
        raise BadArgs("a shell only from the system folders")
    if not is_shell_args(tool, rest):
        raise BadArgs("a shell takes only -l (tmux: attach-session or new-session)")
    if not is_safe_abs_path(cwd):
        raise BadArgs("bad cwd")
    real_cwd, home = fs.realpath(cwd), fs.realpath(fs.home(environ))
    if not fs.is_dir(real_cwd) or not (real_cwd == home or is_allowed_folder(real_cwd, home)):
        raise BadArgs("folder not allowed")
    program = fs.realpath(argv[0])
    if not fs.is_exec_file(program):
        raise BadArgs("argv0 is not an executable file")
    if not is_trusted_program(argv[0], program, environ, fs):
        raise BadArgs("argv0 is in a temporary or world-writable folder")
    # What runs is a shell from the same places, and the arguments fit what runs (a tmux link to
    # bash would read new-session as a script).
    resolved_name = program.rsplit("/", 1)[-1]
    if resolved_name not in SHELLS or (resolved_name == "tmux") != (tool == "tmux") or not is_shell_args(resolved_name, rest):
        raise BadArgs("argv0 must resolve to that shell")
    if program.rsplit("/", 1)[0] not in SHELL_DIRS and not program.startswith(SHELL_STORES):
        raise BadArgs("the shell must resolve inside the system folders")
    return {"cwd": real_cwd, "cols": cols, "rows": rows, "program": program, "argv": list(argv)}


PASSED_ENV = ("HOME", "USER", "LOGNAME", "SHELL", "PATH", "LANG", "SSH_AUTH_SOCK", "TMPDIR")


def child_env(environ):
    """What the child sees: who the user is, PATH, the locale, the ssh agent. Nothing else."""
    env = {k: environ[k] for k in PASSED_ENV if environ.get(k)}
    env.update({k: v for k, v in environ.items() if k.startswith("LC_") and v})
    env.setdefault("PATH", "/usr/local/bin:/usr/bin:/bin")
    if not any(k in env for k in ("LANG", "LC_ALL", "LC_CTYPE")):
        env["LANG"] = "C.UTF-8"
    env["TERM"] = "xterm-256color"
    env["COLORTERM"] = "truecolor"
    return env


def exit_code(status):
    if os.WIFSIGNALED(status):
        return 128 + os.WTERMSIG(status)
    if os.WIFEXITED(status):
        return os.WEXITSTATUS(status)
    return 1


# ---- the session ------------------------------------------------------------

def winsize(cols, rows):
    return struct.pack("HHHH", rows, cols, 0, 0)


def run(req, environ):
    import pty

    env = child_env(environ)
    argv = req["argv"]
    wake_r, wake_w = os.pipe()
    for fd in (wake_r, wake_w):
        os.set_blocking(fd, False)
        os.set_inheritable(fd, False)
    signal.set_wakeup_fd(wake_w, warn_on_full_buffer=False)
    signal.signal(signal.SIGCHLD, lambda *_: None)
    try:
        max_fd = min(os.sysconf("SC_OPEN_MAX"), 65536)
    except (ValueError, OSError):
        max_fd = 4096

    pid, master = pty.fork()
    if pid == 0:
        try:
            fcntl.ioctl(0, termios.TIOCSWINSZ, winsize(req["cols"], req["rows"]))
            for sig in (signal.SIGPIPE, signal.SIGXFSZ, signal.SIGCHLD, signal.SIGINT, signal.SIGHUP, signal.SIGTERM,
                        signal.SIGQUIT, signal.SIGTSTP, signal.SIGTTIN, signal.SIGTTOU):
                signal.signal(sig, signal.SIG_DFL)
            signal.pthread_sigmask(signal.SIG_SETMASK, [])
            os.closerange(3, max_fd)
            os.chdir(req["cwd"])
            os.execve(req["program"], argv, env)
        except BaseException:
            try:
                os.write(2, b"dokki-pty: could not start the program\r\n")
            finally:
                os._exit(127)

    stdin, out = sys.stdin.fileno(), sys.stdout.fileno()
    for fd in (master, stdin, out):
        os.set_blocking(fd, False)
    sel = selectors.DefaultSelector()
    sel.register(stdin, selectors.EVENT_READ, "stdin")
    sel.register(wake_r, selectors.EVENT_READ, "wake")
    decoder = Decoder()
    pending = bytearray()   # host → terminal
    queued = bytearray()    # helper → host
    state = {"stdin": True, "master_open": True, "out": True, "kill_at": None, "status": None,
             "registered": {}, "dropped": False, "ending": False}

    def want(fd, events, tag):
        """Keep `fd` registered for exactly `events` (0: not at all)."""
        current = state["registered"].get(fd, 0)
        if events == current:
            return
        if current and events:
            sel.modify(fd, events, tag)
        elif events:
            sel.register(fd, events, tag)
        else:
            sel.unregister(fd)
        state["registered"][fd] = events

    def update():
        # Read the terminal only while the host keeps up (or the session is ending: a dying
        # program may wait for its output to drain), write it while input waits.
        reading = state["master_open"] and (len(queued) < OUTPUT_HIGH_WATER or state["ending"])
        writing = state["master_open"] and bool(pending)
        want(master, (selectors.EVENT_READ if reading else 0) | (selectors.EVENT_WRITE if writing else 0), "master")
        want(out, selectors.EVENT_WRITE if state["out"] and queued else 0, "out")

    def kill(sig):
        for target in (-pid, pid):
            try:
                os.kill(target, sig)
                return
            except ProcessLookupError:
                continue
            except PermissionError:
                return

    def arm_watchdog():
        """SIGALRM interrupts anything (select, a write): the helper cannot outlive it."""
        if not state.get("watchdog"):
            state["watchdog"] = True

            def expire(*_):
                for sig_target in (-pid, pid):
                    try:
                        os.kill(sig_target, signal.SIGKILL)
                    except OSError:
                        pass
                os._exit(0)

            signal.signal(signal.SIGALRM, expire)
            signal.setitimer(signal.ITIMER_REAL, WATCHDOG)

    def hang_up():
        arm_watchdog()
        state["ending"] = True
        if state["stdin"]:
            state["stdin"] = False
            sel.unregister(stdin)
        if state["kill_at"] is None and state["status"] is None:
            kill(signal.SIGHUP)
            state["kill_at"] = time.monotonic() + HANGUP_GRACE

    def emit(frame):
        if state["out"]:
            queued.extend(frame)

    def flush_out():
        try:
            while queued:
                del queued[:os.write(out, queued)]
        except (BlockingIOError, InterruptedError):
            pass
        except OSError:
            # The host went away: end the session like stdin EOF (never an orphaned child).
            state["out"] = False
            queued.clear()
            hang_up()

    def read_master():
        try:
            data = os.read(master, 16384)
        except (BlockingIOError, InterruptedError):
            return
        except OSError as error:
            if error.errno != errno.EIO:
                raise
            data = b""
        if data:
            # Ending (the helper hung up, or the program is gone) and the host is behind even
            # after a flush: what the program prints while dying is dropped.
            if state["ending"] and len(queued) >= OUTPUT_HIGH_WATER:
                flush_out()
            if not (state["ending"] and len(queued) >= OUTPUT_HIGH_WATER):
                emit(encode_data(data))
        else:
            state["master_open"] = False  # every holder of the terminal closed it

    def read_stdin():
        try:
            chunk = os.read(stdin, 65536)
        except (BlockingIOError, InterruptedError):
            return
        if not chunk:
            hang_up()
            return
        try:
            frames = decoder.feed(chunk)
        except FrameError as error:
            sys.stderr.write("dokki-pty: %s\n" % error.code)
            hang_up()
            return
        for ftype, payload in frames:
            if ftype == T_DATA:
                if len(pending) + len(payload) > MAX_PENDING_INPUT:
                    # The program is not reading: drop the overflow, never stop reading stdin.
                    if not state["dropped"]:
                        sys.stderr.write("dokki-pty: input dropped (program not reading)\n")
                        state["dropped"] = True
                    continue
                pending.extend(payload)
            elif ftype == T_RESIZE:
                cols, rows = struct.unpack(">HH", payload)
                try:
                    fcntl.ioctl(master, termios.TIOCSWINSZ, winsize(cols, rows))
                except OSError:
                    pass
            elif ftype == T_SIGNAL:
                # Only a signal: the program may handle it and go on, so backpressure stays.
                kill(SIGNALS[payload[0]])

    def reap():
        if state["status"] is None:
            try:
                done, status = os.waitpid(pid, os.WNOHANG)
            except ChildProcessError:
                done, status = pid, 0
            if done == pid:
                state["status"] = status
        return state["status"] is not None

    emit(encode(T_READY, struct.pack(">I", pid & 0xFFFFFFFF)))
    update()
    while not reap():
        timeout = 1.0
        if state["kill_at"] is not None:
            left = state["kill_at"] - time.monotonic()
            if left <= 0:
                kill(signal.SIGKILL)
                state["kill_at"] = None
            else:
                timeout = min(timeout, left)
        for key, events in sel.select(timeout):
            if key.data == "wake":
                try:
                    os.read(wake_r, 512)
                except BlockingIOError:
                    pass
            elif key.data == "stdin" and state["stdin"]:
                read_stdin()
            elif key.data == "out":
                flush_out()
            elif key.data == "master":
                if events & selectors.EVENT_READ:
                    read_master()
                if events & selectors.EVENT_WRITE and pending:
                    try:
                        del pending[:os.write(master, pending)]
                    except (BlockingIOError, InterruptedError):
                        pass
                    except OSError:
                        pending.clear()
        update()

    # The program is gone. Its group (pgid == pid) is hung up at once; then, for at most
    # HANGUP_GRACE, the terminal's last output (Linux flushes it asynchronously) is read and
    # passed on until the terminal is closed or quiet and the group is empty; what is left of
    # the group is killed. stdout stays non-blocking: a host that does not read gets
    # OUTPUT_DEADLINE, then the helper exits anyway.
    state["ending"] = True
    arm_watchdog()
    kill_group(pid, signal.SIGHUP)
    deadline = time.monotonic() + HANGUP_GRACE
    quiet = False
    while True:
        left = deadline - time.monotonic()
        if left <= 0:
            break
        flush_out()
        if state["master_open"] and not quiet:
            try:
                ready = select_readable(master, min(0.1, left))
            except InterruptedError:
                continue
            if ready:
                read_master()
            else:
                quiet = True
        elif group_alive(pid):
            time.sleep(min(0.05, left))
        else:
            break
    kill_group(pid, signal.SIGKILL)
    emit(encode(T_EXIT, struct.pack(">i", exit_code(state["status"]))))
    out_deadline = time.monotonic() + OUTPUT_DEADLINE
    while queued and state["out"]:
        flush_out()
        left = out_deadline - time.monotonic()
        if not queued or left <= 0:
            break
        try:
            select_writable(out, min(0.1, left))
        except InterruptedError:
            pass
    return 0


def select_writable(fd, timeout):
    import select
    _, writable, _ = select.select([], [fd], [], timeout)
    return bool(writable)


def select_readable(fd, timeout):
    import select
    readable, _, _ = select.select([fd], [], [], timeout)
    return bool(readable)


def kill_group(pgid, sig):
    try:
        os.killpg(pgid, sig)
    except (ProcessLookupError, PermissionError):
        pass


def group_alive(pgid):
    try:
        os.killpg(pgid, 0)
        return True
    except (ProcessLookupError, PermissionError):
        return False


def main(argv):
    if len(argv) < 1 or argv[0] != "pty-spawn":
        sys.stderr.write("usage: dokki-pty pty-spawn --cwd <path> --cols N --rows N -- <argv...>\n")
        return 2
    try:
        req = parse_spawn(argv[1:], os.environ)
    except BadArgs as error:
        sys.stderr.write("dokki-pty: %s\n" % error)
        return 2
    try:
        return run(req, os.environ)
    except BrokenPipeError:
        return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
