"""Long-running commands: start one, read what it has printed, feed it input, stop it."""

from __future__ import annotations

import os
import select
import signal
import subprocess
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path

OUTPUT_BUFFER_BYTES = 1 << 20
OUTPUT_CAPTURE_BYTES = 8 << 20
STOP_GRACE_SECONDS = 5.0
MAX_LIVE_PROCESSES = 8
STDIN_WRITE_TIMEOUT_SECONDS = 5.0


def _kill_group(pgid: int) -> None:
    for sig in (signal.SIGTERM, signal.SIGKILL):
        try:
            os.killpg(pgid, sig)
        except ProcessLookupError:
            return
        time.sleep(0.2)


def run_bounded(
    argv: list[str], *, cwd: Path, env: dict[str, str], timeout: float, max_bytes: int = OUTPUT_CAPTURE_BYTES
) -> tuple[int | None, str, bool]:
    """Run to completion, capping wall time AND captured bytes. `capture_output=True`
    would buffer without limit: a command printing `yes` reaches gigabytes of RSS
    in seconds, which is a one-call denial of service against the host."""
    popen = subprocess.Popen(
        argv,
        cwd=cwd,
        env=env,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    pgid = os.getpgid(popen.pid)
    stream = popen.stdout
    assert stream is not None
    os.set_blocking(stream.fileno(), False)
    collected = bytearray()
    truncated = False
    deadline = time.monotonic() + timeout
    timed_out = False
    try:
        while True:
            if popen.poll() is not None and not select.select([stream], [], [], 0)[0]:
                break
            if time.monotonic() >= deadline:
                timed_out = True
                break
            if not select.select([stream], [], [], 0.05)[0]:
                continue
            chunk = stream.read(65536)
            if chunk is None:
                continue
            if not chunk:
                if popen.poll() is not None:
                    break
                continue
            room = max_bytes - len(collected)
            if room <= 0:
                truncated = True
                break
            collected += chunk[:room]
    finally:
        if popen.poll() is None:
            _kill_group(pgid)
        stream.close()
        popen.wait()
    exit_code = None if timed_out or truncated else popen.returncode
    return exit_code, collected.decode("utf-8", "replace"), truncated


@dataclass
class ManagedProcess:
    process_id: str
    command: str
    cwd: str
    popen: subprocess.Popen[bytes]
    started_at: float
    # Captured at start: reading it later races the reaper, and after PID reuse
    # os.getpgid would name a group this registry never started.
    pgid: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock)
    _buffer: bytearray = field(default_factory=bytearray)
    _dropped: int = 0

    @property
    def running(self) -> bool:
        return self.popen.poll() is None

    @property
    def exit_code(self) -> int | None:
        return self.popen.poll()

    def _pump(self) -> None:
        stream = self.popen.stdout
        assert stream is not None
        # read1 returns whatever has arrived; read(n) would block until n bytes,
        # hiding a running process's output until it produced 4 KiB or exited.
        for chunk in iter(lambda: stream.read1(4096), b""):
            with self._lock:
                self._buffer += chunk
                overflow = len(self._buffer) - OUTPUT_BUFFER_BYTES
                if overflow > 0:
                    del self._buffer[:overflow]
                    self._dropped += overflow
        stream.close()
        self.popen.wait()

    def read(self, since: int) -> tuple[str, int, int]:
        """Output from absolute offset `since`, the next offset, and how many bytes
        were discarded before the returned slice because the buffer overflowed."""
        with self._lock:
            start = max(since, self._dropped, 0)
            data = bytes(self._buffer[start - self._dropped :])
            return data.decode("utf-8", "replace"), start + len(data), start - since

    def write_stdin(self, text: str) -> None:
        if not self.running:
            raise ValueError(f"process {self.process_id} has already exited")
        stdin = self.popen.stdin
        if stdin is None or stdin.closed:
            raise ValueError(f"process {self.process_id} has no open stdin")
        # A process that never reads stdin fills the 64 KiB pipe and a blocking
        # write hangs forever. On one event loop that freezes every other tool,
        # including the one that would stop this process.
        payload = text.encode("utf-8")
        fd = stdin.fileno()
        os.set_blocking(fd, False)
        deadline = time.monotonic() + STDIN_WRITE_TIMEOUT_SECONDS
        while payload:
            if time.monotonic() >= deadline:
                raise ValueError(
                    f"process {self.process_id} is not reading stdin; "
                    f"{len(payload)} bytes were not delivered"
                )
            if not select.select([], [fd], [], 0.05)[1]:
                continue
            try:
                payload = payload[os.write(fd, payload) :]
            except BlockingIOError:
                continue
            except BrokenPipeError:
                raise ValueError(f"process {self.process_id} closed its stdin") from None

    def stop(self) -> str:
        if not self.running:
            return f"process {self.process_id} had already exited ({self.exit_code})"
        os.killpg(self.pgid, signal.SIGTERM)
        deadline = time.monotonic() + STOP_GRACE_SECONDS
        while time.monotonic() < deadline:
            if not self.running:
                return f"process {self.process_id} stopped"
            time.sleep(0.1)
        os.killpg(self.pgid, signal.SIGKILL)
        self.popen.wait(timeout=STOP_GRACE_SECONDS)
        return f"process {self.process_id} killed after {STOP_GRACE_SECONDS:.0f}s grace"


class ProcessRegistry:
    def __init__(self) -> None:
        self._processes: dict[str, ManagedProcess] = {}
        self._next = 1
        self._lock = threading.Lock()

    def start(self, argv: list[str], *, command: str, cwd: Path, env: dict[str, str]) -> ManagedProcess:
        live = sum(1 for managed in self.all() if managed.running)
        if live >= MAX_LIVE_PROCESSES:
            raise ValueError(
                f"{live} background processes are already running (limit {MAX_LIVE_PROCESSES}); "
                "stop one before starting another"
            )
        popen = subprocess.Popen(
            argv,
            cwd=cwd,
            env=env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        with self._lock:
            process_id = f"proc-{self._next}"
            self._next += 1
            managed = ManagedProcess(
                process_id=process_id,
                command=command,
                cwd=str(cwd),
                popen=popen,
                started_at=time.time(),
                pgid=os.getpgid(popen.pid),
            )
            self._processes[process_id] = managed
        threading.Thread(target=managed._pump, daemon=True).start()
        return managed

    def get(self, process_id: str) -> ManagedProcess:
        try:
            return self._processes[process_id]
        except KeyError:
            raise ValueError(f"no such process: {process_id}") from None

    def all(self) -> list[ManagedProcess]:
        with self._lock:
            return list(self._processes.values())

    def stop_all(self) -> None:
        for managed in self.all():
            if managed.running:
                managed.stop()
