"""Cross-process mirror for the comms control plane (SQLite, stdlib-only).

The JobStore and Bus are per-process — correct for one gunicorn worker, but
the dev service runs ``--workers 3``: a cancel POST lands on process B while
process A serves the stream, and B's store has never heard of the job. This
mirror is the smallest primitive that fixes that without a broker or an
external dependency:

    - every job transition is upserted here (JSON snapshot + cancel flag),
    - a cancel for a job we don't hold locally sets the flag on the row,
    - the process that owns the stream notices the flag on its next token
      (throttled read on the hot path) and fires its local cancel handle,
    - queue views merge live rows from sibling processes into their snapshot.

SQLite in WAL mode handles concurrent single-row writes from a handful of
processes easily. Every operation opens a short-lived connection — no shared
handles across threads, no pooling to get wrong. All methods are best-effort:
a mirror failure degrades to per-process behavior (exactly today's world),
never breaks a chat. After MAX_FAILURES consecutive errors the mirror
disables itself loudly rather than taxing every token with a doomed write.

The db path must be shared by all processes of one service: HUGPY_COMMS_DB
if set, else a per-user file under XDG_RUNTIME_DIR (or /tmp). Central and a
worker agent on the same box sharing one file is harmless — job ids are
uuids and each process still only cancels what it holds a handle for.
"""
from __future__ import annotations

import errno
import json
import logging
import os
import random
import sqlite3
import threading
import time
from typing import Any, Callable, Optional, TypeVar

logger = logging.getLogger(__name__)

_T = TypeVar("_T")

# ---------------------------------------------------------------------------
# EMFILE burst hardening (incident 2026-07-23)
# ---------------------------------------------------------------------------
# On a restart, all gunicorn workers spawn at once and simultaneously open
# store-backed files on the virtiofs mount (/mnt/llm_storage). Under that
# concurrent-open burst the mount MOMENTARILY returns EMFILE ("Too many open
# files") — this is the mount degrading under load, NOT a real per-process
# fd-limit breach (each proc held ~50 of 65536). SQLite surfaces the very same
# transient as OperationalError("unable to open database file"). Both settle in
# well under a second. This helper retries the store-open with exponential
# backoff + small jitter so a burst settles instead of surfacing a 500.
#
# It is DELIBERATELY narrow: only EMFILE OSErrors and the one sqlite "unable to
# open database file" message are retryable. Every other error (a real
# permission fault, a corrupt db, a genuinely different sqlite error) propagates
# immediately — a burst-settle retry must never paper over a real fault.


def _is_emfile(exc: BaseException) -> bool:
    """True for the transient store-open faults the mount throws under an
    open burst: an OSError with errno EMFILE, or the sqlite OperationalError
    whose message is 'unable to open database file' (same root cause, surfaced
    through the sqlite open path). Nothing else is retryable."""
    if isinstance(exc, sqlite3.OperationalError):
        return "unable to open database file" in str(exc).lower()
    if isinstance(exc, OSError):
        return exc.errno == errno.EMFILE
    return False


def retry_on_emfile(
    fn: Callable[[], _T],
    *,
    attempts: int = 5,
    base_delay: float = 0.05,
    max_delay: float = 0.8,
    sleep: Callable[[float], None] = time.sleep,
    rng: Optional[Callable[[], float]] = random.random,
) -> _T:
    """Call ``fn()`` and, on a transient EMFILE / 'unable to open database file'
    store-open fault, retry with exponential backoff + jitter; re-raise the LAST
    error once ``attempts`` are exhausted.

    Only the transient store-open faults (see ``_is_emfile``) are retried — any
    other exception propagates immediately, unretried. Backoff is
    ``min(max_delay, base_delay * 2**i)`` plus up to one base_delay of jitter
    (``rng()`` in [0,1)); pass ``rng=None`` for deterministic (jitter-free)
    backoff in tests, and inject ``sleep`` to make tests fast. ``attempts`` is
    the total number of tries (>=1), so there are at most ``attempts - 1``
    sleeps.
    """
    if attempts < 1:
        attempts = 1
    last: BaseException
    for i in range(attempts):
        try:
            return fn()
        except BaseException as exc:  # noqa: BLE001 — re-raised unless retryable
            if not _is_emfile(exc):
                raise
            last = exc
            if i >= attempts - 1:
                break
            delay = min(max_delay, base_delay * (2 ** i))
            if rng is not None:
                delay += base_delay * rng()
            logger.warning(
                "store-open EMFILE burst (attempt %d/%d): %s — backing off %.3fs",
                i + 1, attempts, exc, delay,
            )
            sleep(delay)
    raise last

_SCHEMA = """
CREATE TABLE IF NOT EXISTS jobs (
    id               TEXT PRIMARY KEY,
    data             TEXT NOT NULL,
    status           TEXT NOT NULL,
    kind             TEXT NOT NULL DEFAULT 'chat',
    cancel_requested INTEGER NOT NULL DEFAULT 0,
    updated          REAL NOT NULL,
    -- Last REAL forward-progress epoch (Job.progressed_at), distinct from
    -- `updated` which any write bumps. The orphan-sweep in prune() ages on
    -- THIS clock so a wedged render (updated by mere views/recomputes) still
    -- goes reapable, while a truly-progressing stream never does.
    progressed_at    REAL
);
"""

# Statuses considered "active" for the wedged-orphan sweep. Mirrors the
# _STALL_ACTIVE spirit in jobs.py: only a job that claims to be doing work can
# be "wedged". A pending/queued job is starved (waiting its turn), not wedged,
# so it is NEVER reaped by the progress-sweep.
_ORPHAN_ACTIVE = ("processing", "streaming", "running")

# Pending-orphan expiry (slice 9, defect 2): a `pending` job with no worker and
# no forward progress this long was NEVER dispatched (model unresolvable, or no
# capable worker) — it will never run, so it transitions to a terminal `expired`
# state. Aged on `progressed_at` (the movement-only clock — NEVER `updated`,
# which view/recompute writes bump; that resurrection bug was fixed once and must
# not return). Env-overridable; default 30 min.
def _pending_expiry_seconds() -> float:
    raw = (os.environ.get("HUGPY_JOB_PENDING_EXPIRY_SECONDS") or "").strip()
    if not raw:
        return 1800.0
    try:
        v = float(raw)
        return v if v > 0 else 1800.0
    except ValueError:
        return 1800.0

MAX_FAILURES = 5


def default_db_path() -> str:
    env = (os.environ.get("HUGPY_COMMS_DB") or "").strip()
    if env:
        return env
    base = os.environ.get("XDG_RUNTIME_DIR") or "/tmp"
    return os.path.join(base, f"hugpy-comms-{os.getuid()}.db")


class SqliteMirror:
    def __init__(self, path: Optional[str] = None,
                 retain_secs: float = 600.0) -> None:
        self.path = path or default_db_path()
        self.retain_secs = retain_secs
        self._failures = 0
        self._disabled = False
        self._init_lock = threading.Lock()
        self._initialized = False

    # -- plumbing ------------------------------------------------------------
    def _connect(self) -> sqlite3.Connection:
        # sqlite3.connect is the store-open point that throws the transient
        # EMFILE / 'unable to open database file' under a restart burst — retry
        # just the open, then run the (handle-local) PRAGMAs normally.
        conn = retry_on_emfile(lambda: sqlite3.connect(self.path, timeout=2.0))
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.execute("PRAGMA busy_timeout=2000")
        return conn

    def _ensure(self) -> bool:
        if self._disabled:
            return False
        if self._initialized:
            return True
        with self._init_lock:
            if self._initialized:
                return True
            try:
                with self._connect() as conn:
                    conn.execute(_SCHEMA)
                    self._migrate(conn)
                self._initialized = True
                return True
            except Exception as exc:
                self._note_failure("init", exc)
                return False

    def _migrate(self, conn: sqlite3.Connection) -> None:
        """Idempotent, safe on an existing populated DB. Adds the
        `progressed_at` column to pre-existing `jobs` tables that predate it
        (a plain CREATE ... IF NOT EXISTS won't alter an already-present
        table). Guarded by a table_info probe so a second run is a no-op and a
        DB already carrying the column never raises."""
        cols = {row[1] for row in conn.execute("PRAGMA table_info(jobs)")}
        if "progressed_at" not in cols:
            # NULL default: old rows have no known progress clock, so they are
            # fail-open (never reaped by the progress-sweep) until a fresh
            # upsert stamps a real progressed_at.
            conn.execute("ALTER TABLE jobs ADD COLUMN progressed_at REAL")

    def _note_failure(self, op: str, exc: Exception) -> None:
        self._failures += 1
        if self._failures >= MAX_FAILURES and not self._disabled:
            self._disabled = True
            logger.error("comms mirror DISABLED after %d failures "
                         "(last: %s during %s) — cross-process cancel/queue "
                         "degrade to per-process until restart",
                         self._failures, exc, op)
        else:
            logger.warning("comms mirror %s failed: %s", op, exc)

    def _ok(self) -> None:
        self._failures = 0

    # -- writes --------------------------------------------------------------
    def upsert(self, job_dict: dict[str, Any]) -> None:
        """Mirror a job snapshot. cancel_requested is only ever raised here,
        never lowered (a resurrection lowers it explicitly via clear_cancel) —
        so a racing sibling's cancel can't be lost under a stale local write."""
        if not self._ensure():
            return
        try:
            # progressed_at is the movement-only clock (jobs.py bumps it on real
            # forward progress, NOT on log-tail/view/recompute writes). Persist
            # it as a real column so prune()'s orphan-sweep can age on it. A
            # missing/garbage value degrades to NULL -> fail-open (never reaped).
            prog = job_dict.get("progressed_at")
            try:
                prog = float(prog) if prog is not None else None
            except (TypeError, ValueError):
                prog = None
            with self._connect() as conn:
                conn.execute(
                    "INSERT INTO jobs (id, data, status, kind, cancel_requested, updated, progressed_at) "
                    "VALUES (?, ?, ?, ?, ?, ?, ?) "
                    "ON CONFLICT(id) DO UPDATE SET "
                    "  data=excluded.data, status=excluded.status, "
                    "  kind=excluded.kind, "
                    "  cancel_requested=MAX(jobs.cancel_requested, excluded.cancel_requested), "
                    "  updated=excluded.updated, "
                    "  progressed_at=excluded.progressed_at",
                    (str(job_dict.get("id")), json.dumps(job_dict),
                     str(job_dict.get("status") or "pending"),
                     str(job_dict.get("kind") or "chat"),
                     1 if job_dict.get("cancel_requested") else 0,
                     time.time(), prog))
            self._ok()
        except Exception as exc:
            self._note_failure("upsert", exc)

    def request_cancel(self, job_id: str) -> bool:
        """Raise the cancel flag on a row we may not hold locally. Returns
        True if a live (non-terminal) row existed to flag."""
        if not self._ensure():
            return False
        try:
            with self._connect() as conn:
                cur = conn.execute(
                    "UPDATE jobs SET cancel_requested=1, updated=? "
                    "WHERE id=? AND status NOT IN ('done','cancelled','failed')",
                    (time.time(), job_id))
            self._ok()
            return cur.rowcount > 0
        except Exception as exc:
            self._note_failure("request_cancel", exc)
            return False

    def force_terminal(self, job_id: str, status: str, message: str = "") -> bool:
        """Force a LIVE mirror row terminal (slice 9) — the authoritative cancel /
        orphan-expiry path for a row NO process holds in memory (the immortal
        pending job that survived restarts). Updates the persisted `status`
        column AND rewrites the row's JSON `data` blob so /llm/jobs reads the
        terminal state consistently. Only acts on a row that is NOT already
        terminal (first-terminal-wins); returns True if a live row was retired.

        Goes through the store's own connection (never a caller-side raw sqlite
        channel) — the DB write discipline stays in one place."""
        if not self._ensure():
            return False
        term = str(status or "cancelled")
        try:
            with self._connect() as conn:
                row = conn.execute(
                    "SELECT data FROM jobs WHERE id=? "
                    "AND status NOT IN ('done','cancelled','failed','expired')",
                    (job_id,)).fetchone()
                if not row:
                    return False           # unknown or already terminal
                try:
                    data = json.loads(row[0]) if row[0] else {}
                except Exception:
                    data = {}
                data["id"] = job_id
                data["status"] = term
                if message:
                    data["message"] = message
                data["ended_ts"] = data.get("ended_ts") or time.time()
                if term == "cancelled":
                    data["cancel_requested"] = True
                conn.execute(
                    "UPDATE jobs SET status=?, data=?, updated=? WHERE id=?",
                    (term, json.dumps(data), time.time(), job_id))
            self._ok()
            return True
        except Exception as exc:
            self._note_failure("force_terminal", exc)
            return False

    def expire_pending_orphans(self) -> list[str]:
        """Transition never-dispatched pending jobs to terminal `expired` (slice 9,
        defect 2). A `pending` row with no worker whose `progressed_at` is older
        than the pending-expiry threshold was never picked up (unresolvable model
        or no capable worker); it will never run. Aged on `progressed_at` — the
        movement-only clock, NEVER `updated` (view/recompute writes bump `updated`,
        and aging on it is exactly the resurrection bug already fixed once). A NULL
        progressed_at is fail-open (never expired here — we only act when we
        positively know it is old). Rewrites the JSON `data` blob so /llm/jobs
        reads the honest terminal state + message. Returns the expired ids."""
        if not self._ensure():
            return []
        cutoff = time.time() - _pending_expiry_seconds()
        msg = ("never dispatched — model unresolvable or no capable worker "
               "(auto-expired)")
        expired: list[str] = []
        try:
            with self._connect() as conn:
                rows = conn.execute(
                    "SELECT id, data FROM jobs WHERE status='pending' "
                    "AND progressed_at IS NOT NULL AND progressed_at < ?",
                    (cutoff,)).fetchall()
                for job_id, data_json in rows:
                    try:
                        data = json.loads(data_json) if data_json else {}
                    except Exception:
                        data = {}
                    # A pending job that somehow has a worker assigned is being
                    # dispatched — leave it (belt-and-braces; the query already
                    # only takes pending, but a worker means an owner exists).
                    if data.get("worker"):
                        continue
                    data["id"] = job_id
                    data["status"] = "expired"
                    data["message"] = msg
                    data["ended_ts"] = data.get("ended_ts") or time.time()
                    conn.execute(
                        "UPDATE jobs SET status='expired', data=?, updated=? "
                        "WHERE id=?",
                        (json.dumps(data), time.time(), job_id))
                    expired.append(job_id)
            self._ok()
        except Exception as exc:
            self._note_failure("expire_pending_orphans", exc)
        return expired

    def clear_cancel(self, job_id: str) -> None:
        """Deliberate resurrection (download retry) starts a fresh run."""
        if not self._ensure():
            return
        try:
            with self._connect() as conn:
                conn.execute("UPDATE jobs SET cancel_requested=0, updated=? "
                             "WHERE id=?", (time.time(), job_id))
            self._ok()
        except Exception as exc:
            self._note_failure("clear_cancel", exc)

    def prune(self) -> None:
        if not self._ensure():
            return
        try:
            now = time.time()
            with self._connect() as conn:
                # (1) Terminal rows past the retain window — correctly keyed on
                # `updated` (any-write clock): once terminal, no more real
                # progress happens, so time-since-any-write is the right age.
                conn.execute(
                    "DELETE FROM jobs WHERE "
                    "status IN ('done','cancelled','failed') AND updated < ?",
                    (now - self.retain_secs,))
                # (2) Wedged-orphan sweep: a job whose owner died mid-stream
                # never goes terminal. It MUST age on real forward-progress
                # silence (`progressed_at`), NOT on `updated` — `updated` is
                # bumped by mere views/sibling-recomputes/API-restart re-reads,
                # which kept these rows immortal. Only reap an ACTIVE row whose
                # progressed_at is genuinely ancient. NULL progressed_at is
                # fail-open (never reaped here) — we only reap when we
                # positively know it's old. A healthy long render bumps
                # progressed_at on every stage, so this window of TOTAL
                # progress silence = genuinely wedged.
                orphan_secs = max(self.retain_secs * 6, 3600.0)
                placeholders = ",".join("?" for _ in _ORPHAN_ACTIVE)
                conn.execute(
                    "DELETE FROM jobs WHERE "
                    f"status IN ({placeholders}) "
                    "AND progressed_at IS NOT NULL AND progressed_at < ?",
                    (*_ORPHAN_ACTIVE, now - orphan_secs))
            self._ok()
        except Exception as exc:
            self._note_failure("prune", exc)

    # -- reads ---------------------------------------------------------------
    def cancel_requested(self, job_id: str) -> bool:
        if not self._ensure():
            return False
        try:
            with self._connect() as conn:
                row = conn.execute(
                    "SELECT cancel_requested FROM jobs WHERE id=?",
                    (job_id,)).fetchone()
            self._ok()
            return bool(row and row[0])
        except Exception as exc:
            self._note_failure("cancel_requested", exc)
            return False

    def flagged_ids(self) -> set:
        """Ids with the cancel flag raised — the owner-side watcher intersects
        these with its local live jobs."""
        if not self._ensure():
            return set()
        try:
            with self._connect() as conn:
                rows = conn.execute(
                    "SELECT id FROM jobs WHERE cancel_requested=1").fetchall()
            self._ok()
            return {r[0] for r in rows}
        except Exception as exc:
            self._note_failure("flagged_ids", exc)
            return set()

    def live_rows(self) -> list[dict[str, Any]]:
        """Snapshots of every live job any process mirrored here."""
        if not self._ensure():
            return []
        try:
            with self._connect() as conn:
                rows = conn.execute(
                    "SELECT data FROM jobs WHERE "
                    "status NOT IN ('done','cancelled','failed')").fetchall()
            self._ok()
        except Exception as exc:
            self._note_failure("live_rows", exc)
            return []
        out = []
        for (data,) in rows:
            try:
                out.append(json.loads(data))
            except Exception:
                continue
        return out

    def terminal_rows(self, kinds: tuple) -> list[dict[str, Any]]:
        """Snapshots of TERMINAL jobs of the given kinds — the additive
        complement to live_rows() (which excludes terminal by design). Callers
        pass a strict kind allowlist (MEDIA_KINDS) so this surfaces a sibling's
        finished media job cross-process WITHOUT touching chat/download terminal
        behavior. Empty kinds -> nothing (never a full-table terminal scan)."""
        kinds = tuple(kinds or ())
        if not kinds:
            return []
        if not self._ensure():
            return []
        try:
            # Placeholders are '?' derived from the count only — the kind values
            # themselves are always bound parameters (no SQL injection surface).
            placeholders = ",".join("?" for _ in kinds)
            with self._connect() as conn:
                rows = conn.execute(
                    "SELECT data FROM jobs WHERE "
                    "status IN ('done','cancelled','failed') "
                    f"AND kind IN ({placeholders})", kinds).fetchall()
            self._ok()
        except Exception as exc:
            self._note_failure("terminal_rows", exc)
            return []
        out = []
        for (data,) in rows:
            try:
                out.append(json.loads(data))
            except Exception:
                continue
        return out
