app/rag/evidence_store.py

"""Canonical SQLite store for structured academic evidence.

Chroma is a retrieval index, not the source of truth. This store owns exact
document structure, page/box provenance, relations, and the low-memory FTS5
lexical index used by hybrid retrieval.
"""

from __future__ import annotations

import json
import os
import sqlite3
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator, Sequence

from app.rag.models import (
    AcademicDocument,
    BoundingBox,
    EvidenceRelation,
    EvidenceUnit,
    normalize_evidence_text,
)


EVIDENCE_SCHEMA_VERSION = 1
DEFAULT_EVIDENCE_ROOT = Path(os.path.expanduser("~/.studybuddy/evidence"))


@dataclass(frozen=True)
class LexicalCandidate:
    evidence_id: str
    rank: float
    score: float


class EvidenceStore:
    def __init__(self, root: Path | str | None = None) -> None:
        configured = Path(root) if root is not None else DEFAULT_EVIDENCE_ROOT
        self.path = configured if configured.suffix in {".db", ".sqlite", ".sqlite3"} else configured / "evidence.sqlite3"
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._initialize()

    @contextmanager
    def _connect(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self.path, timeout=30.0)
        try:
            conn.row_factory = sqlite3.Row
            conn.execute("PRAGMA foreign_keys = ON")
            conn.execute("PRAGMA journal_mode = WAL")
            conn.execute("PRAGMA synchronous = NORMAL")
            yield conn
        finally:
            conn.close()

    def _initialize(self) -> None:
        with self._connect() as conn:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS evidence_meta (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL
                );

                CREATE TABLE IF NOT EXISTS documents (
                    project_id TEXT NOT NULL,
                    document_id TEXT NOT NULL,
                    filename TEXT NOT NULL,
                    title TEXT NOT NULL,
                    authors_json TEXT NOT NULL,
                    abstract TEXT NOT NULL,
                    page_count INTEGER NOT NULL,
                    parse_quality TEXT NOT NULL,
                    quality_flags_json TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    PRIMARY KEY (project_id, document_id)
                );

                CREATE TABLE IF NOT EXISTS evidence_units (
                    evidence_id TEXT PRIMARY KEY,
                    project_id TEXT NOT NULL,
                    document_id TEXT NOT NULL,
                    element_type TEXT NOT NULL,
                    page_start INTEGER NOT NULL,
                    page_end INTEGER NOT NULL,
                    parent_id TEXT,
                    section_path_json TEXT NOT NULL,
                    ordinal INTEGER NOT NULL,
                    bbox_json TEXT,
                    raw_text TEXT NOT NULL,
                    retrieval_text TEXT NOT NULL,
                    caption TEXT NOT NULL,
                    table_markdown TEXT NOT NULL,
                    visual_description TEXT NOT NULL,
                    quality_flags_json TEXT NOT NULL,
                    annotation_ids_json TEXT NOT NULL,
                    source_kind TEXT NOT NULL,
                    metadata_json TEXT NOT NULL,
                    FOREIGN KEY (project_id, document_id)
                        REFERENCES documents(project_id, document_id)
                        ON DELETE CASCADE
                );

                CREATE INDEX IF NOT EXISTS idx_evidence_project_document
                    ON evidence_units(project_id, document_id);
                CREATE INDEX IF NOT EXISTS idx_evidence_page
                    ON evidence_units(project_id, document_id, page_start, page_end);
                CREATE INDEX IF NOT EXISTS idx_evidence_type
                    ON evidence_units(project_id, element_type);

                CREATE TABLE IF NOT EXISTS evidence_relations (
                    source_evidence_id TEXT NOT NULL,
                    target_evidence_id TEXT NOT NULL,
                    relation_type TEXT NOT NULL,
                    confidence REAL NOT NULL,
                    derivation TEXT NOT NULL,
                    PRIMARY KEY (source_evidence_id, target_evidence_id, relation_type),
                    FOREIGN KEY (source_evidence_id) REFERENCES evidence_units(evidence_id) ON DELETE CASCADE,
                    FOREIGN KEY (target_evidence_id) REFERENCES evidence_units(evidence_id) ON DELETE CASCADE
                );

                CREATE VIRTUAL TABLE IF NOT EXISTS evidence_fts USING fts5(
                    evidence_id UNINDEXED,
                    project_id UNINDEXED,
                    document_id UNINDEXED,
                    content,
                    tokenize='unicode61 remove_diacritics 2'
                );
                """
            )
            current = conn.execute("SELECT value FROM evidence_meta WHERE key = 'schema_version'").fetchone()
            if current is None:
                conn.execute(
                    "INSERT INTO evidence_meta(key, value) VALUES ('schema_version', ?)",
                    (str(EVIDENCE_SCHEMA_VERSION),),
                )
            elif int(current["value"]) != EVIDENCE_SCHEMA_VERSION:
                raise RuntimeError(
                    f"unsupported evidence schema {current['value']}; clean rebuild required for {EVIDENCE_SCHEMA_VERSION}"
                )

    def replace_document(
        self,
        document: AcademicDocument,
        units: Sequence[EvidenceUnit],
        relations: Sequence[EvidenceRelation],
    ) -> None:
        self._validate_document_batch(document, units, relations)
        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            self._delete_document_rows(conn, document.project_id, document.document_id)
            conn.execute(
                """
                INSERT INTO documents(
                    project_id, document_id, filename, title, authors_json, abstract,
                    page_count, parse_quality, quality_flags_json, created_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    document.project_id,
                    document.document_id,
                    document.filename,
                    document.title,
                    _json(document.authors),
                    document.abstract,
                    document.page_count,
                    document.parse_quality,
                    _json(document.quality_flags),
                    document.created_at.isoformat(),
                ),
            )
            conn.executemany(
                """
                INSERT INTO evidence_units(
                    evidence_id, project_id, document_id, element_type, page_start,
                    page_end, parent_id, section_path_json, ordinal, bbox_json,
                    raw_text, retrieval_text, caption, table_markdown,
                    visual_description, quality_flags_json, annotation_ids_json,
                    source_kind, metadata_json
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [self._unit_row(unit) for unit in units],
            )
            fts_rows = [
                (unit.evidence_id, unit.project_id, unit.document_id, _fts_content(document, unit))
                for unit in units
                if unit.index_text.strip()
            ]
            if fts_rows:
                conn.executemany(
                    "INSERT INTO evidence_fts(evidence_id, project_id, document_id, content) VALUES (?, ?, ?, ?)",
                    fts_rows,
                )
            if relations:
                conn.executemany(
                    """
                    INSERT INTO evidence_relations(
                        source_evidence_id, target_evidence_id, relation_type, confidence, derivation
                    ) VALUES (?, ?, ?, ?, ?)
                    """,
                    [
                        (
                            relation.source_evidence_id,
                            relation.target_evidence_id,
                            str(relation.relation_type),
                            relation.confidence,
                            relation.derivation,
                        )
                        for relation in relations
                    ],
                )
            conn.commit()

    def get_document(self, project_id: str, document_id: str) -> AcademicDocument | None:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT * FROM documents WHERE project_id = ? AND document_id = ?",
                (project_id, document_id),
            ).fetchone()
        return _row_to_document(row) if row else None

    def list_documents(self, project_id: str) -> list[AcademicDocument]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT * FROM documents WHERE project_id = ? ORDER BY filename COLLATE NOCASE",
                (project_id,),
            ).fetchall()
        return [_row_to_document(row) for row in rows]

    def get_units(
        self,
        project_id: str,
        *,
        evidence_ids: Sequence[str] | None = None,
        document_ids: Sequence[str] | None = None,
    ) -> list[EvidenceUnit]:
        clauses = ["project_id = ?"]
        params: list[object] = [project_id]
        if evidence_ids:
            clauses.append(f"evidence_id IN ({','.join('?' for _ in evidence_ids)})")
            params.extend(evidence_ids)
        if document_ids:
            clauses.append(f"document_id IN ({','.join('?' for _ in document_ids)})")
            params.extend(document_ids)
        sql = f"SELECT * FROM evidence_units WHERE {' AND '.join(clauses)} ORDER BY document_id, ordinal"
        with self._connect() as conn:
            rows = conn.execute(sql, params).fetchall()
        units = [_row_to_unit(row) for row in rows]
        if evidence_ids:
            order = {evidence_id: index for index, evidence_id in enumerate(evidence_ids)}
            units.sort(key=lambda unit: order.get(unit.evidence_id, len(order)))
        return units

    def get_relations(
        self,
        evidence_ids: Sequence[str],
        *,
        outgoing: bool = True,
        incoming: bool = True,
    ) -> list[EvidenceRelation]:
        if not evidence_ids or (not outgoing and not incoming):
            return []
        placeholders = ",".join("?" for _ in evidence_ids)
        clauses: list[str] = []
        params: list[object] = []
        if outgoing:
            clauses.append(f"source_evidence_id IN ({placeholders})")
            params.extend(evidence_ids)
        if incoming:
            clauses.append(f"target_evidence_id IN ({placeholders})")
            params.extend(evidence_ids)
        with self._connect() as conn:
            rows = conn.execute(
                f"SELECT * FROM evidence_relations WHERE {' OR '.join(clauses)}",
                params,
            ).fetchall()
        return [EvidenceRelation(**dict(row)) for row in rows]

    def lexical_search(
        self,
        project_id: str,
        query: str,
        limit: int = 30,
        document_ids: Sequence[str] | None = None,
    ) -> list[LexicalCandidate]:
        terms = _fts_query(query)
        if not project_id or not terms:
            return []
        clauses = ["project_id = ?", "evidence_fts MATCH ?"]
        params: list[object] = [project_id, terms]
        if document_ids:
            clauses.append(f"document_id IN ({','.join('?' for _ in document_ids)})")
            params.extend(document_ids)
        params.append(max(1, int(limit)))
        with self._connect() as conn:
            rows = conn.execute(
                f"""
                SELECT evidence_id, bm25(evidence_fts) AS rank
                FROM evidence_fts
                WHERE {' AND '.join(clauses)}
                ORDER BY rank
                LIMIT ?
                """,
                params,
            ).fetchall()
        return [
            LexicalCandidate(
                evidence_id=row["evidence_id"],
                rank=float(row["rank"]),
                score=1.0 / (1.0 + abs(float(row["rank"]))),
            )
            for row in rows
        ]

    def resolve_selection(
        self,
        project_id: str,
        document_id: str,
        page_number: int,
        boxes: Sequence[BoundingBox] = (),
        text: str = "",
        limit: int = 8,
    ) -> list[str]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT * FROM evidence_units
                WHERE project_id = ? AND document_id = ?
                  AND page_start <= ? AND page_end >= ?
                ORDER BY ordinal
                """,
                (project_id, document_id, page_number, page_number),
            ).fetchall()
        normalized_text = normalize_evidence_text(text).casefold()
        scored: list[tuple[float, int, str]] = []
        for row in rows:
            unit = _row_to_unit(row)
            box_score = 0.0
            if unit.bbox_norm and boxes:
                box_score = max(unit.bbox_norm.intersection_ratio(box) for box in boxes)
            unit_text = normalize_evidence_text(unit.raw_text or unit.caption or unit.table_markdown).casefold()
            text_score = _text_overlap(normalized_text, unit_text)
            score = box_score * 2.0 + text_score
            if score > 0.0:
                scored.append((score, -unit.ordinal, unit.evidence_id))
        scored.sort(reverse=True)
        return [evidence_id for _, _, evidence_id in scored[: max(1, limit)]]

    def resolve_region(self, project_id: str, document_id: str, region_id: str) -> list[str]:
        if not project_id or not document_id or not region_id:
            return []
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT evidence_id FROM evidence_units
                WHERE project_id = ? AND document_id = ?
                  AND json_extract(metadata_json, '$.region_id') = ?
                ORDER BY ordinal
                """,
                (project_id, document_id, region_id),
            ).fetchall()
        return [str(row["evidence_id"]) for row in rows]

    def delete_document(self, project_id: str, document_id: str) -> None:
        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            self._delete_document_rows(conn, project_id, document_id)
            conn.commit()

    def clear_derived_data(self, project_id: str | None = None) -> None:
        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            if project_id:
                document_ids = [
                    row["document_id"]
                    for row in conn.execute(
                        "SELECT document_id FROM documents WHERE project_id = ?", (project_id,)
                    ).fetchall()
                ]
                for document_id in document_ids:
                    self._delete_document_rows(conn, project_id, document_id)
            else:
                conn.execute("DELETE FROM evidence_fts")
                conn.execute("DELETE FROM documents")
            conn.commit()

    def counts(self, project_id: str) -> dict[str, int]:
        with self._connect() as conn:
            documents = conn.execute(
                "SELECT COUNT(*) AS n FROM documents WHERE project_id = ?", (project_id,)
            ).fetchone()["n"]
            units = conn.execute(
                "SELECT COUNT(*) AS n FROM evidence_units WHERE project_id = ?", (project_id,)
            ).fetchone()["n"]
            fts = conn.execute(
                "SELECT COUNT(*) AS n FROM evidence_fts WHERE project_id = ?", (project_id,)
            ).fetchone()["n"]
        return {"documents": int(documents), "evidence_units": int(units), "fts_rows": int(fts)}

    @staticmethod
    def _validate_document_batch(
        document: AcademicDocument,
        units: Sequence[EvidenceUnit],
        relations: Sequence[EvidenceRelation],
    ) -> None:
        ids: set[str] = set()
        for unit in units:
            if unit.project_id != document.project_id or unit.document_id != document.document_id:
                raise ValueError("evidence unit does not belong to document")
            if unit.evidence_id in ids:
                raise ValueError(f"duplicate evidence id: {unit.evidence_id}")
            ids.add(unit.evidence_id)
        for relation in relations:
            if relation.source_evidence_id not in ids or relation.target_evidence_id not in ids:
                raise ValueError("relation endpoint is not part of the document batch")

    @staticmethod
    def _unit_row(unit: EvidenceUnit) -> tuple[object, ...]:
        return (
            unit.evidence_id,
            unit.project_id,
            unit.document_id,
            str(unit.element_type),
            unit.page_start,
            unit.page_end,
            unit.parent_id,
            _json(unit.section_path),
            unit.ordinal,
            _json(unit.bbox_norm.rounded()) if unit.bbox_norm else None,
            unit.raw_text,
            unit.retrieval_text,
            unit.caption,
            unit.table_markdown,
            unit.visual_description,
            _json(unit.quality_flags),
            _json(unit.annotation_ids),
            str(unit.source_kind),
            _json(unit.metadata),
        )

    @staticmethod
    def _delete_document_rows(conn: sqlite3.Connection, project_id: str, document_id: str) -> None:
        conn.execute(
            "DELETE FROM evidence_fts WHERE project_id = ? AND document_id = ?",
            (project_id, document_id),
        )
        conn.execute(
            "DELETE FROM documents WHERE project_id = ? AND document_id = ?",
            (project_id, document_id),
        )


def _json(value: object) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def _loads(value: str | None, fallback: object) -> object:
    return json.loads(value) if value else fallback


def _row_to_document(row: sqlite3.Row) -> AcademicDocument:
    return AcademicDocument(
        project_id=row["project_id"],
        document_id=row["document_id"],
        filename=row["filename"],
        title=row["title"],
        authors=_loads(row["authors_json"], []),
        abstract=row["abstract"],
        page_count=row["page_count"],
        parse_quality=row["parse_quality"],
        quality_flags=_loads(row["quality_flags_json"], []),
        created_at=row["created_at"],
    )


def _row_to_unit(row: sqlite3.Row) -> EvidenceUnit:
    bbox_payload = _loads(row["bbox_json"], None)
    return EvidenceUnit(
        evidence_id=row["evidence_id"],
        project_id=row["project_id"],
        document_id=row["document_id"],
        element_type=row["element_type"],
        page_start=row["page_start"],
        page_end=row["page_end"],
        parent_id=row["parent_id"],
        section_path=_loads(row["section_path_json"], []),
        ordinal=row["ordinal"],
        bbox_norm=BoundingBox(**bbox_payload) if bbox_payload else None,
        raw_text=row["raw_text"],
        retrieval_text=row["retrieval_text"],
        caption=row["caption"],
        table_markdown=row["table_markdown"],
        visual_description=row["visual_description"],
        quality_flags=_loads(row["quality_flags_json"], []),
        annotation_ids=_loads(row["annotation_ids_json"], []),
        source_kind=row["source_kind"],
        metadata=_loads(row["metadata_json"], {}),
    )


def _fts_content(document: AcademicDocument, unit: EvidenceUnit) -> str:
    return "\n".join(
        part
        for part in (
            document.title,
            " > ".join(unit.section_path),
            unit.raw_text,
            unit.caption,
            unit.table_markdown,
            unit.visual_description,
        )
        if part.strip()
    )


def _fts_query(query: str) -> str:
    tokens = [token for token in normalize_evidence_text(query).split(" ") if len(token) > 1]
    escaped = [f'"{token.replace(chr(34), chr(34) * 2)}"' for token in tokens[:32]]
    return " OR ".join(escaped)


def _text_overlap(needle: str, haystack: str) -> float:
    if not needle or not haystack:
        return 0.0
    if needle in haystack:
        return 1.0
    needle_terms = set(needle.split())
    if not needle_terms:
        return 0.0
    return len(needle_terms.intersection(haystack.split())) / len(needle_terms)