"""Timed command execution with admitted-output capture.

This is the measurement primitive every CLI lane shares: run a command under a

monotonic clock, capture full output bytes to disk, and admit only a capped

character window into token estimates (full byte counts are always recorded).

Memory: the child's peak RSS is captured at reap time via os.wait4 (rusage),

so no wrapper binary or sampling thread perturbs the timing. ru_maxrss units

differ by platform (bytes on macOS, KiB on Linux); peak_rss_bytes is always

bytes. None means unmeasured (ADR-0002), never zero — e.g. on platforms

without wait4.

"""

from __future__ import annotations

import os

import signal

import subprocess

import sys

import tempfile

import time

from dataclasses import dataclass

from pathlib import Path

from typing import Any, Optional

_HAS_WAIT4 = hasattr(os, "wait4")

PEAK_RSS_SOURCE = "os_wait4_ru_maxrss" if _HAS_WAIT4 else None

# CPU and block I/O come from the same wait4 rusage as peak RSS: the direct

# child's own usage (on macOS and Linux, including descendants it reaped

# before exiting; daemons it leaves running are not included).

RUSAGE_SOURCE = "os_wait4_child_rusage" if _HAS_WAIT4 else None

@dataclass(frozen=True)

class TimedCapture:

    elapsed_ms: float

    returncode: int

    stdout_text: str

    stdout_bytes: int

    stdout_truncated: bool

    stderr_text: str

    stderr_bytes: int

    stderr_truncated: bool

    peak_rss_bytes: Optional[int] = None

    stdout_full_text: Optional[str] = None

    stdout_full_text_truncated: Optional[bool] = None

    stderr_full_text: Optional[str] = None

    stderr_full_text_truncated: Optional[bool] = None

    timed_out: bool = False

    stdout_raw: bytes | None = None

    cpu_user_ms: Optional[float] = None

    cpu_sys_ms: Optional[float] = None

    io_read_blocks: Optional[int] = None

    io_write_blocks: Optional[int] = None

    major_page_faults: Optional[int] = None

def stop_owned_process(proc: subprocess.Popen[Any]) -> None:

    """Stop the timed child and its owned process group, then always reap it."""

    try:

        if os.name == "posix":

            os.killpg(proc.pid, signal.SIGTERM)

        else:

            proc.terminate()

    except (OSError, ProcessLookupError):

        pass

    try:

        proc.wait(timeout=1.0)

    except subprocess.TimeoutExpired:

        pass

    # The process-group leader may exit on SIGTERM while a descendant ignores

    # it. Probe and kill the owned group even after the leader has been reaped.

    if os.name == "posix":

        try:

            os.killpg(proc.pid, 0)

        except (OSError, ProcessLookupError):

            pass

        else:

            try:

                os.killpg(proc.pid, signal.SIGKILL)

            except (OSError, ProcessLookupError):

                pass

    elif proc.poll() is None:

        try:

            proc.kill()

        except (OSError, ProcessLookupError):

            pass

    if proc.poll() is None:

        proc.wait()

class _RusagePopen(subprocess.Popen):

    """Popen that records the child's rusage when it is reaped.

    Overriding _try_wait keeps a single wait4 syscall as both the reap and the

    measurement; there is no race against the kernel discarding rusage and no

    extra process in the timed path.

    """

    rusage = None

    def _try_wait(self, wait_flags):  # type: ignore[override]

        try:

            pid, sts, rusage = os.wait4(self.pid, wait_flags)

        except ChildProcessError:

            pid = self.pid

            sts = 0

        else:

            if pid == self.pid:

                self.rusage = rusage

        return (pid, sts)

def _peak_rss_bytes(rusage: Any) -> Optional[int]:

    if rusage is None:

        return None

    maxrss = getattr(rusage, "ru_maxrss", None)

    if not maxrss or maxrss < 0:

        return None

    # macOS reports bytes; Linux (and most other Unixes) report KiB.

    return int(maxrss) if sys.platform == "darwin" else int(maxrss) * 1024

def rusage_fields(rusage: Any) -> dict[str, Any]:

    """CPU user/sys ms and block I/O counts; None when unmeasured (ADR-0002).

    ru_inblock/ru_oublock count real block I/O, so a warm page cache yields a

    measured 0, not an unmeasured value.

    """

    if rusage is None:

        return {"cpu_user_ms": None, "cpu_sys_ms": None, "io_read_blocks": None,

                "io_write_blocks": None, "major_page_faults": None}

    return {

        "cpu_user_ms": round(float(rusage.ru_utime) * 1000.0, 3),

        "cpu_sys_ms": round(float(rusage.ru_stime) * 1000.0, 3),

        "io_read_blocks": int(rusage.ru_inblock),

        "io_write_blocks": int(rusage.ru_oublock),

        "major_page_faults": int(rusage.ru_majflt),

    }

def read_admitted_text(handle: Any, char_limit: int) -> tuple[str, int, bool]:

    handle.flush()

    handle.seek(0, os.SEEK_END)

    byte_count = handle.tell()

    handle.seek(0)

    if char_limit <= 0:

        return "", byte_count, byte_count > 0

    raw = handle.read(char_limit * 4 + 16)

    text = raw.decode("utf-8", "replace")[:char_limit]

    return text, byte_count, byte_count > len(text.encode("utf-8", "replace"))

def read_full_text_with_byte_limit(handle: Any, byte_limit: int) -> tuple[str | None, bool]:

    handle.flush()

    handle.seek(0, os.SEEK_END)

    byte_count = handle.tell()

    if byte_limit < 0:

        return None, byte_count > 0

    if byte_count > byte_limit:

        return None, True

    handle.seek(0)

    return handle.read().decode("utf-8", "replace"), False

def run_timed(

    command: list[str],

    cwd: Path,

    env: dict[str, str],

    admitted_output_chars: int,

    full_output_bytes: int | None = None,

    timeout_seconds: float | None = None,

    raw_output_bytes: int | None = None,

) -> TimedCapture:

    if timeout_seconds is not None and timeout_seconds <= 0:

        raise ValueError("timeout_seconds must be positive")

    if raw_output_bytes is not None and raw_output_bytes < 0:

        raise ValueError("raw_output_bytes must be nonnegative")

    popen_cls = _RusagePopen if _HAS_WAIT4 else subprocess.Popen

    start = time.perf_counter()

    with tempfile.TemporaryFile() as stdout_fh, tempfile.TemporaryFile() as stderr_fh:

        proc = popen_cls(

            command,

            cwd=cwd,

            env=env,

            stdout=stdout_fh,

            stderr=stderr_fh,

            start_new_session=timeout_seconds is not None and os.name == "posix",

        )

        timed_out = False

        try:

            returncode = proc.wait(timeout=timeout_seconds)

        except subprocess.TimeoutExpired:

            timed_out = True

            stop_owned_process(proc)

            returncode = 124

        elapsed_ms = (time.perf_counter() - start) * 1000

        stdout_text, stdout_bytes, stdout_truncated = read_admitted_text(stdout_fh, admitted_output_chars)

        stderr_text, stderr_bytes, stderr_truncated = read_admitted_text(stderr_fh, admitted_output_chars)

        stdout_raw = None

        if raw_output_bytes is not None and stdout_bytes <= raw_output_bytes:

            stdout_fh.seek(0)

            stdout_raw = stdout_fh.read()

        stdout_full_text = None

        stdout_full_text_truncated = None

        stderr_full_text = None

        stderr_full_text_truncated = None

        if full_output_bytes is not None:

            stdout_full_text, stdout_full_text_truncated = read_full_text_with_byte_limit(

                stdout_fh, full_output_bytes

            )

            stderr_full_text, stderr_full_text_truncated = read_full_text_with_byte_limit(

                stderr_fh, full_output_bytes

            )

    return TimedCapture(

        elapsed_ms=elapsed_ms,

        returncode=returncode,

        stdout_text=stdout_text,

        stdout_bytes=stdout_bytes,

        stdout_truncated=stdout_truncated,

        stderr_text=stderr_text,

        stderr_bytes=stderr_bytes,

        stderr_truncated=stderr_truncated,

        peak_rss_bytes=_peak_rss_bytes(getattr(proc, "rusage", None)),

        stdout_full_text=stdout_full_text,

        stdout_full_text_truncated=stdout_full_text_truncated,

        stderr_full_text=stderr_full_text,

        stderr_full_text_truncated=stderr_full_text_truncated,

        timed_out=timed_out,

        stdout_raw=stdout_raw,

        **rusage_fields(getattr(proc, "rusage", None)),

    )