"""The warden as a resident UDS daemon — the sole executor, reachable cross-process.

A Unix-domain socket (filesystem-permissioned, no port) carries newline-delimited
JSON requests. Exactly one warden per project is enforced by an ``fcntl`` lock on a
pidfile (the singleton); a stale socket from a crashed warden is reclaimed. Because
all callers (CLI / web / agent / CI) reach the *same* warden process, single-flight
holds across every origin.
"""

from __future__ import annotations

import contextlib
import fcntl
import json
import os
import signal
import socket
import socketserver
import threading
from pathlib import Path
from typing import TYPE_CHECKING, Any

from .leases import RunRequest

if TYPE_CHECKING:
    from collections.abc import Callable

    from .broker import Warden


def _dispatch(warden: Warden, req: dict[str, Any]) -> dict[str, Any]:
    op = req.get("op")
    if op == "ping":
        return {"ok": True, "pong": True, "project": warden.project}
    if op == "status":
        return {"ok": True, **warden.status()}
    if op == "run":
        outcome = warden.submit(
            RunRequest(
                run_id=req["run_id"],
                test_id=req["test_id"],
                config=req.get("config", {}),
                origin=req.get("origin", "cli"),
                who=req.get("who", ""),
                independent=req.get("independent", False),
            )
        )
        return {
            "ok": True,
            "status": outcome.status,
            "run_id": outcome.run_id,
            "position": outcome.position,
        }
    if op == "cancel":
        cancelled = warden.cancel(req["run_id"])
        # The broker stays process-free (unit-testable); the serve path owns os.kill.
        # SIGTERM the child if it's running — the lease then releases the normal way.
        if cancelled.pid:
            # already gone? the lease still releases on the child's natural exit
            with contextlib.suppress(OSError):
                os.kill(cancelled.pid, signal.SIGTERM)
        return {
            "ok": True,
            "found": cancelled.found,
            "run_id": cancelled.run_id,
            "state": cancelled.state,
            "pid": cancelled.pid,
        }
    return {"ok": False, "error": f"unknown op {op!r}"}


class _Handler(socketserver.StreamRequestHandler):
    def handle(self) -> None:
        line = self.rfile.readline()
        if not line:
            return
        warden: Warden = self.server.warden  # type: ignore[attr-defined]
        try:
            response = _dispatch(warden, json.loads(line.decode("utf-8")))
        except Exception as exc:  # never let one bad request take the warden down
            response = {"ok": False, "error": str(exc)}
        self.wfile.write((json.dumps(response) + "\n").encode("utf-8"))


class _Server(socketserver.ThreadingUnixStreamServer):
    daemon_threads = True
    allow_reuse_address = True

    def __init__(self, socket_path: str, warden: Warden) -> None:
        self.warden = warden
        super().__init__(socket_path, _Handler)


def is_live(socket_path: str | Path) -> bool:
    """Does a warden currently answer on this socket? (used to detect a stale socket)."""
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
            sock.settimeout(2)
            sock.connect(str(socket_path))
            sock.sendall(b'{"op": "ping"}\n')
            return b"pong" in sock.recv(4096)
    except OSError:
        return False


class WardenDaemon:
    """Owns the singleton lock + the UDS server for one project's warden."""

    def __init__(self, warden: Warden, socket_path: str | Path) -> None:
        self.warden = warden
        self.socket_path = Path(socket_path)
        self.lock_path = self.socket_path.with_suffix(".lock")
        self._server: _Server | None = None
        self._lock_fh: Any = None
        self._heartbeat_thread: threading.Thread | None = None
        self._heartbeat_stop = threading.Event()

    def _acquire_singleton(self) -> None:
        self.socket_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock_fh = self.lock_path.open("w")
        try:
            fcntl.flock(self._lock_fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError as exc:
            self._lock_fh.close()
            self._lock_fh = None
            raise RuntimeError(
                f"a warden is already running for this project ({self.lock_path})"
            ) from exc
        # we hold the singleton: a leftover socket must be stale -> reclaim it
        if self.socket_path.exists():
            self.socket_path.unlink()

    def serve_forever(self, on_ready: Callable[[], None] | None = None) -> None:
        self._acquire_singleton()
        # on_ready runs AFTER the singleton is held (the old warden is proven gone) and
        # BEFORE we accept requests — the crash-recovery hook lives here.
        if on_ready is not None:
            on_ready()
        self._start_heartbeat()
        self._server = _Server(str(self.socket_path), self.warden)
        try:
            self._server.serve_forever()
        finally:
            self.shutdown()

    def start_background(self, on_ready: Callable[[], None] | None = None) -> None:
        """Start serving on a daemon thread (for embedding / tests)."""
        self._acquire_singleton()
        if on_ready is not None:
            on_ready()
        self._start_heartbeat()
        self._server = _Server(str(self.socket_path), self.warden)
        threading.Thread(target=self._server.serve_forever, daemon=True).start()

    def _start_heartbeat(self) -> None:
        """Refresh the lease heartbeats + re-journal on the warden's ``heartbeat_s``
        interval. Only the live (singleton-holding) warden heartbeats, so a stale journal
        always means a dead warden."""
        if self.warden.heartbeat_s <= 0:
            return
        self._heartbeat_stop = threading.Event()

        def beat() -> None:
            while not self._heartbeat_stop.wait(self.warden.heartbeat_s):
                self.warden.heartbeat()

        self._heartbeat_thread = threading.Thread(target=beat, daemon=True)
        self._heartbeat_thread.start()

    def shutdown(self) -> None:
        self._heartbeat_stop.set()
        if self._heartbeat_thread is not None:
            self._heartbeat_thread.join(timeout=1)
            self._heartbeat_thread = None
        if self._server is not None:
            self._server.shutdown()
            self._server.server_close()
            self._server = None
        if self._lock_fh is not None:
            fcntl.flock(self._lock_fh, fcntl.LOCK_UN)
            self._lock_fh.close()
            self._lock_fh = None
        if self.socket_path.exists():
            self.socket_path.unlink()
