"""The warden broker — the sole executor; single-flight enforced by construction.

The warden is one process, so the acquire decision is an in-process critical section
(a lock), with no file-lock TOCTOU:

    free                              -> grant a lease, spawn the runner
    held, same test_id + config_digest -> coalesce (attach watcher; NO second run)
    held, different work               -> queue (FIFO)

Every outcome — acquire / coalesce / queue / release / reclaim — is an audited event,
the visible proof the lock held. TTL/heartbeat is warden-crash recovery, not the
normal path (the warden owns the child, so it observes completion directly).
"""

from __future__ import annotations

import threading
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from .leases import Lease, LeaseRegistry, RunRequest, age_seconds, now_iso

Emit = Callable[[dict[str, Any]], None]
Runner = Callable[[RunRequest, Emit], str]
#: Called once per *executed* run after the runner returns and the lease is
#: released: ``(request, released-lease-snapshot, result)``. Coalesced/queued
#: requests never execute, so they never fire this — they are recorded in the
#: holder's lease audit instead (the lease's ``watchers``). Lets a recorder write
#: the durable run-record without the broker knowing about ``fed-resources/``.
OnComplete = Callable[[RunRequest, "Lease | None", str], None]
#: Called under the broker lock on every lease change (acquire / release / reclaim /
#: pid-update / heartbeat) with a snapshot of the running leases. Lets ``testctl serve``
#: mirror them to the on-disk journal (``live/leases.json``) for crash recovery without
#: the broker knowing about ``fed-resources/``.
Journal = Callable[[list["Lease"]], None]


@dataclass(frozen=True)
class AcquireOutcome:
    status: str  # "started" | "coalesced" | "queued"
    run_id: str  # the run to watch (the in-flight one, for coalesce)
    position: int | None = None


@dataclass(frozen=True)
class CancelOutcome:
    found: bool
    run_id: str
    pid: int | None = None  # the child to signal (the serve path owns os.kill)
    state: str = "absent"  # "running" | "queued" | "absent"


def _noop_emit(_event: dict[str, Any]) -> None:
    return None


def _noop_complete(_request: RunRequest, _lease: Lease | None, _result: str) -> None:
    return None


def _noop_journal(_leases: list[Lease]) -> None:
    return None


class Warden:
    """A per-project resident broker. Construct once; submit requests to it."""

    def __init__(
        self,
        project: str,
        runner: Runner,
        *,
        serialize: bool = True,
        ttl_s: int = 30,
        heartbeat_s: int = 10,
        on_event: Emit | None = None,
        on_complete: OnComplete | None = None,
        journal: Journal | None = None,
    ) -> None:
        self.project = project
        self.runner = runner
        self.ttl_s = ttl_s
        self.heartbeat_s = heartbeat_s
        self.registry = LeaseRegistry(project=project, serialize=serialize)
        self._lock = threading.RLock()
        self._workers: dict[str, threading.Thread] = {}
        self._results: dict[str, str] = {}
        self._on_event: Emit = on_event or _noop_emit
        self._on_complete: OnComplete = on_complete or _noop_complete
        self._journal: Journal = journal or _noop_journal
        #: run_ids an operator asked to cancel — so the natural release records the
        #: authoritative "cancelled" result rather than the killed child's exit code.
        self._cancelled: set[str] = set()

    def _journal_now(self) -> None:
        """Persist the running leases via the journal hook. Caller MUST hold the lock —
        the snapshot must be consistent with the lease change that triggered it."""
        self._journal(list(self.registry.running.values()))

    # --- the acquire decision (the enforcement) ---
    def submit(self, request: RunRequest) -> AcquireOutcome:
        with self._lock:
            key = request.key(self.project)
            held = self.registry.running.get(key)
            if held is None:
                return self._start(request, key)
            if held.test_id == request.test_id and held.config_digest == request.config_digest:
                held.watchers.append(request.run_id)
                self.registry.audit(
                    "coalesced",
                    test_id=request.test_id,
                    origin=request.origin,
                    attached_to=held.run_id,
                    note="single-flight: attached as watcher, did not start a second run",
                )
                return AcquireOutcome("coalesced", held.run_id)
            self.registry.queue.append(request)
            self.registry.audit("queued", test_id=request.test_id, origin=request.origin)
            return AcquireOutcome("queued", request.run_id, position=len(self.registry.queue))

    def _start(self, request: RunRequest, key: str) -> AcquireOutcome:
        stamp = now_iso()
        lease = Lease(
            run_id=request.run_id,
            key=key,
            test_id=request.test_id,
            config_digest=request.config_digest,
            origin=request.origin,
            who=request.who,
            acquired_at=stamp,
            heartbeat_at=stamp,
            ttl_s=self.ttl_s,
        )
        self.registry.running[key] = lease
        self.registry.audit("lease_acquired", test_id=request.test_id, origin=request.origin)
        self._journal_now()
        worker = threading.Thread(target=self._run, args=(request, lease), daemon=True)
        self._workers[request.run_id] = worker
        worker.start()
        return AcquireOutcome("started", request.run_id)

    def _run(self, request: RunRequest, lease: Lease) -> None:
        def emit(event: dict[str, Any]) -> None:
            # the warden owns the lease, so it lifts the authoritative `progress` from
            # the runner's progress telemetry (status() reports it under the lock).
            if event.get("kind") == "progress" and isinstance(event.get("progress"), int | float):
                with self._lock:
                    lease.progress = float(event["progress"])
            # the runner emits the child's pid (on `test_start`); record it on the lease
            # and re-journal so the crash-recovery journal can kill the orphan if this
            # warden dies. The recorder drops `pid` from the NDJSON (not in its grammar).
            pid = event.get("pid")
            if isinstance(pid, int) and not isinstance(pid, bool):
                with self._lock:
                    lease.pid = pid
                    self._journal_now()
            self._on_event({"run_id": request.run_id, **event})

        try:
            result = self.runner(request, emit)
        except Exception as exc:  # the runner is arbitrary; the warden must survive its failure
            result = "error"
            emit({"kind": "error", "msg": str(exc)})
        with self._lock:
            # A cancelled run's child exits non-zero, but the authoritative result is
            # "cancelled" (operator intent), not the kill's exit code.
            self._results[request.run_id] = (
                "cancelled" if request.run_id in self._cancelled else result
            )
            self._cancelled.discard(request.run_id)
            if self.registry.running.get(lease.key) is lease:
                del self.registry.running[lease.key]
            self.registry.audit("lease_released", test_id=request.test_id)
            self._advance(lease.key)
            self._journal_now()
        # Outside the lock so a slow recorder (file I/O) never blocks the broker; the
        # lease is the released snapshot (its watchers carry any coalesced double-launch).
        try:
            self._on_complete(request, lease, result)
        except Exception as exc:  # a recorder fault must never take the warden down
            emit({"kind": "error", "msg": f"on_complete failed: {exc}"})

    def _advance(self, freed_key: str) -> None:
        if freed_key in self.registry.running:
            return
        for request in list(self.registry.queue):
            if request.key(self.project) == freed_key:
                self.registry.queue.remove(request)
                self._start(request, freed_key)
                return

    # --- operator cancel (signal a run; the queue advances only on real exit) ---
    def cancel(self, run_id: str) -> CancelOutcome:
        """Cancel a run by id — a running lease or a still-queued request.

        For a *running* lease this records the intent + audits ``cancelled`` and returns
        the child ``pid`` for the serve path to signal; it deliberately does **not** drop
        the lease or advance the queue here. The lease releases through the normal
        ``_run`` path when the signalled child actually exits, so single-flight holds —
        no queued run starts while the cancelled child is still alive. A *queued* request
        is simply dropped (no child exists). An unknown id is reported, never an error.
        """
        with self._lock:
            for lease in self.registry.running.values():
                if lease.run_id == run_id:
                    self._cancelled.add(run_id)
                    self.registry.audit(
                        "cancelled",
                        test_id=lease.test_id,
                        origin=lease.origin,
                        note="operator cancel — child signalled; lease releases when it exits",
                    )
                    self._journal_now()
                    return CancelOutcome(True, run_id, pid=lease.pid, state="running")
            for request in list(self.registry.queue):
                if request.run_id == run_id:
                    self.registry.queue.remove(request)
                    self.registry.audit(
                        "cancelled",
                        test_id=request.test_id,
                        origin=request.origin,
                        note="operator cancel — dequeued before start",
                    )
                    self._journal_now()
                    return CancelOutcome(True, run_id, pid=None, state="queued")
            return CancelOutcome(False, run_id, pid=None, state="absent")

    # --- crash recovery (a *new* warden reclaims a dead warden's stale leases) ---
    def reclaim_stale(self, *, now: str | None = None) -> list[str]:
        stamp = now or now_iso()
        reclaimed: list[str] = []
        with self._lock:
            for key, lease in list(self.registry.running.items()):
                if age_seconds(lease.heartbeat_at, stamp) > lease.ttl_s:
                    del self.registry.running[key]
                    self.registry.audit(
                        "lease_reclaimed",
                        test_id=lease.test_id,
                        note=f"stale lease reclaimed (heartbeat aged past {lease.ttl_s}s)",
                    )
                    reclaimed.append(lease.run_id)
                    self._advance(key)
            if reclaimed:
                self._journal_now()
        return reclaimed

    def heartbeat(self) -> None:
        """Refresh every running lease's heartbeat + re-journal — the liveness signal.

        Only the live warden (it holds the daemon's singleton) heartbeats, so a stale
        on-disk journal always means a dead warden. Driven by the daemon on the
        ``heartbeat_s`` interval; a no-op when nothing is running.
        """
        stamp = now_iso()
        with self._lock:
            if not self.registry.running:
                return
            for lease in self.registry.running.values():
                lease.heartbeat_at = stamp
            self._journal_now()

    def recover_from_journal(self, entries: list[dict[str, Any]]) -> list[str]:
        """Reclaim the leases a dead warden left in its on-disk journal (crash recovery).

        The daemon's ``fcntl`` singleton has already proven the previous warden is gone,
        so a leftover journal entry is unconditionally its orphan — reclaim it regardless
        of TTL. Each reclaim is an audited ``lease_reclaimed`` (the visible proof of the
        recovery) and the freed key is advanced so a queued request can start. This is
        process-free — killing the orphaned child by ``pid`` is the serve path's job
        (it owns ``os.kill``) — so it is unit-testable without a real process.
        """
        reclaimed: list[str] = []
        with self._lock:
            for entry in entries:
                key = str(entry.get("key", ""))
                run_id = str(entry.get("run_id", ""))
                if not key or not run_id:
                    continue  # a malformed journal entry recovers nothing; skip it
                self.registry.running.pop(key, None)
                self.registry.audit(
                    "lease_reclaimed",
                    test_id=entry.get("test_id"),
                    note="recovered a dead warden's lease from the journal (crash reclaim)",
                )
                reclaimed.append(run_id)
                self._advance(key)
            if reclaimed:
                self._journal_now()
        return reclaimed

    # --- read surface ---
    def result(self, run_id: str) -> str | None:
        return self._results.get(run_id)

    def status(self, *, broker: str = "up") -> dict[str, Any]:
        with self._lock:
            return self.registry.to_status(broker=broker)

    def wait_idle(self, timeout: float = 5.0) -> bool:
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            with self._lock:
                if not self.registry.running and not self.registry.queue:
                    return True
            time.sleep(0.01)
        return False
