"""Bounded subprocess capture shared by execution, discovery and environment probes."""

from __future__ import annotations

from dataclasses import dataclass
import os
from pathlib import Path
import selectors
import signal
import subprocess
import threading
import time
from typing import Callable


@dataclass(frozen=True)
class OutputSnapshot:
    """Immutable bounded output; raw bytes are decoded only at the display boundary."""

    stdout: bytes
    stderr: bytes
    truncated: bool
    omitted_stdout_bytes: int
    omitted_stderr_bytes: int
    stdout_head_bytes: int
    stderr_head_bytes: int
    final: bool

    def text(self, stream: str = "stdout") -> str:
        """Render each retained segment separately, never joining cut UTF-8 sequences.

        This is not secret redaction. Callers must redact before publishing or
        persisting a snapshot, including withholding partial live lines as needed.
        """
        if stream not in {"stdout", "stderr"}:
            raise ValueError("Output stream must be stdout or stderr")
        data = getattr(self, stream)
        omitted = getattr(self, "omitted_" + stream + "_bytes")
        if not omitted:
            return data.decode("utf-8", errors="replace")
        split = getattr(self, stream + "_head_bytes")
        return (
            data[:split].decode("utf-8", errors="replace")
            + f"\n[{omitted} output bytes omitted between retained segments]\n"
            + data[split:].decode("utf-8", errors="replace")
        )


@dataclass(frozen=True)
class Capture:
    stdout: bytes
    stderr: bytes
    returncode: int
    status: str
    truncated: bool
    omitted_stdout_bytes: int = 0
    omitted_stderr_bytes: int = 0
    stdout_head_bytes: int = 0
    stderr_head_bytes: int = 0

    def snapshot(self) -> OutputSnapshot:
        return OutputSnapshot(
            self.stdout,
            self.stderr,
            self.truncated,
            self.omitted_stdout_bytes,
            self.omitted_stderr_bytes,
            self.stdout_head_bytes,
            self.stderr_head_bytes,
            True,
        )


class _OutputBuffer:
    def __init__(self, limit: int, mode: str):
        self.limit, self.mode, self.total = limit, mode, 0
        self.head, self.tail = bytearray(), bytearray()
        self.head_limit = limit if mode == "prefix" else limit // 2
        self.tail_limit = limit - self.head_limit

    def append(self, data: bytes) -> None:
        self.total += len(data)
        count = min(len(data), self.head_limit - len(self.head))
        self.head.extend(data[:count])
        remaining = data[count:]
        if self.tail_limit and remaining:
            if len(remaining) >= self.tail_limit:
                self.tail[:] = remaining[-self.tail_limit :]
            else:
                excess = max(0, len(self.tail) + len(remaining) - self.tail_limit)
                del self.tail[:excess]
                self.tail.extend(remaining)

    @property
    def omitted(self) -> int:
        return self.total - len(self.head) - len(self.tail)

    def value(self) -> bytes:
        return bytes(self.head) + bytes(self.tail)


def group_exists(pid: int) -> bool:
    try:
        os.killpg(pid, 0)
        return True
    except ProcessLookupError:
        return False


def stop_group(process: subprocess.Popen) -> None:
    try:
        os.killpg(process.pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    deadline = time.monotonic() + 0.5
    while time.monotonic() < deadline and group_exists(process.pid):
        process.poll()
        time.sleep(0.02)
    try:
        os.killpg(process.pid, signal.SIGKILL)
    except ProcessLookupError:
        pass
    process.wait(timeout=3)


def run_captured(
    argv: list[str],
    *,
    cwd: Path | str | None = None,
    env: dict[str, str] | None = None,
    timeout: float = 30,
    max_stdout: int = 1024 * 1024,
    max_stderr: int = 65536,
    input_bytes: bytes | None = None,
    cancelled: threading.Event | None = None,
    merge_stderr: bool = False,
    terminate_on_limit: bool = True,
    capture_mode: str = "prefix",
    on_output: Callable[[OutputSnapshot], None] | None = None,
) -> Capture:
    """No shell, capped memory, finite deadline, owned process-group cleanup.

    Prefix capture preserves probe semantics. Head/tail capture retains the
    beginning and actual latest output within the same byte bound. Optional
    callbacks receive a first-data snapshot, at most one dirty update per second,
    and a final snapshot. Callback failures propagate after owned-process cleanup.

    This is not a sandbox: a deliberately detached new process session is outside
    the original process group. Only run trusted project commands.
    """
    if capture_mode not in {"prefix", "head_tail"}:
        raise ValueError("Capture mode must be prefix or head_tail")
    if any(type(value) is not int or value < 0 for value in (max_stdout, max_stderr)):
        raise ValueError("Output bounds must be nonnegative integers")
    buffers = {
        "stdout": _OutputBuffer(max_stdout, capture_mode),
        "stderr": _OutputBuffer(max_stderr, capture_mode),
    }
    last_notified, notified_total = None, -1

    def snapshot(final: bool) -> OutputSnapshot:
        out, err = buffers["stdout"], buffers["stderr"]
        return OutputSnapshot(
            out.value(),
            err.value(),
            bool(out.omitted or err.omitted),
            out.omitted,
            err.omitted,
            len(out.head),
            len(err.head),
            final,
        )

    def notify(final: bool = False) -> None:
        nonlocal last_notified, notified_total
        if on_output is None:
            return
        total = sum(buffer.total for buffer in buffers.values())
        instant = time.monotonic()
        if final or (
            total > 0
            and total != notified_total
            and (last_notified is None or instant - last_notified >= 1)
        ):
            on_output(snapshot(final))
            last_notified, notified_total = instant, total

    process = subprocess.Popen(
        argv,
        cwd=cwd,
        env=env,
        stdin=subprocess.PIPE if input_bytes is not None else subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT if merge_stderr else subprocess.PIPE,
        start_new_session=True,
    )
    pending_input = memoryview(input_bytes or b"")
    selector = None
    try:
        selector = selectors.DefaultSelector()
        assert process.stdout is not None
        selector.register(process.stdout, selectors.EVENT_READ, "stdout")
        if process.stderr:
            selector.register(process.stderr, selectors.EVENT_READ, "stderr")
        if process.stdin:
            if pending_input:
                selector.register(process.stdin, selectors.EVENT_WRITE, "stdin")
            else:
                process.stdin.close()
    except BaseException:
        stop_group(process)
        if selector is not None:
            selector.close()
        for handle in (process.stdin, process.stdout, process.stderr):
            if handle and not handle.closed:
                handle.close()
        raise
    deadline = time.monotonic() + timeout
    exited_at = None
    state = None
    try:
        while selector.get_map() or process.poll() is None:
            if cancelled and cancelled.is_set():
                state = "interrupted"
                break
            if time.monotonic() >= deadline:
                state = "timed_out"
                break
            for key, _ in selector.select(
                timeout=min(0.1, max(0, deadline - time.monotonic()))
            ):
                if key.data == "stdin":
                    try:
                        written = os.write(key.fd, pending_input[:4096])
                        pending_input = pending_input[written:]
                    except BrokenPipeError:
                        pending_input = memoryview(b"")
                    if not pending_input:
                        selector.unregister(key.fileobj)
                        key.fileobj.close()
                    continue
                data = os.read(key.fd, 65536)
                if not data:
                    selector.unregister(key.fileobj)
                    continue
                target = buffers[key.data]
                target.append(data)
                if target.omitted and terminate_on_limit:
                    state = "output_limit"
                    break
            notify()
            if state:
                break
            if process.poll() is not None:
                exited_at = exited_at or time.monotonic()
                if selector.get_map() and time.monotonic() - exited_at >= 0.5:
                    state = "leaked_process"
                    break
        if state:
            stop_group(process)
        else:
            process.wait(timeout=3)
            # A child can redirect all output then outlive an otherwise successful parent.
            grace = time.monotonic() + 0.1
            while group_exists(process.pid) and time.monotonic() < grace:
                time.sleep(0.01)
            if group_exists(process.pid):
                state = "leaked_process"
                stop_group(process)
            else:
                state = "passed" if process.returncode == 0 else "failed"
        notify(final=True)
    finally:
        if process.poll() is None or group_exists(process.pid):
            stop_group(process)
        selector.close()
        for handle in (process.stdin, process.stdout, process.stderr):
            if handle and not handle.closed:
                handle.close()
    output = snapshot(True)
    return Capture(
        output.stdout,
        output.stderr,
        process.returncode,
        state,
        output.truncated,
        output.omitted_stdout_bytes,
        output.omitted_stderr_bytes,
        output.stdout_head_bytes,
        output.stderr_head_bytes,
    )


def require_capture(argv: list[str], **kwargs) -> Capture:
    result = run_captured(argv, **kwargs)
    if result.status in {"timed_out", "interrupted", "output_limit", "leaked_process"}:
        raise ValueError(
            f"Bounded subprocess did not complete normally: {result.status}"
        )
    return result