src/store/db.py

"""Local SQLite working store for inference samples and reviews."""
from __future__ import annotations

import json
import os
import sqlite3
from contextlib import closing
from typing import Any

from . import hf_sync

_PERSISTENT_ROOT = "/data"
_PULLED = False


def _resolve_db_path() -> str:
    """Pick a writable store DB path without crashing on startup."""
    preferred = os.environ.get("STORE_DB", "/data/store.db")
    parent = os.path.dirname(os.path.abspath(preferred))
    try:
        os.makedirs(parent, exist_ok=True)
        if os.access(parent, os.W_OK):
            return preferred
    except OSError:
        pass
    home = os.environ.get("HOME") or os.path.expanduser("~") or "."
    return os.path.join(home, "store.db")


_DB_PATH = _resolve_db_path()


def db_status() -> tuple[str, bool]:
    """Return (resolved_db_path, is_persistent)."""
    path = os.path.abspath(_DB_PATH)
    root = os.path.abspath(_PERSISTENT_ROOT)
    return _DB_PATH, path.startswith(root + os.sep) or path == root


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(_DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS samples (
            id             INTEGER PRIMARY KEY AUTOINCREMENT,
            ts             TEXT NOT NULL,
            host           TEXT,
            inchikey       TEXT NOT NULL,
            guest_name     TEXT,
            model          TEXT NOT NULL,
            prompt_version TEXT NOT NULL,
            prompt_label   TEXT,
            batch          TEXT,
            prompt         TEXT,
            results_json   TEXT NOT NULL,
            true_logka     REAL,
            UNIQUE(inchikey, model, prompt_version)
        )
        """
    )
    # Migrate DBs created before `batch`/`prompt` became their own atomic
    # columns (previously packed into prompt_label as "<prompt>·<batch>").
    _cols = {r[1] for r in conn.execute("PRAGMA table_info(samples)")}
    if "batch" not in _cols:
        conn.execute("ALTER TABLE samples ADD COLUMN batch TEXT")
    if "prompt" not in _cols:
        conn.execute("ALTER TABLE samples ADD COLUMN prompt TEXT")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS reviews (
            id             INTEGER PRIMARY KEY AUTOINCREMENT,
            ts             TEXT NOT NULL,
            inchikey       TEXT,
            guest_name     TEXT,
            model          TEXT,
            prompt_version TEXT,
            rating         TEXT,
            comment        TEXT,
            reviewer       TEXT
        )
        """
    )
    return conn


def _init_db() -> None:
    global _PULLED
    try:
        with closing(_connect()) as conn:
            conn.commit()
        if not _PULLED:
            _PULLED = True
            hf_sync.pull_into_sqlite(_DB_PATH)
    except Exception:
        # The resolver should pick a usable path, but startup must remain robust.
        pass


def _decode_sample(row: sqlite3.Row) -> dict[str, Any]:
    out = dict(row)
    try:
        out["results"] = json.loads(out.get("results_json") or "{}")
    except json.JSONDecodeError:
        out["results"] = {}
    return out


def get_sample(inchikey: str, model: str, prompt_version: str) -> dict | None:
    """Return one cached sample with parsed results, or None."""
    with closing(_connect()) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            """
            SELECT * FROM samples
            WHERE inchikey = ? AND model = ? AND prompt_version = ?
            """,
            (inchikey, model, prompt_version),
        ).fetchone()
    return _decode_sample(row) if row is not None else None


def get_sample_any(inchikey: str, model: str) -> dict | None:
    """Return the stored sample for a guest+model regardless of prompt_version
    (most recent first). Used by the read-only Review 'Run' so it surfaces the
    canonical batch prediction (stored as 'v5') instead of missing on the UI's
    hashed prompt key. Tries the given model, then its normalized form."""
    with closing(_connect()) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            """
            SELECT * FROM samples
            WHERE inchikey = ? AND model IN (?, ?)
            ORDER BY ts DESC LIMIT 1
            """,
            (inchikey, model, norm_model(model)),
        ).fetchone()
    return _decode_sample(row) if row is not None else None


def norm_model(model: str | None) -> str:
    """Canonical model id: drop any provider prefix ("openai/gpt-5.5" -> "gpt-5.5")."""
    return (model or "").split("/")[-1].strip()


def norm_prompt(prompt_version: str | None) -> str:
    """Atomic prompt version: the leading token before any ':' suffix
    ("v5:gpt-5.5" -> "v5", "v1:openai/gpt-5.5:hash" -> "v1")."""
    return (prompt_version or "").split(":")[0].strip()


def fmt_ts(ts: str | None) -> str:
    """Render an ISO-8601 timestamp as 'YYYY-MM-DD:HHmm' (UTC, minute precision)."""
    s = (ts or "").strip()
    if not s:
        return ""
    try:
        from datetime import datetime

        return datetime.fromisoformat(s.replace("Z", "+00:00")).strftime("%Y-%m-%d:%H%M")
    except ValueError:
        return s[:16]


def put_sample(row: dict) -> None:
    """Upsert one inference sample and schedule durable HF sync."""
    results_json = row.get("results_json")
    if results_json is None and "results" in row:
        results_json = json.dumps(row.get("results") or {})
    with closing(_connect()) as conn, conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO samples (
                ts, host, inchikey, guest_name, model, prompt_version,
                prompt_label, batch, prompt, results_json, true_logka
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                row.get("ts", ""),
                row.get("host", ""),
                row.get("inchikey", ""),
                row.get("guest_name", ""),
                norm_model(row.get("model", "")),
                norm_prompt(row.get("prompt_version", "")),
                row.get("prompt_label", ""),
                row.get("batch", ""),
                row.get("prompt") or norm_prompt(row.get("prompt_version", "")),
                results_json or "{}",
                row.get("true_logka"),
            ),
        )
    hf_sync.schedule_push(_DB_PATH)


def save_review(
    *,
    inchikey: str,
    guest_name: str,
    model: str,
    prompt_version: str,
    rating: str,
    comment: str,
    reviewer: str,
    ts: str,
) -> int:
    """Upsert one review and schedule durable HF sync.

    A reviewer has at most one review per scored guest
    (inchikey, model, prompt_version, reviewer): re-submitting edits the existing
    row's rating/comment/timestamp instead of appending a duplicate, so the
    Feedback view stays one-row-per-reviewer and the Review tab can be used to
    edit a past review.
    """
    nm, npv = norm_model(model), norm_prompt(prompt_version)
    with closing(_connect()) as conn, conn:
        existing = conn.execute(
            "SELECT id FROM reviews WHERE COALESCE(inchikey,'')=? AND COALESCE(model,'')=? "
            "AND COALESCE(prompt_version,'')=? AND COALESCE(reviewer,'')=? ORDER BY id DESC LIMIT 1",
            (inchikey or "", nm, npv, reviewer or ""),
        ).fetchone()
        if existing is not None:
            row_id = int(existing[0])
            conn.execute(
                "UPDATE reviews SET ts=?, guest_name=?, rating=?, comment=? WHERE id=?",
                (ts, guest_name, rating, comment, row_id),
            )
        else:
            cur = conn.execute(
                """
                INSERT INTO reviews (
                    ts, inchikey, guest_name, model, prompt_version,
                    rating, comment, reviewer
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (ts, inchikey, guest_name, nm, npv, rating, comment, reviewer),
            )
            row_id = int(cur.lastrowid or 0)
    hf_sync.schedule_push(_DB_PATH)
    return row_id


def list_samples() -> list[dict]:
    """Return samples with latest matching review fields for admin views."""
    query = """
        SELECT
            s.inchikey,
            s.guest_name,
            s.host,
            s.model,
            s.prompt_version,
            s.prompt_label,
            s.batch,
            s.prompt,
            s.results_json,
            s.true_logka,
            r.rating,
            r.comment,
            s.ts
        FROM samples s
        LEFT JOIN reviews r
          ON r.id = (
              SELECT rr.id
              FROM reviews rr
              WHERE COALESCE(rr.inchikey, '') = COALESCE(s.inchikey, '')
                AND COALESCE(rr.model, '') = COALESCE(s.model, '')
                AND COALESCE(rr.prompt_version, '') = COALESCE(s.prompt_version, '')
              ORDER BY rr.id DESC
              LIMIT 1
          )
        ORDER BY s.id DESC
    """
    with closing(_connect()) as conn:
        conn.row_factory = sqlite3.Row
        rows = []
        for row in conn.execute(query):
            d = dict(row)
            try:
                results = json.loads(d.pop("results_json") or "{}")
            except json.JSONDecodeError:
                results = {}
            combined = results.get("combined") or {}
            d["combined_pred"] = combined.get("pred")
            d["combined_tldr"] = combined.get("tldr") or ""
            # Normalize on read so legacy rows join/display consistently.
            d["model"] = norm_model(d.get("model"))
            d["prompt_version"] = norm_prompt(d.get("prompt_version"))
            d["prompt"] = d.get("prompt") or d["prompt_version"]
            rows.append(d)
        return rows


def pull_remote() -> None:
    """Re-pull the HF store dataset into local SQLite (refresh durable predictions)."""
    hf_sync.pull_into_sqlite(_DB_PATH)


def export_reviews() -> list[dict]:
    """Return all reviews newest first."""
    with closing(_connect()) as conn:
        conn.row_factory = sqlite3.Row
        return [dict(r) for r in conn.execute("SELECT * FROM reviews ORDER BY id DESC")]


def review_map() -> dict[tuple[str, str, str], dict[str, str]]:
    """Aggregate every review into {(inchikey, model, prompt_version): {reviewer: rating}}.

    Multi-reviewer surface for the merged Data board: each sample maps to the set
    of reviewers who rated it and their latest rating. Reviews are append-only, so
    when a reviewer rates the same sample twice we keep the newest (highest id).
    An empty/blank reviewer is bucketed under "anon" so it still shows.
    """
    out: dict[tuple[str, str, str], dict[str, str]] = {}
    # Ascending id so later (newer) rows overwrite earlier ones per reviewer.
    with closing(_connect()) as conn:
        conn.row_factory = sqlite3.Row
        rows = conn.execute(
            "SELECT inchikey, model, prompt_version, rating, reviewer "
            "FROM reviews ORDER BY id ASC"
        )
        for r in rows:
            key = (r["inchikey"] or "", norm_model(r["model"]), norm_prompt(r["prompt_version"]))
            reviewer = (r["reviewer"] or "").strip() or "anon"
            out.setdefault(key, {})[reviewer] = r["rating"] or ""
    return out


_init_db()