db.py

"""
SQLite interface for paper cache, extraction results, and match outcomes.

Schema:
  papers:        arXiv metadata + pipeline status + citation metadata
  concepts:      extracted concepts per paper (from LLM deconstruction)
  reductions:    match results per concept (from compositional matching engine)
  formalism_kb:  cached KB entries for fast lookup (mirrors YAML)
  pipeline_runs: audit log of pipeline executions

Thread-safe: uses WAL mode. Single-writer by design (pipeline is sequential).
"""

from __future__ import annotations

import json
import sqlite3
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------

SCHEMA_SQL = """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;

CREATE TABLE IF NOT EXISTS papers (
    arxiv_id        TEXT PRIMARY KEY,
    title           TEXT NOT NULL,
    abstract        TEXT NOT NULL,
    authors         TEXT NOT NULL,          -- JSON array of strings
    categories      TEXT NOT NULL,          -- JSON array of strings
    published       TEXT NOT NULL,          -- ISO 8601
    updated         TEXT NOT NULL,          -- ISO 8601
    pdf_url         TEXT,
    -- Pipeline status
    status          TEXT NOT NULL DEFAULT 'ingested',  -- ingested|triaged|extracted|matched|displayed|skipped|error
    triage_passed   INTEGER,               -- 1 = novelty claim detected, 0 = skipped, NULL = not triaged
    triage_reason   TEXT,                   -- why it passed or was skipped
    ingestion_ts    TEXT NOT NULL DEFAULT (datetime('now')),
    extraction_ts   TEXT,
    matching_ts     TEXT,
    error_message   TEXT,
    citation_count  INTEGER DEFAULT 0,
    citation_fetched_ts TEXT
);

CREATE TABLE IF NOT EXISTS concepts (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    paper_arxiv_id  TEXT NOT NULL REFERENCES papers(arxiv_id) ON DELETE CASCADE,
    name            TEXT NOT NULL,          -- the term the paper uses
    is_claimed_novel INTEGER NOT NULL DEFAULT 0,
    claimed_novelty_text TEXT,
    mathematical_operation TEXT NOT NULL,
    domain          TEXT,
    codomain        TEXT,
    objective       TEXT,
    constraints     TEXT,                   -- JSON array
    canonical_analog TEXT,
    deconstructive_move TEXT,
    confidence      TEXT NOT NULL DEFAULT 'low',  -- high|medium|low
    confidence_rationale TEXT,
    flags           TEXT,                   -- JSON array
    -- Extraction metadata
    extraction_json TEXT NOT NULL,          -- full concept JSON from LLM
    extracted_at    TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS reductions (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    concept_id      INTEGER NOT NULL REFERENCES concepts(id) ON DELETE CASCADE,
    paper_arxiv_id  TEXT NOT NULL REFERENCES papers(arxiv_id) ON DELETE CASCADE,
    concept_name    TEXT NOT NULL,
    result_type     TEXT NOT NULL,          -- identity|compositional|analogy|unknown|confused
    reduction       TEXT NOT NULL,          -- e.g., "Kernel CCA ∘ neuralize ∘ predict_in_codomain"
    canonical_analog TEXT,
    genuine_delta   TEXT,
    micro           TEXT,
    meso            TEXT,
    macro           TEXT,
    confidence      REAL NOT NULL DEFAULT 0.0,
    display         TEXT NOT NULL,          -- sous rature formatted string
    notes           TEXT,                   -- JSON array
    match_json      TEXT NOT NULL,          -- full MatchResult as JSON
    matched_at      TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS formalism_kb (
    id              TEXT PRIMARY KEY,       -- matches formalism.id in YAML
    name            TEXT NOT NULL,
    signature_json  TEXT NOT NULL,          -- JSON: {operation, domain, codomain, objective_family}
    meso_type       TEXT,
    macro_type      TEXT,
    canonical_reference TEXT,
    researchor_artifact_id TEXT,
    researchor_mental_model_id TEXT,
    status          TEXT NOT NULL DEFAULT 'seed',
    cached_at       TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS pipeline_runs (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    run_type        TEXT NOT NULL,          -- daily|single|retroactive|manual
    started_at      TEXT NOT NULL DEFAULT (datetime('now')),
    finished_at     TEXT,
    papers_ingested INTEGER DEFAULT 0,
    papers_triaged  INTEGER DEFAULT 0,
    papers_extracted INTEGER DEFAULT 0,
    papers_matched  INTEGER DEFAULT 0,
    papers_error    INTEGER DEFAULT 0,
    status          TEXT NOT NULL DEFAULT 'running',  -- running|completed|failed
    error_message   TEXT,
    config_json     TEXT                    -- snapshot of pipeline config at run time
);

CREATE INDEX IF NOT EXISTS idx_papers_status ON papers(status);
CREATE INDEX IF NOT EXISTS idx_papers_updated ON papers(updated);
CREATE INDEX IF NOT EXISTS idx_concepts_paper ON concepts(paper_arxiv_id);
CREATE INDEX IF NOT EXISTS idx_reductions_paper ON reductions(paper_arxiv_id);
CREATE INDEX IF NOT EXISTS idx_reductions_concept ON reductions(concept_id);
CREATE INDEX IF NOT EXISTS idx_reductions_type ON reductions(result_type);
CREATE INDEX IF NOT EXISTS idx_pipeline_runs_started ON pipeline_runs(started_at);

-- Migration: add citation columns (safe to run on existing DBs)
ALTER TABLE papers ADD COLUMN citation_count INTEGER DEFAULT 0;
ALTER TABLE papers ADD COLUMN citation_fetched_ts TEXT;
"""


# ---------------------------------------------------------------------------
# Database wrapper
# ---------------------------------------------------------------------------

@dataclass
class Database:
    """SQLite database interface for the Différance Engine pipeline."""

    path: Path
    _conn: sqlite3.Connection | None = field(default=None, repr=False, init=False)

    def __post_init__(self):
        self.path = Path(self.path)
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def connect(self):
        """Open connection and ensure schema exists."""
        if self._conn is not None:
            return
        self._conn = sqlite3.connect(str(self.path))
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(SCHEMA_SQL)
        self._conn.commit()

    def close(self):
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, *args):
        self.close()

    # ---- Papers ----

    def paper_exists(self, arxiv_id: str) -> bool:
        self.connect()
        row = self._conn.execute("SELECT 1 FROM papers WHERE arxiv_id = ?", (arxiv_id,)).fetchone()
        return row is not None

    def insert_paper(self, paper: dict) -> bool:
        """Insert a paper from arXiv API parsed data. Returns True if new."""
        self.connect()
        if self.paper_exists(paper["arxiv_id"]):
            return False
        self._conn.execute(
            """INSERT INTO papers (arxiv_id, title, abstract, authors, categories,
               published, updated, pdf_url, status)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'ingested')""",
            (
                paper["arxiv_id"],
                paper["title"],
                paper["abstract"],
                json.dumps(paper.get("authors", [])),
                json.dumps(paper.get("categories", [])),
                paper.get("published", ""),
                paper.get("updated", ""),
                paper.get("pdf_url", ""),
            ),
        )
        self._conn.commit()
        return True

    def update_triage(self, arxiv_id: str, passed: bool, reason: str = ""):
        self.connect()
        self._conn.execute(
            """UPDATE papers SET triage_passed = ?, triage_reason = ?,
               status = CASE WHEN ? THEN 'triaged' ELSE 'skipped' END
               WHERE arxiv_id = ?""",
            (1 if passed else 0, reason, passed, arxiv_id),
        )
        self._conn.commit()

    def update_status(self, arxiv_id: str, status: str, error: str = ""):
        self.connect()
        ts = datetime.now(timezone.utc).isoformat()
        field = {"extracted": "extraction_ts", "matched": "matching_ts",
                 "displayed": "matching_ts"}.get(status, "")
        if field:
            self._conn.execute(
                f"UPDATE papers SET status = ?, {field} = ?, error_message = ? WHERE arxiv_id = ?",
                (status, ts, error, arxiv_id),
            )
        else:
            self._conn.execute(
                "UPDATE papers SET status = ?, error_message = ? WHERE arxiv_id = ?",
                (status, error, arxiv_id),
            )
        self._conn.commit()

    def update_citation(self, arxiv_id: str, count: int):
        self.connect()
        try:
            self._conn.execute(
                "UPDATE papers SET citation_count = ?, citation_fetched_ts = ? WHERE arxiv_id = ?",
                (count, "now", arxiv_id),
            )
            self._conn.commit()
        except sqlite3.OperationalError:
            # Column might not exist yet
            try:
                self._conn.execute("ALTER TABLE papers ADD COLUMN citation_count INTEGER DEFAULT 0")
                self._conn.execute("ALTER TABLE papers ADD COLUMN citation_fetched_ts TEXT")
                self._conn.execute(
                    "UPDATE papers SET citation_count = ?, citation_fetched_ts = ? WHERE arxiv_id = ?",
                    (count, "now", arxiv_id),
                )
                self._conn.commit()
            except Exception:
                pass

    def get_papers_by_status(self, status: str, limit: int = 100) -> list[dict]:
        self.connect()
        rows = self._conn.execute(
            "SELECT * FROM papers WHERE status = ? ORDER BY updated DESC LIMIT ?",
            (status, limit),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_papers_needing_extraction(self, limit: int = 10) -> list[dict]:
        self.connect()
        rows = self._conn.execute(
            "SELECT * FROM papers WHERE status = 'triaged' AND triage_passed = 1 ORDER BY updated DESC LIMIT ?",
            (limit,),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_paper(self, arxiv_id: str) -> dict | None:
        self.connect()
        row = self._conn.execute("SELECT * FROM papers WHERE arxiv_id = ?", (arxiv_id,)).fetchone()
        return _row_to_dict(row) if row else None

    def find_paper(self, arxiv_id: str) -> dict | None:
        """Look up a paper by arXiv ID, trying version-suffix variations.

        arXiv IDs can be stored with version suffixes (e.g. 2301.07093v1)
        but users may query without them (2301.07093). This tries exact
        match first, then strips/adds version suffixes.
        """
        # 1. Exact match
        paper = self.get_paper(arxiv_id)
        if paper:
            return paper

        # 2. User provided no version — try v1, v2, v3
        import re
        if not re.search(r'v\d+$', arxiv_id):
            for v in range(1, 4):
                paper = self.get_paper(f"{arxiv_id}v{v}")
                if paper:
                    return paper
        else:
            # 3. User provided version — try stripping it
            base = re.sub(r'v\d+$', '', arxiv_id)
            paper = self.get_paper(base)
            if paper:
                return paper

        # 4. LIKE prefix match (finds any version of this paper)
        self.connect()
        row = self._conn.execute(
            "SELECT * FROM papers WHERE arxiv_id LIKE ? || '%' LIMIT 1",
            (arxiv_id,),
        ).fetchone()
        return _row_to_dict(row) if row else None

    def count_by_status(self) -> dict[str, int]:
        self.connect()
        rows = self._conn.execute(
            "SELECT status, COUNT(*) as cnt FROM papers GROUP BY status"
        ).fetchall()
        return {r["status"]: r["cnt"] for r in rows}

    # ---- Concepts ----

    def insert_concept(self, paper_arxiv_id: str, concept: dict):
        self.connect()
        self._conn.execute(
            """INSERT INTO concepts (paper_arxiv_id, name, is_claimed_novel,
               claimed_novelty_text, mathematical_operation, domain, codomain,
               objective, constraints, canonical_analog, deconstructive_move,
               confidence, confidence_rationale, flags, extraction_json)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                paper_arxiv_id,
                concept.get("name", ""),
                1 if concept.get("is_claimed_novel") else 0,
                concept.get("claimed_novelty_text"),
                concept.get("mathematical_operation", ""),
                concept.get("domain"),
                concept.get("codomain"),
                concept.get("objective"),
                json.dumps(concept.get("constraints", [])),
                concept.get("canonical_analog"),
                concept.get("deconstructive_move"),
                concept.get("confidence", "low"),
                concept.get("confidence_rationale"),
                json.dumps(concept.get("flags", [])),
                json.dumps(concept),
            ),
        )
        self._conn.commit()
        return self._conn.execute("SELECT last_insert_rowid()").fetchone()[0]

    def delete_concepts_for_paper(self, arxiv_id: str):
        self.connect()
        self._conn.execute("DELETE FROM concepts WHERE paper_arxiv_id = ?", (arxiv_id,))
        self._conn.commit()

    # ---- Reductions ----

    def insert_reduction(self, concept_id: int, paper_arxiv_id: str, match: dict):
        self.connect()
        self._conn.execute(
            """INSERT INTO reductions (concept_id, paper_arxiv_id, concept_name,
               result_type, reduction, canonical_analog, genuine_delta, micro,
               meso, macro, confidence, display, notes, match_json)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                concept_id, paper_arxiv_id,
                match.get("concept_name", ""),
                match.get("result_type", ""),
                match.get("reduction", ""),
                match.get("canonical_analog", ""),
                match.get("genuine_delta", ""),
                match.get("micro", ""),
                match.get("meso", ""),
                match.get("macro", ""),
                match.get("confidence", 0.0),
                match.get("display", ""),
                json.dumps(match.get("notes", [])),
                json.dumps(match),
            ),
        )
        self._conn.commit()

    def delete_reductions_for_paper(self, arxiv_id: str):
        self.connect()
        self._conn.execute("DELETE FROM reductions WHERE paper_arxiv_id = ?", (arxiv_id,))
        self._conn.commit()

    def get_reductions_for_paper(self, arxiv_id: str) -> list[dict]:
        self.connect()
        rows = self._conn.execute(
            "SELECT * FROM reductions WHERE paper_arxiv_id = ? ORDER BY id",
            (arxiv_id,),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_displayable_papers(self, limit: int = 50) -> list[dict]:
        """Papers ready for display: matched results exist."""
        self.connect()
        rows = self._conn.execute(
            """SELECT p.* FROM papers p
               WHERE p.status = 'matched'
               ORDER BY p.updated DESC LIMIT ?""",
            (limit,),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    # ---- Cross-reference indexing ----

    def get_canonical_analogs(self) -> list[dict]:
        """All distinct canonical analogs cited, with paper counts."""
        self.connect()
        rows = self._conn.execute(
            """SELECT canonical_analog, COUNT(*) as paper_count,
                      GROUP_CONCAT(DISTINCT paper_arxiv_id) as paper_ids
               FROM reductions
               WHERE canonical_analog IS NOT NULL AND canonical_analog != ''
               GROUP BY canonical_analog
               ORDER BY paper_count DESC"""
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_papers_by_analog(self, analog: str, limit: int = 20) -> list[dict]:
        """Papers that share a canonical analog (fuzzy match)."""
        self.connect()
        rows = self._conn.execute(
            """SELECT DISTINCT p.* FROM papers p
               JOIN reductions r ON r.paper_arxiv_id = p.arxiv_id
               WHERE r.canonical_analog LIKE ?
               ORDER BY p.updated DESC LIMIT ?""",
            (f"%{analog}%", limit),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_papers_by_move(self, move: str, limit: int = 20) -> list[dict]:
        """Papers whose concepts use a specific deconstructive move."""
        self.connect()
        rows = self._conn.execute(
            """SELECT DISTINCT p.*, c.deconstructive_move, c.name as concept_name
               FROM papers p
               JOIN concepts c ON c.paper_arxiv_id = p.arxiv_id
               WHERE c.deconstructive_move = ?
               ORDER BY p.updated DESC LIMIT ?""",
            (move, limit),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_move_counts(self) -> list[dict]:
        """Count of each deconstructive move across all concepts."""
        self.connect()
        rows = self._conn.execute(
            """SELECT deconstructive_move, COUNT(*) as cnt
               FROM concepts
               WHERE deconstructive_move IS NOT NULL
               GROUP BY deconstructive_move
               ORDER BY cnt DESC"""
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    def get_papers_with_unknown(self, limit: int = 100) -> list[dict]:
        """Papers with at least one UNKNOWN reduction (retroactive queue)."""
        self.connect()
        rows = self._conn.execute(
            """SELECT DISTINCT p.* FROM papers p
               JOIN reductions r ON r.paper_arxiv_id = p.arxiv_id
               WHERE r.result_type = 'unknown'
               ORDER BY p.updated DESC LIMIT ?""",
            (limit,),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    # ---- KB cache ----

    def sync_kb_cache(self, formalisms: list[dict]):
        """Sync the cached KB table from the current formalisms list."""
        self.connect()
        self._conn.execute("DELETE FROM formalism_kb")
        for fm in formalisms:
            self._conn.execute(
                """INSERT OR REPLACE INTO formalism_kb
                   (id, name, signature_json, meso_type, macro_type,
                    canonical_reference, researchor_artifact_id,
                    researchor_mental_model_id, status)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    fm["id"],
                    fm["name"],
                    json.dumps(fm.get("signature", {})),
                    fm.get("meso_type"),
                    fm.get("macro_type"),
                    fm.get("canonical_reference"),
                    fm.get("researchor_artifact_id"),
                    fm.get("researchor_mental_model_id"),
                    fm.get("status", "seed"),
                ),
            )
        self._conn.commit()

    # ---- Pipeline runs ----

    def start_run(self, run_type: str, config: dict | None = None) -> int:
        self.connect()
        cur = self._conn.execute(
            "INSERT INTO pipeline_runs (run_type, config_json) VALUES (?, ?)",
            (run_type, json.dumps(config) if config else "{}"),
        )
        self._conn.commit()
        return cur.lastrowid

    def finish_run(self, run_id: int, counts: dict, error: str = ""):
        self.connect()
        status = "failed" if error else "completed"
        self._conn.execute(
            """UPDATE pipeline_runs SET finished_at = ?, papers_ingested = ?,
               papers_triaged = ?, papers_extracted = ?, papers_matched = ?,
               papers_error = ?, status = ?, error_message = ?
               WHERE id = ?""",
            (
                datetime.now(timezone.utc).isoformat(),
                counts.get("ingested", 0),
                counts.get("triaged", 0),
                counts.get("extracted", 0),
                counts.get("matched", 0),
                counts.get("errors", 0),
                status,
                error,
                run_id,
            ),
        )
        self._conn.commit()

    def recent_runs(self, limit: int = 10) -> list[dict]:
        self.connect()
        rows = self._conn.execute(
            "SELECT * FROM pipeline_runs ORDER BY started_at DESC LIMIT ?",
            (limit,),
        ).fetchall()
        return [_row_to_dict(r) for r in rows]

    # ---- Stats ----

    def stats(self) -> dict:
        self.connect()
        counts = self.count_by_status()
        total = sum(counts.values())

        # Reduction type breakdown
        red_rows = self._conn.execute(
            "SELECT result_type, COUNT(*) as cnt FROM reductions GROUP BY result_type"
        ).fetchall()
        reduction_counts = {r["result_type"]: r["cnt"] for r in red_rows}

        # Top deconstructive moves
        move_rows = self._conn.execute(
            "SELECT deconstructive_move, COUNT(*) as cnt FROM concepts WHERE deconstructive_move IS NOT NULL GROUP BY deconstructive_move ORDER BY cnt DESC LIMIT 5"
        ).fetchall()

        return {
            "total_papers": total,
            "by_status": counts,
            "total_concepts": sum(
                r["cnt"] for r in self._conn.execute("SELECT COUNT(*) as cnt FROM concepts").fetchall()
            ),
            "total_reductions": sum(reduction_counts.values()),
            "reductions_by_type": reduction_counts,
            "reduction_rate": (
                (reduction_counts.get("identity", 0) + reduction_counts.get("compositional", 0))
                / max(sum(reduction_counts.values()), 1)
            ),
            "top_moves": [{"move": r["deconstructive_move"], "count": r["cnt"]} for r in move_rows],
            "recent_runs": self.recent_runs(3),
        }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _row_to_dict(row: sqlite3.Row | None) -> dict | None:
    if row is None:
        return None
    d = dict(row)
    # Deserialize JSON fields
    for field in ("authors", "categories", "constraints", "flags", "notes"):
        if field in d and isinstance(d[field], str):
            try:
                d[field] = json.loads(d[field])
            except (json.JSONDecodeError, TypeError):
                pass
    return d


# ---------------------------------------------------------------------------
# Convenience
# ---------------------------------------------------------------------------

def get_db(path: Path | str | None = None) -> Database:
    """Get a Database instance for the default data directory."""
    if path is None:
        path = Path(__file__).resolve().parent.parent / "data" / "papers.db"
    return Database(Path(path))