"""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"}

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
"""

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

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

write_lock = threading.Lock()


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 pump_output(master_fd: int, proc: subprocess.Popen[bytes]) -> None:
    try:
        while True:
            try:
                ready, _, _ = select.select([master_fd], [], [], 0.5)
            except (OSError, ValueError):
                break
            if not ready:
                if proc.poll() is not None:
                    break
                continue
            try:
                chunk = os.read(master_fd, 16_384)
            except OSError as error:
                if error.errno == errno.EINTR:
                    continue
                break
            if not chunk:
                break
            if 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

    reader = threading.Thread(target=pump_output, args=(master_fd, proc), 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_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:
        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())