"""Subprocess registry + output pumping for the ``abk dashboard`` cockpit.

Ported near-verbatim from the donor's ``detectkit/ui/jobs.py``
(``docs/specs/m11-implementation-plan.md`` DASH-1). The dashboard server never
runs the pipeline in-process — every Run / Unlock / Clean / Explore click
spawns the real ``abk`` CLI as a subprocess, exactly as if typed at a
terminal. This module only tracks those subprocesses: pumping their merged
stdout/stderr into an in-memory line buffer the page polls, and reporting
status/return code. Nothing here touches the DB, the pipeline lock, or
``abkit.stats``.

Deliberate deviations from the donor (DASH-1 steps 1/3/4/5, plus two hazards
the DASH-1 review reproduced in the donor's own shape):

* **The kind vocabulary is abkit's** — ``run``/``unlock``/``clean``/``explore``
  (the donor's ``autotune``/``tune`` have no abkit equivalent). ``explore``
  plays the donor's ``tune`` role: it is excluded from the one-at-a-time
  pipeline gate (concurrent cockpits on *different* experiments are safe)
  while being deduped per experiment, because two cockpits on the *same*
  experiment race the same Apply-rewrites-YAML hazard. The vocabulary lives in
  :data:`JOB_KINDS` / :data:`PIPELINE_KINDS`, and the gate asks whether a kind
  is IN the pipeline set where the donor asked whether it was ``!= "tune"``.
  That flip is why both spawn entry points **validate against the vocabulary**:
  under the donor's blacklist an unknown kind erred toward being gated, under a
  whitelist it would silently escape the gate instead — so the set and the
  validation are one mechanism, not a constant plus a nicety.
* **``Job.experiment`` replaces the donor's ``Job.metric``.** abkit's whole
  dashboard grain is the experiment (every button is experiment-scoped), so
  the dedup key is a purpose-built field rather than an overloaded one. Unlike
  the donor's ``metric`` — which only ever fed ``running_tune_for`` — it is
  accepted by :meth:`JobManager.spawn_pipeline` too and rides along in both
  snapshot shapes, so a client can label a job from a structured field instead
  of parsing ``label``. It is only ever as populated as the caller makes it:
  a spawn that passes no ``experiment`` (a multi-experiment ``abk run``, say)
  reports ``None``, and a label is the honest fallback there.
* **:meth:`JobManager.wait_for_line` is drop-aware.** The donor tracked its
  scan position as an index into the *current* buffer, which stops advancing
  once the buffer hits :data:`_MAX_LINES` (each append pops the front, so the
  length is pinned) — the watcher then goes permanently blind. It now counts
  ABSOLUTE line indices, the scheme :meth:`JobManager.snapshot` already used,
  so a job that is chattier than the cap before it prints the line being
  waited for (``Explore: <url>``) is still matched instead of failing at the
  timeout.
* **A spawn racing :meth:`JobManager.shutdown` cannot orphan its child.**
  ``shutdown`` snapshots the registry and reaps that snapshot, so in the donor
  a job appended afterwards survives the teardown — reparented, untracked, and
  in abkit's case potentially an ``abk run`` still holding the pipeline lock
  (reproduced 300/300 in review). The registry now latches closed under the
  same lock that takes the snapshot: a later :meth:`JobManager.spawn` kills the
  child it just created and raises :class:`JobManagerClosed`, so every spawned
  process is either in the snapshot or never left alive.
* **A job that exits successfully is never reported as ``stopped``.** The pump
  reaps the process before it takes ``job.lock``, so in the donor a Stop click
  landing in that window flips an already-successful run to ``stopped`` even
  though its ``terminate()`` hit an exited process and did nothing. Status now
  reads ``stopped`` only when a stop was requested **and** the exit was not a
  clean one. The rule cannot separate that race from a job that caught SIGTERM
  and chose to exit 0 — nothing observable distinguishes them — so it prefers
  the report that cannot be a lie about success: both read ``done``. No abkit
  verb installs a SIGTERM handler today; if one ever does and the distinction
  starts to matter, it has to come from the job's own cooperation.
  :meth:`JobManager.shutdown` marks the jobs it kills as stop-requested for the
  same reason — a teardown is deliberate, so those read ``stopped``, not
  ``failed`` (the donor left them indistinguishable from a crash).
* **A spawned job never shares the cockpit's stdin** (``stdin=DEVNULL``; DASH-4).
  The donor leaves it inherited, so a child that ever prompts reads the
  operator's terminal out from under the dashboard — invisibly, since the child's
  prompt goes to the pumped pipe and the keystrokes never reach the shell the
  cockpit was launched from. abkit has one prompting path today (``abk clean
  --orphaned-experiments``, which no route spawns), and on a closed stdin
  ``click.confirm`` aborts loudly: the job reads ``failed`` with the reason in
  its own output, which is the outcome a launcher wants.
* **Per-experiment dedup is atomic** (:meth:`JobManager.spawn_deduped`; DASH-4).
  ``running_job_for`` on its own is a check-then-act across request threads, so
  a double-clicked Explore button starts two cockpits on one experiment — the
  exact Apply-rewrites-the-YAML race the dedup exists to prevent (each session
  writes the whole file from its own startup snapshot). The check and the spawn
  now happen under the same gate :meth:`spawn_pipeline` uses, which is the
  DASH-1 TOCTOU fix applied to the other spawn entry point.
"""

from __future__ import annotations

import secrets
import subprocess
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# Cap on retained stdout/stderr lines per job — a runaway/verbose process must
# not grow memory unboundedly; the oldest lines are dropped once the cap is hit.
_MAX_LINES = 5000
# Cap on retained jobs — the drawer only ever needs recent history.
_MAX_JOBS = 20
# Grace period between SIGTERM and SIGKILL when stopping a job.
_STOP_GRACE_SECONDS = 5.0

#: Job kinds the dashboard spawns (abkit CLI verbs; DASH-1 step 1).
JOB_KINDS = frozenset({"run", "unlock", "clean", "explore"})
#: The kinds that serialize against each other: they all mutate the same
#: DB-level state, so the cockpit only ever lets one of them run at a time.
#: ``explore`` is deliberately absent (DASH-1 step 3) — it only reads results
#: and writes YAML on an explicit Apply.
PIPELINE_KINDS = frozenset({"run", "unlock", "clean"})


class JobManagerClosed(RuntimeError):
    """Raised by :meth:`JobManager.spawn` once :meth:`JobManager.shutdown` ran.

    A job route racing the dashboard's teardown must answer "not now" (a 503)
    rather than leave a subprocess nobody will ever reap.
    """


def _now_ms() -> int:
    return int(datetime.now(timezone.utc).timestamp() * 1000)


def _terminate_and_reap(proc: subprocess.Popen[str]) -> None:
    """SIGTERM now; SIGKILL from a daemon thread if the grace period lapses.

    Never blocks the caller: a request handler asking for a stop must not wait
    out :data:`_STOP_GRACE_SECONDS` before replying.

    The kill branch reaps too. Where a pump thread exists it would do that
    anyway, but the rejected-spawn path in :meth:`JobManager.spawn` has no
    pump — and without the second wait the killed child stays a zombie for as
    long as the dashboard lives.
    """
    try:
        proc.terminate()
    except Exception:
        pass

    def _grace() -> None:
        try:
            proc.wait(timeout=_STOP_GRACE_SECONDS)
        except subprocess.TimeoutExpired:
            try:
                proc.kill()
                proc.wait()  # SIGKILL is not deniable; this returns promptly
            except Exception:
                pass

    threading.Thread(target=_grace, daemon=True).start()


@dataclass
class Job:
    """One spawned subprocess and its captured output.

    All mutable fields (``lines``, ``status``, ``returncode``, ``url``,
    ``finished_at``, ``stop_requested``) are guarded by ``lock`` — the pump
    thread writes them, request-handler threads read/poll them.
    """

    id: str
    kind: str  # "run" | "unlock" | "clean" | "explore"
    label: str
    argv: list[str]
    proc: subprocess.Popen[str]
    started_at: int
    # explore jobs: the experiment the cockpit was opened on (dedup key).
    experiment: str | None = None
    lock: threading.Lock = field(default_factory=threading.Lock)
    lines: list[str] = field(default_factory=list)
    # Count of lines dropped off the front of ``lines`` once the buffer cap is
    # hit. Poll offsets are ABSOLUTE line indices (dropped + buffered), so a
    # verbose job keeps streaming past the cap instead of the poller's offset
    # pinning at the buffer length and never advancing again.
    dropped: int = 0
    truncated: bool = False
    status: str = "running"  # running | done | failed | stopped
    returncode: int | None = None
    url: str | None = None
    finished_at: int | None = None
    stop_requested: bool = False


def _pump(job: Job) -> None:
    """Read *job*'s merged stdout line by line into its buffer until it exits."""
    assert job.proc.stdout is not None
    try:
        for raw_line in job.proc.stdout:
            line = raw_line.rstrip("\n")
            with job.lock:
                job.lines.append(line)
                if len(job.lines) > _MAX_LINES:
                    job.lines.pop(0)
                    job.dropped += 1
                    job.truncated = True
    finally:
        returncode = job.proc.wait()
        with job.lock:
            job.returncode = returncode
            # The reap above happens OUTSIDE job.lock, so a stop() landing in
            # that window sets stop_requested on a process that had already
            # exited under its own power (its terminate() reached nothing).
            # A clean exit therefore outranks the flag: "stopped" is reserved
            # for a stop request the job did not survive intact.
            if job.stop_requested and returncode != 0:
                job.status = "stopped"
            else:
                job.status = "done" if returncode == 0 else "failed"
            job.finished_at = _now_ms()


class JobManager:
    """In-memory registry of spawned CLI subprocesses (last :data:`_MAX_JOBS`)."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # Serializes every check-and-spawn: the one-at-a-time pipeline gate
        # (spawn_pipeline, run/unlock/clean) and the per-experiment explore
        # dedup (spawn_deduped), so two near-simultaneous POSTs can't both
        # observe "nothing running" and each start a subprocess. Separate from
        # _lock: spawn() acquires _lock internally, so holding _gate around it
        # must not self-deadlock. Held only for the microseconds of a check plus
        # a Popen — never while a job runs.
        self._gate = threading.Lock()
        self._jobs: list[Job] = []
        # Latched by shutdown() under _lock, in the same critical section that
        # snapshots the registry — that is what makes "every spawned child is
        # either in the snapshot or killed by its own spawn()" airtight.
        self._closed = False

    def spawn(
        self,
        kind: str,
        label: str,
        argv: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        experiment: str | None = None,
    ) -> Job:
        """Start *argv* as a subprocess and begin pumping its output.

        Never called while holding a DB lock — spawning is fire-and-forget;
        the caller returns the job id immediately and polls for progress.

        Raises ``ValueError`` for a *kind* outside :data:`JOB_KINDS`: an
        unknown kind would sail past :meth:`pipeline_active`'s gate and
        :meth:`running_job_for`'s dedup alike, which is precisely the silent
        divergence the vocabulary constant exists to prevent. Raises
        :class:`JobManagerClosed` once :meth:`shutdown` has run — the child is
        killed and reaped before the raise, never left behind. ``Popen``'s own
        errors pass through untouched (``FileNotFoundError`` for a missing
        executable, ``NotADirectoryError``/``FileNotFoundError`` for a bad
        *cwd*): they happen before the job exists, so nothing is registered.
        """
        if kind not in JOB_KINDS:
            raise ValueError(f"unknown job kind {kind!r} (expected one of {sorted(JOB_KINDS)})")
        proc = subprocess.Popen(
            argv,
            cwd=str(cwd),
            env=env,
            # Never the cockpit's own stdin (module docstring): a prompting child
            # would silently eat the operator's terminal input. EOF makes it
            # fail loudly instead.
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            errors="replace",
            bufsize=1,
        )
        job = Job(
            id=secrets.token_hex(4),
            kind=kind,
            label=label,
            argv=list(argv),
            proc=proc,
            started_at=_now_ms(),
            experiment=experiment,
        )
        with self._lock:
            closed = self._closed
            if not closed:
                self._jobs.append(job)
                if len(self._jobs) > _MAX_JOBS:
                    # Evict the oldest *finished* job — never a running one,
                    # which would orphan its process (untracked by stop()/
                    # shutdown()). The scan walks PAST a running job to find a
                    # finished one behind it; if everything is still running
                    # the registry briefly exceeds the cap rather than losing
                    # track of a live subprocess.
                    for i, old in enumerate(self._jobs):
                        with old.lock:
                            still_running = old.status == "running"
                        if not still_running:
                            self._jobs.pop(i)
                            break
        if closed:
            _terminate_and_reap(proc)
            raise JobManagerClosed(
                f"the job registry is shut down; {kind!r} was not started "
                "(its process has been terminated)"
            )
        threading.Thread(target=_pump, args=(job,), daemon=True).start()
        return job

    def get(self, job_id: str) -> Job | None:
        with self._lock:
            for job in self._jobs:
                if job.id == job_id:
                    return job
        return None

    def set_url(self, job: Job, url: str) -> None:
        """Record the cockpit URL an ``explore`` job reported on its stdout."""
        with job.lock:
            job.url = url

    def snapshot(self, job: Job, offset: int = 0) -> dict[str, Any]:
        """One job's poll reply, lines included.

        ``offset`` / ``next_offset`` are ABSOLUTE line indices over the job's
        whole lifetime (``dropped + buffered``), not indices into the current
        buffer — otherwise a job more verbose than :data:`_MAX_LINES` would pin
        the poller's offset at the buffer length and the stream would go silent
        forever. Lines that already fell off the front are simply gone,
        matching a live terminal's scrollback — and ``dropped``/``truncated``
        say so on the wire, so a log drawer can render "earlier output was
        discarded" instead of inferring it from a hole between ``offset`` and
        ``len(lines)``. (:meth:`list_snapshots` carries neither: it is the
        chip summary, and it has no lines to be missing from.)
        """
        with job.lock:
            lines = list(job.lines)
            dropped = job.dropped
            truncated = job.truncated
            status = job.status
            returncode = job.returncode
            url = job.url
        total = dropped + len(lines)
        start = max(dropped, min(offset, total))
        return {
            "id": job.id,
            "kind": job.kind,
            "label": job.label,
            "experiment": job.experiment,
            "status": status,
            "returncode": returncode,
            "url": url,
            "next_offset": total,
            "dropped": dropped,
            "truncated": truncated,
            "lines": lines[start - dropped :],
        }

    def list_snapshots(self) -> list[dict[str, Any]]:
        """Every job's summary (no ``lines``), newest first."""
        with self._lock:
            jobs = list(self._jobs)
        out = []
        for job in reversed(jobs):
            with job.lock:
                out.append(
                    {
                        "id": job.id,
                        "kind": job.kind,
                        "label": job.label,
                        "experiment": job.experiment,
                        "status": job.status,
                        "returncode": job.returncode,
                        "url": job.url,
                        "started_at": job.started_at,
                        "finished_at": job.finished_at,
                    }
                )
        return out

    def spawn_pipeline(
        self,
        kind: str,
        label: str,
        argv: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        experiment: str | None = None,
    ) -> Job | None:
        """Atomically spawn a pipeline job, or return ``None`` if one is running.

        The busy check and the spawn happen under one gate, so two
        near-simultaneous requests can't both observe "idle" and each start a
        subprocess (the plain check-then-``spawn()`` sequence is a TOCTOU race
        across the server's request threads).

        Raises ``ValueError`` for a *kind* outside :data:`PIPELINE_KINDS` —
        ``explore`` is not gated one-at-a-time and belongs to
        :meth:`spawn_deduped`, so accepting it here would return a job the gate
        never actually protected. (Not :meth:`spawn`: that is the ungated
        primitive, and reaching for it directly is the check-then-act race
        :meth:`spawn_deduped` exists to close.)

        *experiment* is optional here (a run may span a whole selection) but
        pass it whenever the job targets exactly one, so the registry can label
        the job from a field instead of a formatted string.

        Propagates :class:`JobManagerClosed` from :meth:`spawn` on a shut-down
        registry — ``None`` means "busy, try later", which a teardown is not.
        """
        if kind not in PIPELINE_KINDS:
            raise ValueError(
                f"{kind!r} is not a pipeline job kind "
                f"(expected one of {sorted(PIPELINE_KINDS)}); use spawn_deduped() instead"
            )
        with self._gate:
            if self.pipeline_active():
                return None
            return self.spawn(kind, label, argv, cwd=cwd, env=env, experiment=experiment)

    def spawn_deduped(
        self,
        kind: str,
        label: str,
        argv: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        experiment: str,
    ) -> tuple[Job, bool]:
        """The live *kind* job for *experiment*, or a newly spawned one.

        Returns ``(job, created)`` — ``created=False`` means the caller got the
        job that was already running, which is what makes a second Explore click
        reopen the live cockpit instead of racing it.

        The lookup and the spawn happen under the SAME gate
        :meth:`spawn_pipeline` uses, because :meth:`running_job_for` followed by
        :meth:`spawn` is a check-then-act across the server's request threads: a
        double-clicked button (browsers do fire both) would pass the check twice
        and start two ``abk explore`` sessions on one experiment — each holding
        its own startup snapshot of the YAML, so whichever Applies second
        silently drops the other's changes. That is the hazard the dedup exists
        for, so the dedup has to be atomic.

        Raises ``ValueError`` for a *kind* in :data:`PIPELINE_KINDS`: those are
        gated one-at-a-time and must go through :meth:`spawn_pipeline`, so
        deduping one here would hand back a job that gate never saw. Propagates
        :class:`JobManagerClosed` from :meth:`spawn`.
        """
        if kind in PIPELINE_KINDS:
            raise ValueError(
                f"{kind!r} is a pipeline job kind (one at a time, not deduped "
                f"per experiment); use spawn_pipeline() instead"
            )
        with self._gate:
            existing = self.running_job_for(kind, experiment)
            if existing is not None:
                return existing, False
            job = self.spawn(kind, label, argv, cwd=cwd, env=env, experiment=experiment)
            return job, True

    def url_for(self, job: Job) -> str | None:
        """*job*'s recorded cockpit URL, read under its lock.

        The read half of :meth:`set_url`, so a caller never has to reach into
        ``job.lock`` itself (the pump thread writes these fields).
        """
        with job.lock:
            return job.url

    def running_job_for(self, kind: str, experiment: str) -> Job | None:
        """The still-running *kind* job for *experiment*, if any.

        The abkit analog of the donor's ``running_tune_for(metric)``: dedup for
        ``POST /api/explore``, keyed on ``(kind='explore', experiment=name)``
        so a second click reopens the live cockpit instead of racing it. Call it
        through :meth:`spawn_deduped` when the answer decides whether to spawn —
        on its own it is a TOCTOU check.
        """
        with self._lock:
            jobs = list(self._jobs)
        for job in jobs:
            with job.lock:
                if job.kind == kind and job.experiment == experiment and job.status == "running":
                    return job
        return None

    def pipeline_active(self) -> bool:
        """True when a :data:`PIPELINE_KINDS` job (run/unlock/clean) is still running.

        ``explore`` jobs are excluded: multiple concurrent cockpits are fine
        (each explores a different experiment, and explore takes no pipeline
        lock), but Run/Unlock/Clean all mutate the same DB-level state, so the
        dashboard only ever lets one of those run at a time.
        """
        with self._lock:
            jobs = list(self._jobs)
        for job in jobs:
            with job.lock:
                if job.kind in PIPELINE_KINDS and job.status == "running":
                    return True
        return False

    def stop(self, job_id: str) -> bool:
        """Terminate a running job (grace period, then kill). False if not running."""
        job = self.get(job_id)
        if job is None:
            return False
        with job.lock:
            if job.status != "running":
                return False
            job.stop_requested = True
        _terminate_and_reap(job.proc)
        return True

    def wait_for_line(
        self, job: Job, predicate: Callable[[str], bool], timeout: float
    ) -> str | None:
        """Block (polling) until a line matching *predicate* appears, or *timeout*.

        Returns ``None`` on timeout or if the job stops running before a
        matching line appears. Used by ``POST /api/explore`` to wait for the
        ``Explore: <url>`` line ``serve_explore`` echoes.

        The scan position is an ABSOLUTE line index, like
        :meth:`snapshot`'s offsets: tracking an index into the live buffer
        (the donor's shape) stops advancing the moment the buffer saturates at
        :data:`_MAX_LINES`, and the watcher never sees another line. Lines
        dropped before this method got to them are gone, not re-scanned.
        """
        deadline = time.monotonic() + timeout
        checked = 0
        while True:
            with job.lock:
                dropped = job.dropped
                buffered = list(job.lines)
                status = job.status
            total = dropped + len(buffered)
            start = max(checked, dropped)
            new_lines = buffered[start - dropped :] if start < total else []
            checked = total
            for line in new_lines:
                if predicate(line):
                    return line
            if status != "running":
                return None
            if time.monotonic() >= deadline:
                return None
            time.sleep(0.05)

    def shutdown(self) -> None:
        """Terminate every still-running job (grace period, then kill).

        Latches the registry closed in the same critical section that takes
        the snapshot, so a concurrent :meth:`spawn` either landed in this
        snapshot or refuses and kills its own child — no process outlives the
        teardown untracked. Idempotent; a manager stays closed.

        Unlike :meth:`stop`'s per-job grace thread, the grace period here is
        ONE budget shared by every job in iteration order: total teardown
        latency stays bounded by :data:`_STOP_GRACE_SECONDS` (the donor's
        choice, kept), at the price of jobs late in the order getting little
        SIGTERM grace before the SIGKILL. Every one of them is still reaped.
        """
        with self._lock:
            self._closed = True
            jobs = list(self._jobs)
        running = []
        for job in jobs:
            with job.lock:
                if job.status != "running":
                    continue
                # A teardown is as deliberate as a Stop click: without this the
                # pump would read the SIGTERM returncode as a crash and report
                # "failed" for every job the dashboard itself killed.
                job.stop_requested = True
            try:
                job.proc.terminate()
            except Exception:
                pass
            running.append(job)
        deadline = time.monotonic() + _STOP_GRACE_SECONDS
        for job in running:
            remaining = max(0.0, deadline - time.monotonic())
            try:
                job.proc.wait(timeout=remaining)
            except subprocess.TimeoutExpired:
                try:
                    job.proc.kill()
                except Exception:
                    pass
