"""Run an interactive shell in a kernel pseudo-terminal over 4-byte-framed stdio.

The first stdin frame is a UTF-8 JSON object:
    {"cwd": "/path", "cols": 80, "rows": 24, "shell": "/bin/zsh", "credit": 8}

Outbound frames (stdout):
    0x00 + utf8-json: {"pid": int, "cwd": str, "shell": str}
    0x01 + bytes:     raw PTY output bytes
    0x02 + int32-be:  process exit status (negative for signals)

Inbound frames (stdin):
    0x01 + bytes:              write bytes to PTY master
    0x02 + uint16-be*2:        resize PTY to (cols, rows)
    0x03:                      terminate the PTY process group
    0x04 + uint8:              send signal number to foreground process group
    0x05 + uint16-be:          allow that many more output frames

With "credit": N in the spec, at most N output frames are sent before the
daemon grants more, and while none are allowed the PTY is not read, so a fast
producer blocks on the PTY instead of filling the daemon's mailbox. Without
it (a daemon from before flow control) output is never held back.
"""

from __future__ import annotations

import contextlib
import errno
import fcntl
import json
import os
import pty
import select
import shutil
import signal
import struct
import subprocess
import sys
import termios
import threading
import time

FRAME_READY = 0
FRAME_DATA = 1
FRAME_RESIZE = 2
FRAME_EXIT = 2
FRAME_STOP = 3
FRAME_SIGNAL = 4
FRAME_CREDIT = 5

STRIPPED_ENV = (
    "ALBEDO_DAEMON_TOKEN",
    "ALBEDO_INSPECT",
    "ALBEDO_PARENT_PID",
    "ALBEDO_TEST_DAEMON",
    "ALBEDO_PRELOAD",
    "PYTHONUNBUFFERED",
)

write_lock = threading.Lock()


class Credit:
    """Output frames the daemon allows; None means unlimited."""

    def __init__(self, frames: int | None) -> None:
        self.frames = frames
        self.closed = False
        self.changed = threading.Condition()

    def grant(self, frames: int) -> None:
        with self.changed:
            if self.frames is not None:
                self.frames += frames
            self.changed.notify()

    def close(self) -> None:
        with self.changed:
            self.closed = True
            self.changed.notify()

    def take(self) -> bool:
        with self.changed:
            while self.frames == 0 and not self.closed:
                self.changed.wait()
            if self.closed:
                return False
            if self.frames is not None:
                self.frames -= 1
            return True


def send_frame(kind: int, payload: bytes) -> bool:
    frame = struct.pack("!IB", len(payload) + 1, kind) + payload
    with write_lock:
        view = memoryview(frame)
        while view:
            try:
                written = os.write(1, view)
            except OSError:
                return False
            if written <= 0:
                return False
            view = view[written:]
    return True


def read_exact(fd: int, size: int) -> bytes | None:
    chunks = bytearray()
    while len(chunks) < size:
        try:
            chunk = os.read(fd, size - len(chunks))
        except OSError as error:
            if error.errno == errno.EINTR:
                continue
            return None
        if not chunk:
            return None
        chunks.extend(chunk)
    return bytes(chunks)


def read_frame(fd: int) -> bytes | None:
    header = read_exact(fd, 4)
    if header is None:
        return None
    (length,) = struct.unpack("!I", header)
    if length == 0 or length > 1_048_576:
        return None
    return read_exact(fd, length)


def resolve_cwd(requested: object) -> str:
    if isinstance(requested, str) and requested.strip():
        expanded = os.path.abspath(os.path.expanduser(requested.strip()))
        if os.path.isdir(expanded):
            return expanded
    home = os.path.expanduser("~")
    return home if os.path.isdir(home) else "/"


def resolve_shell(requested: object) -> str:
    candidates: list[str] = []
    if isinstance(requested, str) and requested.strip():
        candidates.append(requested.strip())
    env_shell = os.environ.get("SHELL", "").strip()
    if env_shell:
        candidates.append(env_shell)
    candidates.extend(("/bin/zsh", "/bin/bash", "/bin/sh", "sh"))
    for candidate in candidates:
        resolved = (
            candidate
            if os.path.isabs(candidate) and os.access(candidate, os.X_OK)
            else shutil.which(candidate)
        )
        if resolved and os.access(resolved, os.X_OK):
            return resolved
    return "/bin/sh"


def build_env(cwd: str, shell: str) -> dict[str, str]:
    env = {k: v for k, v in os.environ.items() if k not in STRIPPED_ENV}
    env["TERM"] = "xterm-256color"
    env["COLORTERM"] = "truecolor"
    env["TERM_PROGRAM"] = "albedo"
    env["PWD"] = cwd
    env["SHELL"] = shell
    env.setdefault("HOME", os.path.expanduser("~"))
    env.setdefault("LANG", "en_US.UTF-8")
    return env


def set_winsize(fd: int, cols: int, rows: int) -> None:
    c = max(2, min(500, cols))
    r = max(1, min(200, rows))
    with contextlib.suppress(OSError):
        fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", r, c, 0, 0))


def terminate_group(pid: int) -> None:
    for sig, wait_s in (
        (signal.SIGHUP, 0.08),
        (signal.SIGTERM, 0.12),
        (signal.SIGKILL, 0.05),
    ):
        try:
            os.killpg(pid, sig)
        except ProcessLookupError:
            return
        except PermissionError:
            with contextlib.suppress(OSError):
                os.kill(pid, sig)
        deadline = time.monotonic() + wait_s
        while time.monotonic() < deadline:
            try:
                waited, _ = os.waitpid(pid, os.WNOHANG)
                if waited == pid:
                    return
            except ChildProcessError:
                return
            time.sleep(0.01)


def write_all(fd: int, data: bytes) -> None:
    view = memoryview(data)
    while view:
        try:
            written = os.write(fd, view)
        except OSError as error:
            if error.errno in (errno.EINTR, errno.EAGAIN):
                time.sleep(0.005)
                continue
            return
        if written <= 0:
            return
        view = view[written:]


def read_output(master_fd: int, proc: subprocess.Popen[bytes]) -> bytes:
    """Return the next PTY output, or b"" once the shell is gone."""
    while True:
        try:
            ready, _, _ = select.select([master_fd], [], [], 0.5)
        except (OSError, ValueError):
            return b""
        if not ready:
            if proc.poll() is not None:
                return b""
            continue
        try:
            return os.read(master_fd, 16_384)
        except OSError as error:
            if error.errno != errno.EINTR:
                return b""


def pump_output(master_fd: int, proc: subprocess.Popen[bytes], credit: Credit) -> None:
    try:
        while credit.take():
            chunk = read_output(master_fd, proc)
            if not chunk or not send_frame(FRAME_DATA, chunk):
                break
    finally:
        try:
            code = proc.wait(timeout=0.5)
        except subprocess.TimeoutExpired:
            terminate_group(proc.pid)
            code = proc.wait()
        terminate_group(proc.pid)
        send_frame(FRAME_EXIT, struct.pack("!i", int(code)))


def main() -> int:
    init_raw = read_frame(0)
    if init_raw is None:
        return 1
    try:
        spec = json.loads(init_raw.decode("utf-8"))
    except Exception:
        return 1
    if not isinstance(spec, dict):
        return 1

    cwd = resolve_cwd(spec.get("cwd"))
    shell = resolve_shell(spec.get("shell"))
    cols = spec.get("cols") if isinstance(spec.get("cols"), int) else 80
    rows = spec.get("rows") if isinstance(spec.get("rows"), int) else 24

    master_fd, slave_fd = pty.openpty()
    set_winsize(master_fd, cols, rows)

    def child_setup() -> None:
        os.setsid()
        with contextlib.suppress(OSError):
            fcntl.ioctl(0, termios.TIOCSCTTY, 0)

    try:
        proc = subprocess.Popen(
            [shell],
            stdin=slave_fd,
            stdout=slave_fd,
            stderr=slave_fd,
            cwd=cwd,
            env=build_env(cwd, shell),
            preexec_fn=child_setup,
            close_fds=True,
        )
    except Exception:
        os.close(master_fd)
        os.close(slave_fd)
        return 1
    finally:
        with contextlib.suppress(OSError):
            os.close(slave_fd)

    ready_payload = json.dumps(
        {"pid": proc.pid, "cwd": cwd, "shell": os.path.basename(shell)}
    ).encode("utf-8")
    if not send_frame(FRAME_READY, ready_payload):
        terminate_group(proc.pid)
        os.close(master_fd)
        return 1

    window = spec.get("credit")
    credit = Credit(window if isinstance(window, int) and window > 0 else None)
    reader = threading.Thread(
        target=pump_output, args=(master_fd, proc, credit), daemon=True
    )
    reader.start()

    try:
        while proc.poll() is None:
            try:
                ready, _, _ = select.select([0], [], [], 0.2)
            except (OSError, ValueError):
                break
            if not ready:
                continue
            frame = read_frame(0)
            if frame is None:
                break
            kind = frame[0]
            body = frame[1:]
            if kind == FRAME_DATA and body:
                write_all(master_fd, body)
            elif kind == FRAME_RESIZE and len(body) == 4:
                new_cols, new_rows = struct.unpack("!HH", body)
                set_winsize(master_fd, new_cols, new_rows)
            elif kind == FRAME_STOP:
                break
            elif kind == FRAME_CREDIT and len(body) == 2:
                credit.grant(struct.unpack("!H", body)[0])
            elif kind == FRAME_SIGNAL and len(body) == 1:
                sig = body[0]
                try:
                    fg_pgid = os.tcgetpgrp(master_fd)
                except OSError:
                    fg_pgid = proc.pid
                with contextlib.suppress(OSError):
                    os.killpg(fg_pgid, sig)
    finally:
        credit.close()
        terminate_group(proc.pid)
        with contextlib.suppress(OSError):
            os.close(master_fd)
        reader.join(timeout=1.0)
    return 0


if __name__ == "__main__":
    sys.exit(main())