"""Core DB primitives for the ``proficiency`` table.

Kept separate from the public writer surface so the empirical-Bayes
recomputation module can share the same single write path without creating
a circular import.
"""

from __future__ import annotations

import sqlite3
from collections.abc import Generator
from contextlib import contextmanager
from datetime import datetime, timezone

from config import RouterConfig
from proficiency import blend


@contextmanager
def _proficiency_tx(
    conn: sqlite3.Connection,
) -> Generator[None, None, None]:
    """Run a block inside a single SQLite transaction.

    If the caller has already started a transaction, this manager is a no-op
    except for rolling back on exception.  Otherwise it ensures columns are
    present, issues one ``conn.commit()`` at successful exit, and rolls back on
    error.  Python's sqlite3 module implicitly begins a transaction before the
    first data-modifying statement, so we do not issue an explicit ``BEGIN``.
    """
    own_transaction = not conn.in_transaction
    if own_transaction:
        ensure_columns(conn)
    try:
        yield
        if own_transaction:
            conn.commit()
    except Exception:
        if own_transaction:
            conn.rollback()
        raise


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


def ensure_columns(conn: sqlite3.Connection) -> None:
    """Add columns an older router.db predates. Idempotent, and cheap.

    schema.sql is CREATE TABLE IF NOT EXISTS, so it defines a NEW database and
    silently does nothing to an existing one. Anything added later therefore
    needs this, or the first write against a database created last week fails
    with "no such column".

    This function may run DDL; callers that need atomic behaviour should call
    it *before* beginning an explicit transaction.
    """
    columns = {row[1] for row in conn.execute("PRAGMA table_info(proficiency)")}
    migrated = False

    if "inherited_from" not in columns:
        conn.execute("ALTER TABLE proficiency ADD COLUMN inherited_from TEXT")
        conn.commit()
        _backfill_inherited(conn)
    if "outcome_score" not in columns:
        conn.execute("ALTER TABLE proficiency ADD COLUMN outcome_score REAL")
        migrated = True
    if "outcome_samples" not in columns:
        conn.execute(
            "ALTER TABLE proficiency ADD COLUMN outcome_samples INTEGER DEFAULT 0"
        )
        migrated = True

    if migrated and not conn.in_transaction:
        conn.commit()


def _backfill_inherited(conn: sqlite3.Connection) -> None:
    """One-time: mark rows this harness could not have measured directly.

    ADD COLUMN gives every existing row NULL, which reads as "measured here"
    and would leave exactly the rows this provenance was added for frozen
    forever -- the migration would ship the fix and none of the repair.

    Provenance that was never recorded cannot be recovered in general. It can
    be for the rows that matter, and not by guessing: ``eval_identities``
    selects ``latency_class='standard'`` rows, plus flex rows that have NO
    standard equivalent. So a flex row WITH a standard equivalent was never a
    candidate for direct evaluation, whatever its sample count says. That is
    the harness's own selection rule read backwards.

    Anything else keeps NULL, which is the safe direction: it means "do not
    overwrite", so a real measurement is never lost to this.
    """
    pairs = conn.execute(
        """
        SELECT v.model_id, v.provider, s.model_id
          FROM models v
          JOIN models s
            ON s.base_model_id = v.base_model_id
           AND s.provider = v.provider
           AND s.reasoning_mode = v.reasoning_mode
           AND s.context_variant = v.context_variant
         WHERE v.latency_class = 'flex' AND s.latency_class = 'standard'
        """
    ).fetchall()
    for variant_id, provider, source_id in pairs:
        conn.execute(
            """
            UPDATE proficiency SET inherited_from = ?
             WHERE model_id = ? AND provider = ? AND inherited_from IS NULL
            """,
            (source_id, variant_id, provider),
        )
    conn.commit()


def _read_row(
    conn: sqlite3.Connection, model_id: str, provider: str, category: str
) -> sqlite3.Row | None:
    conn.row_factory = sqlite3.Row
    return conn.execute(
        """
        SELECT leaderboard_score, self_eval_score, self_eval_samples,
               outcome_score, outcome_samples, inherited_from
          FROM proficiency
         WHERE model_id = ? AND provider = ? AND category = ?
        """,
        (model_id, provider, category),
    ).fetchone()


def _write(
    conn: sqlite3.Connection,
    cfg: RouterConfig,
    model_id: str,
    provider: str,
    category: str,
    leaderboard_score: float | None,
    self_eval_score: float | None,
    self_eval_samples: int,
    outcome_score: float | None = None,
    outcome_samples: int = 0,
    inherited_from: str | None = None,
    *,
    _blended_override: bool = False,
    _blended_score: float | None = None,
    _blended_source: str | None = None,
) -> None:
    ensure_columns(conn)

    if _blended_override:
        blended, source = _blended_score, _blended_source
    else:
        blended, source = blend(
            leaderboard_score,
            self_eval_score,
            self_eval_samples,
            leaderboard_weight=cfg.proficiency.leaderboard_weight,
            self_eval_weight=cfg.proficiency.self_eval_weight,
            min_samples=cfg.proficiency.self_eval_min_samples,
        )
    conn.execute(
        """
        INSERT INTO proficiency (
            model_id, provider, category, leaderboard_score,
            self_eval_score, self_eval_samples, blended_score, source,
            inherited_from, outcome_score, outcome_samples, last_updated
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(model_id, provider, category) DO UPDATE SET
            leaderboard_score = excluded.leaderboard_score,
            self_eval_score = excluded.self_eval_score,
            self_eval_samples = excluded.self_eval_samples,
            blended_score = excluded.blended_score,
            source = excluded.source,
            inherited_from = excluded.inherited_from,
            outcome_score = excluded.outcome_score,
            outcome_samples = excluded.outcome_samples,
            last_updated = excluded.last_updated
        """,
        (
            model_id,
            provider,
            category,
            leaderboard_score,
            self_eval_score,
            self_eval_samples,
            blended,
            source,
            inherited_from,
            outcome_score,
            outcome_samples,
            _now(),
        ),
    )