cleveragents-core/tools/_pr_state_cache.py

"""SQLite-backed store for the PR State Warmer's snapshot of open PRs.

The :mod:`pr_state_warmer` sidecar polls Forgejo's ``/pulls?state=open``
endpoint every ``PR_STATE_WARMER_INTERVAL_S`` seconds (default 30s),
walks every page, and writes the full PR object (number, head.sha,
labels, body, mergeable, updated_at, etc.) here. Dispatchers READ from
this cache instead of calling Forgejo per cycle.

v3 (2026-05-17): adds ``comments_refreshed_updated_at`` to drive the
warmer's comments-cache refresh queue from persistent state, not
from the per-cycle ``upsert_prs`` delta. Stores the PR's
``updated_at`` value at the time of the refresh — pending refreshes
are then ``WHERE updated_at != COALESCE(comments_refreshed_updated_at, '')``.
Using the PR's own ``updated_at`` (string equality) rather than a
wall-clock stamp avoids two failure modes: (1) every cycle's
unchanged-upsert advances ``last_seen_at`` which would trigger
spurious refreshes, and (2) Forgejo's ``updated_at`` format and
our wall-clock format have different string-sort semantics. The
change is ADDITIVE — the migration runs ``ALTER TABLE ADD COLUMN``
rather than DROP/CREATE so a v2 → v3 upgrade preserves data and a
deploy with version skew can co-exist (a v2 reader sees the table
without the new column, which is harmless because v2 readers never
SELECT it).

Why this design
---------------

The dispatcher's per-cycle ``/pulls`` call was both:

1. **Truncated at 50 PRs** — single page, no pagination. Past 50 open
   PRs the 51st+ were silently invisible.
2. **Flaky** — Forgejo's response builds 16 head+base repo+owner
   profiles inline (127 KB for 8 PRs). First request after the
   server's hot-cache window expires takes 24-30s; dispatcher's
   short timeout fires and serves stale.

The warmer eliminates both problems: pagination is a warmer-side
detail; Forgejo's hot cache stays permanently primed by the 30s
poll cadence (shorter than the ~60-90s cache TTL), so cold rebuilds
collapse to ~1/day. Dispatchers do a local SQLite read (~0ms).

Storage
-------

SQLite at ``/tmp/cleveragents-pr-state/state.sqlite3`` (overridable
via ``PR_STATE_CACHE_DIR``). WAL mode so the warmer-writer and
dispatcher-readers can coexist without lock contention.

Schema::

    CREATE TABLE pr_state (
      owner           TEXT NOT NULL,    -- repo owner (e.g. 'drew')
      repo            TEXT NOT NULL,    -- repo name (e.g. 'cleveragents-core')
      number          INTEGER NOT NULL, -- PR number within (owner, repo)
      body_json       TEXT NOT NULL,    -- full PR object as JSON
      head_sha        TEXT NOT NULL,
      state           TEXT NOT NULL,    -- 'open' / 'closed' / 'merged'
      labels_json     TEXT NOT NULL,    -- JSON list of label-name strings
      updated_at      TEXT NOT NULL,    -- last-modified at Forgejo
      first_seen_at   TEXT NOT NULL,    -- when the warmer first added the row
      last_seen_at    TEXT NOT NULL,    -- last warmer cycle that found it in /pulls
      vanished_at     TEXT,             -- NULL while open; set when removed from /pulls
      PRIMARY KEY (owner, repo, number)
    );

The ``(owner, repo)`` discriminator prevents a warmer/dispatcher
cfg mismatch (e.g. running tests against the canonical repo while
the prod warmer wrote rows for the fork) from silently serving the
wrong-repo PRs. Every read takes ``(owner, repo)`` and queries
within that scope only.

Vanish semantics: when a PR disappears from the warmer's enumeration
(closed/merged externally, or the operator deleted it), the row is
NOT immediately deleted — instead ``vanished_at`` is stamped on the
first cycle that misses it AND the row keeps for a grace period
(default 7 days). This lets late-arriving dispatcher cycles still
see the PR's last known state for telemetry/audit purposes.

Errors
------

Operational errors raise :class:`PRStateCacheError`. SQLite errors
propagate as ``sqlite3.DatabaseError`` so a corrupt DB is
distinguishable from a missing row.
"""
from __future__ import annotations

import datetime as _dt
import fcntl
import functools
import json
import logging
import os
import sqlite3
import threading
import time as _time
from pathlib import Path
from typing import Any

_logger = logging.getLogger("pr_state_cache")

# v2 added the (owner, repo) primary-key prefix so a warmer/dispatcher
# cfg mismatch can't silently serve wrong-repo PRs.
# v3 added comments_refreshed_at to drive deferral state from
# persistent storage instead of in-memory.
# Bump when changing schema. Migration policy below for upgrade paths.
SCHEMA_VERSION = 3

_DEFAULT_CACHE_DIR = Path("/tmp/cleveragents-pr-state")
_CACHE_DIR_ENV = "PR_STATE_CACHE_DIR"
_DISABLE_ENV = "PR_STATE_CACHE_DISABLE"

# How long a vanished row sticks around before janitor sweeps it.
# A week is generous — most operators close PRs and don't care, but
# audit/telemetry walks may want to see the last-known state of a
# recently-closed PR.
_VANISHED_GRACE_S_DEFAULT = 7 * 24 * 3600
_VANISHED_GRACE_S_ENV = "PR_STATE_CACHE_VANISHED_GRACE_S"


class PRStateCacheError(RuntimeError):
    """Operational error from the PR-state cache."""


def cache_dir() -> Path:
    return Path(os.environ.get(_CACHE_DIR_ENV) or str(_DEFAULT_CACHE_DIR))


def cache_path() -> Path:
    return cache_dir() / "state.sqlite3"


def is_disabled() -> bool:
    raw = os.environ.get(_DISABLE_ENV, "").strip().lower()
    return raw in {"1", "true", "yes", "on"}


def _vanished_grace_s() -> int:
    raw = os.environ.get(_VANISHED_GRACE_S_ENV)
    if raw:
        try:
            return max(1, int(raw))
        except ValueError:
            pass
    return _VANISHED_GRACE_S_DEFAULT


def _now() -> str:
    return _dt.datetime.now(_dt.timezone.utc).isoformat()


def _normalize_updated_at(value: str) -> str:
    """Strip sub-second precision and canonicalize the timezone marker
    so two values representing the same instant compare equal.

    Why: ``comments_refreshed_updated_at`` is matched against
    ``pr_state.updated_at`` via string equality in the pending-refresh
    query. Forgejo's ``updated_at`` can emit different representations
    of the same instant across versions / proxies (``Z`` vs ``+00:00``,
    microseconds vs whole-second). Without normalization a single
    representation flap triggers infinite re-refresh for every PR.

    Conservative: input that doesn't look like an ISO-8601 timestamp
    is returned untouched."""
    if not value:
        return value
    t_pos = value.find("T")
    if t_pos < 0:
        return value
    # Find the timezone suffix and split.
    tz_pos = -1
    for sep in ("+", "-"):
        candidate = value.find(sep, t_pos)
        if candidate > t_pos:
            tz_pos = candidate
            break
    if tz_pos > 0:
        dt_part = value[:tz_pos]
        tz_part = "+00:00" if value[tz_pos:] in ("+00:00", "+0000") else value[tz_pos:]
    elif value and value[-1] in ("Z", "z"):
        dt_part = value[:-1]
        tz_part = "+00:00"  # canonicalize Z → +00:00
    else:
        dt_part = value
        tz_part = ""
    dot = dt_part.find(".")
    if dot >= 0:
        dt_part = dt_part[:dot]
    return dt_part + tz_part


# SQLite caps placeholders per statement at SQLITE_MAX_VARIABLE_NUMBER
# (32766 on modern builds). ``mark_vanished`` builds an IN-clause from
# the seen-numbers set; if the warmer's pagination cap is raised past
# ~650 pages × 50/page = 32 500 PRs, a single IN-clause would overflow.
# Chunk well under the limit (3 reserved slots for owner/repo/now).
_MARK_VANISHED_CHUNK = 32_000

_MIGRATION_LOCK_FILENAME = "migration.lock"
# Bounded retry on the cross-process migration lock. A stale lock
# (previous holder SIGKILL'd) would block every subsequent `_connect()`
# forever under unbounded `LOCK_EX`. With LOCK_NB + retry we surface
# a clear timeout error after N seconds instead.
_MIGRATION_LOCK_TIMEOUT_S = 30
_MIGRATION_LOCK_POLL_S = 0.5

# Per-process migration state. Schema migration runs at most ONCE per
# process — the first connect performs the migration if needed and
# stamps ``_initialized = True``. Subsequent connects skip the
# migration block entirely.
#
# Thread safety: a ``threading.Lock`` guards the check/set of
# ``_initialized`` so two threads in the same process can't both
# observe "not initialised", both re-migrate, and stamp the flag
# while a third thread sees a partially-migrated state. Used in
# concert with the cross-process file lock.
#
# Cross-process safety: an OS-level ``fcntl.flock`` on
# ``<cache_dir>/migration.lock`` serializes the migration window so N
# dispatcher processes + the warmer can't race destructive DDL. The
# migration policy is also restricted (see ``_run_migration_locked``)
# to ADDITIVE changes whenever possible — a DROP only happens on the
# pre-v2 cases where the column shape is fundamentally incompatible.
#
# Re-heal: if a public-API call gets ``sqlite3.OperationalError`` for
# ``no such table`` or ``no such column``, we clear ``_initialized``
# and retry once. Catches the externally-replaced/corrupted-DB case
# without requiring a process restart. A ``sqlite3.DatabaseError``
# (file corruption) also triggers re-heal but additionally quarantines
# the broken file.
_initialized = False
_init_lock = threading.Lock()


def _connect() -> sqlite3.Connection:
    """Open a SQLite connection, running schema migration ONCE per
    process. Subsequent calls in the same process skip the migration
    block.

    The migration block is guarded by:
      - ``_init_lock`` (in-process): ensures only one thread per
        process runs the migration body; other threads block until
        it returns.
      - ``fcntl.flock`` (cross-process, inside ``_run_migration_locked``):
        ensures only one process across the host runs destructive DDL.
    """
    global _initialized
    target = cache_path()
    target.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(
        str(target),
        timeout=10.0,
        isolation_level=None,  # autocommit; explicit BEGIN/COMMIT below
    )
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = NORMAL")
    if not _initialized:
        with _init_lock:
            # Re-check under the lock — another thread might have
            # initialised while we waited.
            if not _initialized:
                _run_migration_locked(conn, target.parent)
                _initialized = True
    return conn


def _run_migration_locked(conn: sqlite3.Connection, lock_dir: Path) -> None:
    """Acquire an exclusive file lock and run migrations.

    Policy:
      - ``on_disk == SCHEMA_VERSION``: no-op (another process raced
        ahead while we waited for the lock).
      - ``on_disk == 0`` (pre-versioning): DROP + CREATE. The pre-v2
        column shape is fundamentally incompatible (no ``owner`` /
        ``repo`` columns) so an ALTER cannot preserve the rows
        meaningfully.
      - ``0 < on_disk < SCHEMA_VERSION``: additive ALTER TABLE ADD
        COLUMN per known step. Preserves data across upgrade.
      - ``on_disk > SCHEMA_VERSION``: REFUSE to mutate. A newer
        version has already touched the DB; downgrading would lose
        data the newer process expects. Log + leave alone — caller
        will likely see a schema-shape error and we let it surface
        rather than papering over with a DROP.
    """
    lock_path = lock_dir / _MIGRATION_LOCK_FILENAME
    with open(lock_path, "w") as lock_handle:
        # Bounded retry instead of unbounded LOCK_EX: a stale lock
        # (previous holder SIGKILL'd, NFS hang, etc.) would otherwise
        # block every subsequent _connect() forever with no signal.
        # LOCK_NB + poll gives us a clear timeout error and lets the
        # operator see WHICH process is stuck.
        deadline = _time.monotonic() + _MIGRATION_LOCK_TIMEOUT_S
        while True:
            try:
                fcntl.flock(
                    lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB,
                )
                break
            except BlockingIOError:
                if _time.monotonic() >= deadline:
                    # Read holder PID from the lock file if any.
                    try:
                        holder = lock_path.read_text().strip() or "(unknown)"
                    except OSError:
                        holder = "(read failed)"
                    raise PRStateCacheError(
                        f"pr_state cache migration lock at {lock_path} "
                        f"held by another process (pid={holder}) for "
                        f">{_MIGRATION_LOCK_TIMEOUT_S}s; refusing to "
                        f"wait further. Inspect with: lsof {lock_path}"
                    ) from None
                _time.sleep(_MIGRATION_LOCK_POLL_S)
        # Stamp our PID so a future timeout can identify the holder.
        try:
            lock_handle.seek(0)
            lock_handle.truncate()
            lock_handle.write(f"{os.getpid()}\n")
            lock_handle.flush()
        except OSError:
            pass  # cosmetic
        try:
            on_disk = int(conn.execute("PRAGMA user_version").fetchone()[0])
            # Detect the externally-dropped case: user_version stamp
            # is still correct (it lives in the SQLite header, not
            # the table) but the table itself is gone. Treat the same
            # as a fresh DB — CREATE TABLE puts us back in shape and
            # the user_version stamp is already correct.
            table_exists = conn.execute(
                "SELECT 1 FROM sqlite_master "
                "WHERE type='table' AND name='pr_state'"
            ).fetchone() is not None
            if on_disk == SCHEMA_VERSION and table_exists:
                # Another process won the race (or migration is
                # already current); nothing to do.
                _ensure_indices(conn)
                return
            if on_disk == SCHEMA_VERSION and not table_exists:
                _logger.warning(
                    "pr_state cache: user_version matches but table "
                    "is missing — recreating (external drop / "
                    "corruption recovery)",
                )
                _create_table_at_latest(conn)
                _ensure_indices(conn)
                return
            if on_disk > SCHEMA_VERSION:
                _logger.error(
                    "pr_state cache user_version=%s is NEWER than this "
                    "process knows about (expected %s). Refusing to "
                    "mutate — a downgrade would corrupt data the newer "
                    "process expects. Upgrade this process or wipe the "
                    "cache dir (%s) manually if you really want to "
                    "downgrade.",
                    on_disk, SCHEMA_VERSION, lock_dir,
                )
                return
            # Versions < 2 used a fundamentally different schema
            # (no owner/repo columns) — ALTER cannot preserve the
            # rows meaningfully. From v2 onward, every step is
            # additive, so the upgrade preserves data.
            if on_disk < 2:
                _logger.info(
                    "pr_state cache: legacy DB found (user_version=%s); "
                    "destructive rebuild required — pre-v2 schemas "
                    "lack the owner/repo columns",
                    on_disk,
                )
                conn.execute("DROP TABLE IF EXISTS pr_state")
                _create_table_at_latest(conn)
            else:
                # 2 <= on_disk < SCHEMA_VERSION → additive upgrade.
                _logger.info(
                    "pr_state cache: additive upgrade from v%s to v%s",
                    on_disk, SCHEMA_VERSION,
                )
                _apply_additive_upgrades(conn, on_disk)
            conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
            _ensure_indices(conn)
        finally:
            # flock auto-releases on fd close (the `with open(...)`
            # block); explicit release is redundant but documents
            # intent for the reader.
            fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)


def _create_table_at_latest(conn: sqlite3.Connection) -> None:
    """Build the table in its current shape. Called from a fresh-DB
    path or after a pre-versioning DROP."""
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS pr_state (
            owner                  TEXT NOT NULL,
            repo                   TEXT NOT NULL,
            number                 INTEGER NOT NULL,
            body_json              TEXT NOT NULL,
            head_sha               TEXT NOT NULL,
            state                  TEXT NOT NULL,
            labels_json            TEXT NOT NULL,
            updated_at                       TEXT NOT NULL,
            first_seen_at                    TEXT NOT NULL,
            last_seen_at                     TEXT NOT NULL,
            vanished_at                      TEXT,
            comments_refreshed_updated_at    TEXT,
            PRIMARY KEY (owner, repo, number)
        )
        """
    )


def _apply_additive_upgrades(conn: sqlite3.Connection, from_version: int) -> None:
    """Apply ALTER TABLE steps from ``from_version`` up to
    ``SCHEMA_VERSION``. Each step is non-destructive (ADD COLUMN
    only), so a deploy with version skew can co-exist: a v2 reader
    sees the v3 table without the new column and ignores it; a v3
    reader sees its own column populated by the v3 writer."""
    # v2 → v3: add comments_refreshed_updated_at.
    if from_version < 3:
        try:
            conn.execute(
                "ALTER TABLE pr_state "
                "ADD COLUMN comments_refreshed_updated_at TEXT"
            )
        except sqlite3.OperationalError as exc:
            # ``duplicate column name`` is benign — concurrent
            # migration must have raced (shouldn't with the lock,
            # but be defensive).
            if "duplicate column" not in str(exc).lower():
                raise


def _ensure_indices(conn: sqlite3.Connection) -> None:
    """Create the indices we depend on. ``IF NOT EXISTS`` makes this
    idempotent so callers can run it on every migration outcome."""
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_pr_state_state ON pr_state(state)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_pr_state_updated_at "
        "ON pr_state(updated_at)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_pr_state_vanished_at "
        "ON pr_state(vanished_at)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_pr_state_comments_refreshed "
        "ON pr_state(comments_refreshed_updated_at)"
    )


_CORRUPT_DB_MARKERS = (
    "malformed", "database disk image is malformed",
    "file is not a database", "not a database",
    "database is corrupt",
)


def _quarantine_corrupt_db() -> None:
    """Rename the SQLite file out of the way so the next ``_connect()``
    starts fresh. Called when ``sqlite3.DatabaseError`` indicates
    structural corruption (header damage, truncated WAL).

    Best-effort: a failure to rename logs + propagates so the caller
    can decide. We rename rather than delete so the file is available
    for forensics; the warmer's next cycle re-seeds within seconds."""
    target = cache_path()
    if not target.exists():
        return
    ts = _dt.datetime.now(_dt.timezone.utc).strftime("%Y%m%dT%H%M%S")
    quarantine = target.with_suffix(f".sqlite3.corrupt.{ts}")
    try:
        target.rename(quarantine)
        _logger.error(
            "pr_state cache: quarantined corrupt DB to %s; next cycle "
            "will re-seed from a fresh file", quarantine,
        )
    except OSError as exc:
        _logger.error(
            "pr_state cache: FAILED to quarantine corrupt DB at %s: %s "
            "— manual intervention required",
            target, exc,
        )
        raise


def _reheal_and_retry(exc: sqlite3.Error) -> bool:
    """Decide whether a SQLite error warrants clearing the
    per-process ``_initialized`` flag and letting the caller retry.

    Two paths:
      1. ``OperationalError: no such table`` / ``no such column`` —
         the DB file was replaced or its schema mutated out from under
         a long-running warmer. Re-arm; retry against the re-migrated
         schema.
      2. ``DatabaseError`` with corruption marker — the file itself
         is broken. Quarantine it and re-arm; the retry will start
         from a fresh file.

    Returns True iff caller should retry. False if the exception
    isn't one we know how to recover from.
    """
    global _initialized
    msg = str(exc).lower()
    if isinstance(exc, sqlite3.OperationalError) and (
        "no such table" in msg or "no such column" in msg
    ):
        with _init_lock:
            _initialized = False
        _logger.warning(
            "pr_state cache: re-arming migration after operational "
            "error (%s) — likely external DB replacement / schema drop",
            exc,
        )
        return True
    if isinstance(exc, sqlite3.DatabaseError) and any(
        marker in msg for marker in _CORRUPT_DB_MARKERS
    ):
        try:
            _quarantine_corrupt_db()
        except OSError:
            return False
        with _init_lock:
            _initialized = False
        return True
    return False


# ─── Public API ─────────────────────────────────────────────────────


def _with_reheal(fn):
    """Wrap a public-API call so a recoverable SQLite error triggers
    a one-shot re-arm + retry. Targets two failure classes:

      - ``OperationalError`` "no such table"/"no such column" —
        externally-replaced DB; re-migrate against the file as-is.
      - ``DatabaseError`` corruption markers — quarantine the
        broken file, re-migrate from scratch.

    Single retry — if the second attempt also fails, the error
    propagates so the caller sees a real signal instead of an
    infinite loop. ``DatabaseError`` is a SUPERCLASS of
    ``OperationalError`` so the catch order matters: subclass first.
    """
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except sqlite3.OperationalError as exc:
            if _reheal_and_retry(exc):
                return fn(*args, **kwargs)
            raise
        except sqlite3.DatabaseError as exc:
            if _reheal_and_retry(exc):
                return fn(*args, **kwargs)
            raise
    return wrapper


@_with_reheal
def upsert_prs(
    prs: list[dict[str, Any]], *, owner: str, repo: str,
) -> dict[str, Any]:
    """Write each PR's state within ``(owner, repo)``. Returns counts
    AND the changed-number list (union of inserted + updated). See
    module docstring for the schema details.

    Caller passes the full PR object dicts returned by Forgejo's
    ``/pulls`` endpoint."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    if not prs:
        return {"inserted": 0, "updated": 0, "unchanged": 0, "changed_numbers": []}
    now = _now()
    counts = {"inserted": 0, "updated": 0, "unchanged": 0}
    changed: list[int] = []
    conn = _connect()
    try:
        with conn:
            for pr in prs:
                if not isinstance(pr, dict):
                    continue
                try:
                    number = int(pr.get("number"))
                except (TypeError, ValueError):
                    continue
                if number <= 0:
                    continue
                head_sha = ""
                head = pr.get("head")
                if isinstance(head, dict):
                    head_sha = str(head.get("sha") or "")
                state = str(pr.get("state") or "open")
                labels = []
                raw_labels = pr.get("labels") or []
                if isinstance(raw_labels, list):
                    labels = [
                        str(l.get("name") or "")
                        for l in raw_labels
                        if isinstance(l, dict) and l.get("name")
                    ]
                # Canonicalize on write so the pending-comments-refresh
                # query (string equality vs comments_refreshed_updated_at)
                # is robust to Forgejo's timezone-marker drift.
                updated_at = _normalize_updated_at(str(pr.get("updated_at") or ""))
                body_json = json.dumps(pr, default=str, sort_keys=True)
                labels_json = json.dumps(labels)
                cached = conn.execute(
                    "SELECT updated_at FROM pr_state "
                    "WHERE owner = ? AND repo = ? AND number = ?",
                    (owner, repo, number),
                ).fetchone()
                if cached is None:
                    conn.execute(
                        """
                        INSERT INTO pr_state (
                            owner, repo, number, body_json, head_sha, state,
                            labels_json, updated_at, first_seen_at,
                            last_seen_at, vanished_at
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL)
                        """,
                        (
                            owner, repo, number, body_json, head_sha, state,
                            labels_json, updated_at, now, now,
                        ),
                    )
                    counts["inserted"] += 1
                    changed.append(number)
                elif cached["updated_at"] != updated_at:
                    conn.execute(
                        """
                        UPDATE pr_state SET
                            body_json     = ?,
                            head_sha      = ?,
                            state         = ?,
                            labels_json   = ?,
                            updated_at    = ?,
                            last_seen_at  = ?,
                            vanished_at   = NULL
                        WHERE owner = ? AND repo = ? AND number = ?
                        """,
                        (
                            body_json, head_sha, state, labels_json,
                            updated_at, now, owner, repo, number,
                        ),
                    )
                    counts["updated"] += 1
                    changed.append(number)
                else:
                    # Refresh body_json + labels_json even on the
                    # ``updated_at`` unchanged branch — Forgejo
                    # mutations like label add/remove do NOT bump
                    # ``updated_at``, so without this the warmer's
                    # cached body would carry stale labels for as
                    # long as the PR sat idle, and downstream
                    # consumers (cycle-cap, claim sweeps, filter
                    # exclusions) would never see them. Cheap: same
                    # JSON we already computed above.
                    conn.execute(
                        """
                        UPDATE pr_state SET
                            body_json    = ?,
                            labels_json  = ?,
                            last_seen_at = ?,
                            vanished_at  = NULL
                        WHERE owner = ? AND repo = ? AND number = ?
                        """,
                        (body_json, labels_json, now, owner, repo, number),
                    )
                    counts["unchanged"] += 1
    finally:
        conn.close()
    return {**counts, "changed_numbers": changed}


@_with_reheal
def mark_vanished(
    seen_numbers: set[int], *, owner: str, repo: str,
) -> int:
    """Mark any open-state row within ``(owner, repo)`` whose
    ``number`` is NOT in ``seen_numbers`` as vanished. Returns the
    number of rows newly marked vanished.

    Scoped to ``(owner, repo)`` so a multi-repo warmer doesn't
    accidentally vanish PRs from a sibling repo on a single poll.

    Implementation note: ``seen_numbers`` larger than
    ``_MARK_VANISHED_CHUNK`` would overflow SQLite's per-statement
    placeholder cap. We invert the semantics for large sets — write
    a temp table of seen IDs, then ``NOT IN (SELECT ...)`` — same
    result, no placeholder limit. The large-set branch runs both the
    temp-table populate AND the UPDATE inside an explicit
    BEGIN/COMMIT pair so a signal mid-batch can't leave the
    intermediate temp data without its accompanying mutation.
    """
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    now = _now()
    conn = _connect()
    try:
        if not seen_numbers:
            with conn:
                cur = conn.execute(
                    "UPDATE pr_state SET vanished_at = ? "
                    "WHERE owner = ? AND repo = ? AND vanished_at IS NULL",
                    (now, owner, repo),
                )
                return int(cur.rowcount or 0)
        if len(seen_numbers) <= _MARK_VANISHED_CHUNK:
            with conn:
                placeholders = ",".join("?" * len(seen_numbers))
                cur = conn.execute(
                    f"""
                    UPDATE pr_state SET vanished_at = ?
                    WHERE owner = ? AND repo = ?
                      AND vanished_at IS NULL
                      AND number NOT IN ({placeholders})
                    """,
                    (now, owner, repo, *seen_numbers),
                )
                return int(cur.rowcount or 0)
        # Large set: stage in a TEMP table to dodge the placeholder
        # cap. Autocommit + ``with conn:`` is a no-op (no implicit
        # BEGIN), so we drive the transaction explicitly to make the
        # whole temp-populate + UPDATE atomic. Failure mid-batch
        # rolls back the temp inserts AND any partial UPDATE; success
        # commits both as one unit. TEMP tables are connection-scoped
        # so cleanup happens on close.
        conn.execute("BEGIN")
        try:
            conn.execute("DROP TABLE IF EXISTS _seen_pr_numbers")
            conn.execute(
                "CREATE TEMP TABLE _seen_pr_numbers (n INTEGER PRIMARY KEY)"
            )
            conn.executemany(
                "INSERT OR IGNORE INTO _seen_pr_numbers (n) VALUES (?)",
                ((int(n),) for n in seen_numbers),
            )
            cur = conn.execute(
                """
                UPDATE pr_state SET vanished_at = ?
                WHERE owner = ? AND repo = ?
                  AND vanished_at IS NULL
                  AND number NOT IN (SELECT n FROM _seen_pr_numbers)
                """,
                (now, owner, repo),
            )
            rowcount = int(cur.rowcount or 0)
            conn.execute("COMMIT")
            return rowcount
        except BaseException:
            try:
                conn.execute("ROLLBACK")
            except sqlite3.DatabaseError:
                pass
            raise
    finally:
        conn.close()


@_with_reheal
def list_open_prs(*, owner: str, repo: str) -> list[dict[str, Any]]:
    """Return the full PR objects for every cached row within
    ``(owner, repo)`` that is currently open, sorted by
    ``updated_at`` desc (matches Forgejo's ``?sort=newest``)."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    conn = _connect()
    try:
        rows = conn.execute(
            """
            SELECT body_json FROM pr_state
            WHERE owner = ? AND repo = ? AND vanished_at IS NULL
            ORDER BY updated_at DESC
            """,
            (owner, repo),
        ).fetchall()
    finally:
        conn.close()
    out: list[dict[str, Any]] = []
    for r in rows:
        try:
            out.append(json.loads(r["body_json"]))
        except (ValueError, TypeError) as exc:
            _logger.warning("malformed body_json in pr_state cache: %s", exc)
    return out


@_with_reheal
def get_pr(number: int, *, owner: str, repo: str) -> dict[str, Any] | None:
    """Return one PR's full object within ``(owner, repo)`` by number,
    or None if not cached (or if the row is marked vanished)."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    conn = _connect()
    try:
        row = conn.execute(
            "SELECT body_json FROM pr_state "
            "WHERE owner = ? AND repo = ? AND number = ? "
            "AND vanished_at IS NULL",
            (owner, repo, int(number)),
        ).fetchone()
    finally:
        conn.close()
    if row is None:
        return None
    try:
        return json.loads(row["body_json"])
    except (ValueError, TypeError):
        return None


@_with_reheal
def count_rows(
    *, owner: str | None = None, repo: str | None = None,
) -> dict[str, int]:
    """Operator-facing diagnostic: ``{"total", "open", "vanished"}``.

    When ``owner`` + ``repo`` are both given, counts are scoped to
    that repo. Both omitted: global counts across every cached repo."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    conn = _connect()
    try:
        where = ""
        args: tuple[Any, ...] = ()
        if owner is not None and repo is not None:
            where = " WHERE owner = ? AND repo = ?"
            args = (owner, repo)
        total = conn.execute(
            f"SELECT COUNT(*) AS c FROM pr_state{where}", args,
        ).fetchone()["c"]
        open_where = where + (" AND " if where else " WHERE ") + "vanished_at IS NULL"
        open_n = conn.execute(
            f"SELECT COUNT(*) AS c FROM pr_state{open_where}", args,
        ).fetchone()["c"]
    finally:
        conn.close()
    return {"total": int(total), "open": int(open_n), "vanished": int(total - open_n)}


@_with_reheal
def latest_write_at(*, owner: str, repo: str) -> str | None:
    """Most-recent ``last_seen_at`` across all rows within
    ``(owner, repo)`` — i.e. when the warmer last successfully
    polled. Consumers use this as a staleness signal: if the
    timestamp is older than ~2× warmer interval, the warmer is
    likely dead and the cache is stale.

    Returns ``None`` when no rows exist (cold cache)."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    conn = _connect()
    try:
        row = conn.execute(
            "SELECT MAX(last_seen_at) AS latest FROM pr_state "
            "WHERE owner = ? AND repo = ?",
            (owner, repo),
        ).fetchone()
    finally:
        conn.close()
    if row is None or row["latest"] is None:
        return None
    return str(row["latest"])


@_with_reheal
def mark_comments_refreshed(
    refreshed: list[tuple[int, str]],
    *,
    owner: str,
    repo: str,
) -> int:
    """Stamp each row's ``comments_refreshed_updated_at`` to the PR's
    ``updated_at`` value at the moment of refresh. Returns rowcount.

    ``refreshed`` is a list of ``(pr_number, pr_updated_at)`` tuples
    — the caller passes the ``updated_at`` value the comments cache
    was just refreshed against, so a subsequent Forgejo-side change
    to the PR bumps ``pr_state.updated_at`` past the stamped value
    and the pending-refresh query re-surfaces the PR.
    """
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    if not refreshed:
        return 0
    # Normalize the stamped values so a Forgejo timezone-format flap
    # (``Z`` vs ``+00:00``, microseconds vs none) on the same instant
    # doesn't make the next ``list_pending`` query treat the row as
    # "changed since refresh" forever.
    rows_to_stamp = [
        (_normalize_updated_at(str(updated_at)), owner, repo, int(number))
        for number, updated_at in refreshed
    ]
    conn = _connect()
    try:
        # With ``isolation_level=None`` (autocommit) the ``with conn:``
        # context is a no-op — it does NOT wrap the loop in a single
        # transaction. Drive BEGIN/COMMIT explicitly so the N updates
        # are atomic: a failure on row K rolls back rows 0..K-1 too,
        # and the warmer's retry won't see partial state.
        conn.execute("BEGIN")
        try:
            conn.executemany(
                """
                UPDATE pr_state SET comments_refreshed_updated_at = ?
                WHERE owner = ? AND repo = ? AND number = ?
                """,
                rows_to_stamp,
            )
            # executemany on a CONNECTION-level cursor doesn't expose
            # a stable per-statement rowcount across drivers; run a
            # confirmation SELECT for the diag count.
            placeholders = ",".join("?" * len(refreshed))
            row = conn.execute(
                f"""
                SELECT COUNT(*) AS c FROM pr_state
                WHERE owner = ? AND repo = ?
                  AND number IN ({placeholders})
                """,
                (owner, repo, *(int(n) for n, _ in refreshed)),
            ).fetchone()
            conn.execute("COMMIT")
            return int(row["c"] or 0)
        except BaseException:
            try:
                conn.execute("ROLLBACK")
            except sqlite3.DatabaseError:
                pass
            raise
    finally:
        conn.close()


@_with_reheal
def count_pending_comments_refresh(*, owner: str, repo: str) -> int:
    """Count rows whose comments-cache refresh is stale or unwritten.
    Lets the warmer report accurate ``comments_deferred`` in its diag
    without materialising every pending row into memory.

    Raises ``PRStateCacheError`` when the cache is disabled — matches
    every other public-API call (was silently returning 0 in an
    earlier revision, which hid the disabled state in operator diags)."""
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    conn = _connect()
    try:
        row = conn.execute(
            """
            SELECT COUNT(*) AS c FROM pr_state
            WHERE owner = ? AND repo = ?
              AND vanished_at IS NULL
              AND updated_at != COALESCE(comments_refreshed_updated_at, '')
            """,
            (owner, repo),
        ).fetchone()
    finally:
        conn.close()
    return int(row["c"] or 0)


@_with_reheal
def list_pending_comments_refresh(
    *, owner: str, repo: str, limit: int,
) -> list[tuple[int, str]]:
    """Return ``(pr_number, updated_at)`` tuples for rows within
    ``(owner, repo)`` whose comments-cache refresh has not seen the
    current ``updated_at`` value — i.e. either the row has never been
    refreshed (NULL stamp) or the PR has been touched on Forgejo
    since the last refresh.

    Ordered by ``updated_at`` desc (most-actionable first), capped
    at ``limit``. Drives the warmer's refresh queue from PERSISTENT
    state rather than the per-cycle ``upsert_prs`` delta: a burst
    that overflows the per-cycle cap surfaces deferred rows next
    cycle even if nothing else changes (the old in-memory deferral
    lost them on warmer restart AND on the next cycle, because
    ``upsert_prs`` had already updated the cached ``updated_at`` so
    the per-cycle delta no longer flagged them).
    """
    if is_disabled():
        raise PRStateCacheError("PR-state cache is disabled via PR_STATE_CACHE_DISABLE")
    if limit <= 0:
        return []
    conn = _connect()
    try:
        rows = conn.execute(
            """
            SELECT number, updated_at FROM pr_state
            WHERE owner = ? AND repo = ?
              AND vanished_at IS NULL
              AND updated_at != COALESCE(comments_refreshed_updated_at, '')
            ORDER BY updated_at DESC
            LIMIT ?
            """,
            (owner, repo, int(limit)),
        ).fetchall()
    finally:
        conn.close()
    return [(int(r["number"]), str(r["updated_at"])) for r in rows]


def janitor() -> int:
    """Delete vanished rows older than ``_vanished_grace_s()``.
    Returns the number of rows removed. Best-effort: never raises."""
    if is_disabled():
        return 0
    cutoff = (
        _dt.datetime.now(_dt.timezone.utc)
        - _dt.timedelta(seconds=_vanished_grace_s())
    ).isoformat()
    conn: sqlite3.Connection | None = None
    try:
        conn = _connect()
        with conn:
            cur = conn.execute(
                "DELETE FROM pr_state WHERE vanished_at IS NOT NULL "
                "AND vanished_at <= ?",
                (cutoff,),
            )
            return int(cur.rowcount or 0)
    except (OSError, sqlite3.DatabaseError) as exc:
        _logger.warning("pr-state janitor swallowed error: %s", exc)
        return 0
    finally:
        if conn is not None:
            try:
                conn.close()
            except sqlite3.DatabaseError:
                pass


__all__ = (
    "PRStateCacheError",
    "SCHEMA_VERSION",
    "cache_dir",
    "cache_path",
    "count_rows",
    "get_pr",
    "is_disabled",
    "janitor",
    "latest_write_at",
    "count_pending_comments_refresh",
    "list_open_prs",
    "list_pending_comments_refresh",
    "mark_comments_refreshed",
    "mark_vanished",
    "upsert_prs",
)