stack/src/lead/review/store.py

"""Append-only gold labels. The latest row for a key is the current one."""

from __future__ import annotations

import sqlite3
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

from lead.review.codes import DECISIONS, PASSES, PROJECTS, STAGES, reason_codes

_DDL = """
CREATE TABLE IF NOT EXISTS label (
    id INTEGER PRIMARY KEY,
    zotero_key TEXT NOT NULL,
    project TEXT NOT NULL,
    stage TEXT NOT NULL,
    pass INTEGER NOT NULL,
    decision TEXT NOT NULL,
    reason TEXT NOT NULL,
    rationale TEXT NOT NULL,
    actor TEXT NOT NULL,
    created_at TEXT NOT NULL
);
"""


@dataclass(frozen=True)
class Label:
    zotero_key: str
    project: str
    stage: str
    pass_n: int
    decision: str
    reason: str
    rationale: str
    actor: str
    created_at: str


class Store:
    """One sqlite file of human labels. Opens per call and commits."""

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        con = self._open()
        try:
            con.execute(_DDL)
            con.commit()
        finally:
            con.close()

    def _open(self) -> sqlite3.Connection:
        con = sqlite3.connect(self.path)
        con.row_factory = sqlite3.Row
        return con

    def record(
        self,
        *,
        zotero_key: str,
        project: str,
        stage: str,
        pass_n: int,
        decision: str,
        reason: str = "",
        rationale: str,
        actor: str,
    ) -> None:
        """Append one label. A second label for the same key supersedes it."""
        key = zotero_key.strip()
        if not key:
            raise ValueError("zotero key is empty")
        if project not in PROJECTS:
            raise ValueError(project)
        if stage not in STAGES:
            raise ValueError(stage)
        if pass_n not in PASSES:
            raise ValueError(pass_n)
        if decision not in DECISIONS:
            raise ValueError(decision)
        why = reason.strip()
        if decision == "exclude":
            if why not in reason_codes(project):
                raise ValueError(f"{why or '(blank)'} is not in the {project} codebook")
        elif why:
            raise ValueError("a reason is recorded only on exclude")
        text = rationale.strip()
        if not text:
            raise ValueError("rationale is empty")
        who = actor.strip() or "reviewer"
        stamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        con = self._open()
        try:
            con.execute(
                """INSERT INTO label (
                       zotero_key, project, stage, pass, decision, reason,
                       rationale, actor, created_at
                   ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (key, project, stage, pass_n, decision, why, text, who, stamp),
            )
            con.commit()
        finally:
            con.close()

    def latest(self, project: str, stage: str, pass_n: int) -> dict[str, str]:
        """Current decision per zotero key. A later row replaces an earlier one."""
        con = self._open()
        try:
            rows = con.execute(
                """SELECT zotero_key, decision FROM label
                   WHERE project = ? AND stage = ? AND pass = ?
                   ORDER BY id""",
                (project, stage, pass_n),
            ).fetchall()
        finally:
            con.close()
        out: dict[str, str] = {}
        for row in rows:
            out[row["zotero_key"]] = row["decision"]
        return out

    def count(self, project: str, stage: str, pass_n: int) -> int:
        return len(self.latest(project, stage, pass_n))

    def history(self, project: str, stage: str, pass_n: int) -> tuple[Label, ...]:
        """Every row, oldest first. The page shows the tail."""
        con = self._open()
        try:
            rows = con.execute(
                """SELECT zotero_key, project, stage, pass, decision, reason,
                          rationale, actor, created_at
                   FROM label
                   WHERE project = ? AND stage = ? AND pass = ?
                   ORDER BY id""",
                (project, stage, pass_n),
            ).fetchall()
        finally:
            con.close()
        return tuple(
            Label(
                zotero_key=row["zotero_key"],
                project=row["project"],
                stage=row["stage"],
                pass_n=row["pass"],
                decision=row["decision"],
                reason=row["reason"],
                rationale=row["rationale"],
                actor=row["actor"],
                created_at=row["created_at"],
            )
            for row in rows
        )