DocForge2/src/docforge/index.py

"""Atomic derived SQLite index and deterministic read operations."""

from __future__ import annotations

import hashlib
import json
import os
import sqlite3
import tempfile
from collections import deque
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path

from .errors import DocForgeError
from .models import Edge, Node, ProjectService, ProjectSnapshot
from .project import project_root_fingerprint

INDEX_SCHEMA_VERSION = 1
APPLICATION_ID = 1_146_683_778


def _node_hash(nodes: tuple[Node, ...]) -> str:
    payload = json.dumps(
        [node.as_dict() for node in nodes], sort_keys=True, separators=(",", ":")
    ).encode()
    return hashlib.sha256(payload).hexdigest()


def _edge_hash(edges: tuple[Edge, ...]) -> str:
    payload = json.dumps(
        [edge.as_dict() for edge in edges], sort_keys=True, separators=(",", ":")
    ).encode()
    return hashlib.sha256(payload).hexdigest()


def _connect_read_only(path: Path) -> sqlite3.Connection:
    if not path.is_file():
        raise DocForgeError("missing_index", "Derived index does not exist; run build first")
    connection = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    connection.row_factory = sqlite3.Row
    return connection


@contextmanager
def _read_connection(path: Path) -> Generator[sqlite3.Connection, None, None]:
    connection: sqlite3.Connection | None = None
    try:
        connection = _connect_read_only(path)
        yield connection
    except DocForgeError:
        raise
    except (sqlite3.Error, json.JSONDecodeError, KeyError, TypeError) as error:
        raise DocForgeError("invalid_index", "Derived index is corrupt or unreadable") from error
    finally:
        if connection is not None:
            connection.close()


def _status(snapshot: ProjectSnapshot) -> dict[str, object]:
    return {
        "project_id": snapshot.descriptor.project_id,
        "project_root_fingerprint": project_root_fingerprint(snapshot.descriptor.root),
        "revision": snapshot.revision,
        "source_hash": snapshot.source_hash,
        "node_hash": _node_hash(snapshot.nodes),
        "node_count": len(snapshot.nodes),
        "edge_hash": _edge_hash(snapshot.edges),
        "edge_count": len(snapshot.edges),
        "index_schema_version": INDEX_SCHEMA_VERSION,
        "adapter": snapshot.descriptor.adapter,
        "status": "ok",
    }


class ProjectIndex:
    """A disposable index that always checks current canonical source before queries."""

    def __init__(self, project: ProjectService) -> None:
        self.project = project

    @property
    def path(self) -> Path:
        return self.project.descriptor.index_path

    def build(self) -> dict[str, object]:
        snapshot = self.project.load()
        status = _status(snapshot)
        cache_root = snapshot.descriptor.cache_root
        cache_root.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(
            prefix="index-", suffix=".sqlite3", dir=cache_root, delete=False
        ) as descriptor:
            temporary = Path(descriptor.name)
        try:
            connection = sqlite3.connect(temporary)
            try:
                connection.execute(f"PRAGMA application_id={APPLICATION_ID}")
                connection.execute(f"PRAGMA user_version={INDEX_SCHEMA_VERSION}")
                connection.executescript(
                    """
                    CREATE TABLE metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL);
                    CREATE TABLE nodes (
                        node_id TEXT PRIMARY KEY,
                        title TEXT NOT NULL,
                        family TEXT NOT NULL,
                        authority TEXT NOT NULL,
                        status TEXT NOT NULL,
                        tags_json TEXT NOT NULL,
                        summary TEXT NOT NULL,
                        content TEXT NOT NULL,
                        source_path TEXT NOT NULL,
                        source_anchor TEXT,
                        content_hash TEXT NOT NULL
                    );
                    CREATE TABLE edges (
                        source_id TEXT NOT NULL,
                        relation TEXT NOT NULL,
                        target_id TEXT NOT NULL,
                        PRIMARY KEY (source_id, relation, target_id)
                    );
                    CREATE INDEX edges_target ON edges(target_id, relation, source_id);
                    CREATE VIRTUAL TABLE node_fts USING fts5(
                        node_id UNINDEXED, title, summary, content, tags
                    );
                    """
                )
                connection.executemany(
                    "INSERT INTO metadata(key, value) VALUES (?, ?)",
                    sorted((key, str(value)) for key, value in status.items()),
                )
                connection.executemany(
                    """
                    INSERT INTO nodes VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    [
                        (
                            node.node_id,
                            node.title,
                            node.family,
                            node.authority,
                            node.status,
                            json.dumps(node.tags, separators=(",", ":")),
                            node.summary,
                            node.content,
                            node.source_path,
                            node.source_anchor,
                            node.content_hash,
                        )
                        for node in snapshot.nodes
                    ],
                )
                connection.executemany(
                    "INSERT INTO edges VALUES (?, ?, ?)",
                    [(edge.source_id, edge.relation, edge.target_id) for edge in snapshot.edges],
                )
                connection.executemany(
                    "INSERT INTO node_fts VALUES (?, ?, ?, ?, ?)",
                    [
                        (
                            node.node_id,
                            node.title,
                            node.summary,
                            node.content,
                            " ".join(node.tags),
                        )
                        for node in snapshot.nodes
                    ],
                )
                connection.commit()
                integrity = connection.execute("PRAGMA integrity_check").fetchone()
                if integrity is None or integrity[0] != "ok":
                    raise DocForgeError("index_failure", "New index failed SQLite integrity check")
            finally:
                connection.close()
            current = self.project.load()
            if current.source_hash != snapshot.source_hash or current.revision != snapshot.revision:
                raise DocForgeError("source_changed", "Canonical source changed during index build")
            os.replace(temporary, self.path)
        except sqlite3.Error as error:
            temporary.unlink(missing_ok=True)
            raise DocForgeError("index_failure", "Could not build the derived index") from error
        except Exception:
            temporary.unlink(missing_ok=True)
            raise
        return {**status, "database": str(self.path)}

    def check(self) -> dict[str, object]:
        snapshot = self.project.load()
        expected = _status(snapshot)
        with _read_connection(self.path) as connection:
            application_id = connection.execute("PRAGMA application_id").fetchone()[0]
            schema_version = connection.execute("PRAGMA user_version").fetchone()[0]
            if application_id != APPLICATION_ID or schema_version != INDEX_SCHEMA_VERSION:
                raise DocForgeError("invalid_index", "Derived index has an unsupported schema")
            metadata = dict(connection.execute("SELECT key, value FROM metadata"))
            for key in (
                "project_id",
                "project_root_fingerprint",
                "revision",
                "source_hash",
                "node_hash",
                "node_count",
                "edge_hash",
                "edge_count",
                "index_schema_version",
                "adapter",
            ):
                if metadata.get(key) != str(expected[key]):
                    raise DocForgeError(
                        "stale_index", "Derived index does not match canonical source", field=key
                    )
            integrity = connection.execute("PRAGMA integrity_check").fetchone()
            if integrity is None or integrity[0] != "ok":
                raise DocForgeError("invalid_index", "Derived index failed SQLite integrity check")
            indexed_nodes = tuple(
                _row_to_node(row)
                for row in connection.execute("SELECT * FROM nodes ORDER BY node_id")
            )
            indexed_edges = tuple(
                Edge(*row)
                for row in connection.execute(
                    "SELECT source_id, relation, target_id FROM edges "
                    "ORDER BY source_id, relation, target_id"
                )
            )
            fts_count = connection.execute("SELECT COUNT(*) FROM node_fts").fetchone()[0]
            if (
                indexed_nodes != snapshot.nodes
                or indexed_edges != snapshot.edges
                or fts_count != len(snapshot.nodes)
            ):
                raise DocForgeError("invalid_index", "Derived index rows do not match source")
        return {**expected, "database": str(self.path)}

    def get_node(self, node_id: str) -> dict[str, object]:
        checked = self.check()
        with _read_connection(self.path) as connection:
            row = connection.execute("SELECT * FROM nodes WHERE node_id = ?", (node_id,)).fetchone()
        if row is None:
            raise DocForgeError(
                "missing_node", "No node has the requested stable ID", node_id=node_id
            )
        return self._result(checked, node=_row_to_node(row).as_dict())

    def search(self, query: str, *, limit: int | None = None) -> dict[str, object]:
        checked = self.check()
        limits = self.project.descriptor.limits
        if not query.strip() or len(query) > limits.max_query_chars:
            raise DocForgeError("invalid_query", "Search query is empty or exceeds its limit")
        bounded = _bounded_limit(limit, limits.max_results, default=20)
        terms = re_tokenize(query)
        if not terms:
            raise DocForgeError("invalid_query", "Search query contains no searchable text")
        expression = " AND ".join(f'"{term.replace(chr(34), chr(34) * 2)}"' for term in terms)
        with _read_connection(self.path) as connection:
            rows = connection.execute(
                """
                SELECT nodes.*, bm25(node_fts) AS rank,
                       snippet(node_fts, 3, '[', ']', ' … ', 18) AS snippet
                FROM node_fts JOIN nodes USING(node_id)
                WHERE node_fts MATCH ?
                ORDER BY rank, nodes.node_id
                LIMIT ?
                """,
                (expression, bounded),
            ).fetchall()
        results: list[dict[str, object]] = []
        for row in rows:
            payload = _row_to_node(row).as_dict(include_content=False)
            payload.update({"rank": row["rank"], "snippet": row["snippet"]})
            results.append(payload)
        return self._result(checked, query=query, count=len(results), results=results)

    def filter_nodes(
        self,
        *,
        family: str | None = None,
        authority: str | None = None,
        status: str | None = None,
        tag: str | None = None,
        limit: int | None = None,
    ) -> dict[str, object]:
        checked = self.check()
        bounded = _bounded_limit(limit, self.project.descriptor.limits.max_results, default=100)
        clauses: list[str] = []
        values: list[object] = []
        for field, value in (("family", family), ("authority", authority), ("status", status)):
            if value is not None:
                clauses.append(f"{field} = ?")
                values.append(value)
        if tag is not None:
            clauses.append("EXISTS (SELECT 1 FROM json_each(tags_json) WHERE value = ?)")
            values.append(tag)
        where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
        with _read_connection(self.path) as connection:
            rows = connection.execute(
                f"SELECT * FROM nodes {where} ORDER BY node_id LIMIT ?", (*values, bounded)
            ).fetchall()
        results = [_row_to_node(row).as_dict(include_content=False) for row in rows]
        return self._result(checked, count=len(results), results=results)

    def backlinks(self, node_id: str, *, relation: str | None = None) -> dict[str, object]:
        return self._edges(node_id, incoming=True, relation=relation)

    def dependencies(self, node_id: str, *, depth: int = 2) -> dict[str, object]:
        return self._traverse(node_id, incoming=False, depth=depth, relation="depends_on")

    def impact(self, node_id: str, *, depth: int = 2) -> dict[str, object]:
        return self._traverse(node_id, incoming=True, depth=depth, relation=None)

    def _edges(self, node_id: str, *, incoming: bool, relation: str | None) -> dict[str, object]:
        checked = self.check()
        self._require_node(node_id)
        source_column = "target_id" if incoming else "source_id"
        relation_clause = " AND relation = ?" if relation is not None else ""
        values: tuple[object, ...] = (node_id, relation) if relation is not None else (node_id,)
        with _read_connection(self.path) as connection:
            rows = connection.execute(
                f"SELECT source_id, relation, target_id FROM edges "
                f"WHERE {source_column} = ?{relation_clause} "
                "ORDER BY source_id, relation, target_id",
                values,
            ).fetchall()
        return self._result(checked, edges=[Edge(*row).as_dict() for row in rows])

    def _traverse(
        self, node_id: str, *, incoming: bool, depth: int, relation: str | None
    ) -> dict[str, object]:
        checked = self.check()
        self._require_node(node_id)
        maximum = self.project.descriptor.limits.max_traversal_depth
        if type(depth) is not int or depth < 0 or depth > maximum:
            raise DocForgeError("invalid_depth", "Traversal depth is outside the configured limit")
        with _read_connection(self.path) as connection:
            edges = tuple(
                Edge(*row)
                for row in connection.execute(
                    "SELECT source_id, relation, target_id FROM edges "
                    "ORDER BY source_id, relation, target_id"
                )
            )
        queue: deque[tuple[str, int, tuple[str, ...]]] = deque([(node_id, 0, (node_id,))])
        seen = {node_id}
        results: list[dict[str, object]] = []
        while queue:
            current, current_depth, path = queue.popleft()
            if current_depth >= depth:
                continue
            candidates = [
                edge
                for edge in edges
                if (relation is None or edge.relation == relation)
                and ((edge.target_id if incoming else edge.source_id) == current)
            ]
            for edge in candidates:
                target = edge.source_id if incoming else edge.target_id
                if target in seen:
                    continue
                seen.add(target)
                target_path = (*path, target)
                results.append(
                    {
                        "node_id": target,
                        "depth": current_depth + 1,
                        "relation": edge.relation,
                        "path": target_path,
                    }
                )
                queue.append((target, current_depth + 1, target_path))
        return self._result(checked, root=node_id, depth=depth, count=len(results), results=results)

    def _require_node(self, node_id: str) -> None:
        with _read_connection(self.path) as connection:
            exists = connection.execute(
                "SELECT 1 FROM nodes WHERE node_id = ?", (node_id,)
            ).fetchone()
        if exists is None:
            raise DocForgeError(
                "missing_node", "No node has the requested stable ID", node_id=node_id
            )

    def _result(self, checked: dict[str, object], **payload: object) -> dict[str, object]:
        after = self.check()
        if (
            after["source_hash"] != checked["source_hash"]
            or after["revision"] != checked["revision"]
        ):
            raise DocForgeError("source_changed", "Canonical source changed during the query")
        return {
            "status": "ok",
            "project_id": checked["project_id"],
            "project_root_fingerprint": checked["project_root_fingerprint"],
            "revision": checked["revision"],
            "source_hash": checked["source_hash"],
            "adapter": checked["adapter"],
            **payload,
        }


def _row_to_node(row: sqlite3.Row) -> Node:
    return Node(
        node_id=row["node_id"],
        title=row["title"],
        family=row["family"],
        authority=row["authority"],
        status=row["status"],
        tags=tuple(json.loads(row["tags_json"])),
        summary=row["summary"],
        content=row["content"],
        source_path=row["source_path"],
        source_anchor=row["source_anchor"],
        content_hash=row["content_hash"],
    )


def _bounded_limit(value: int | None, maximum: int, *, default: int) -> int:
    if value is None:
        return min(default, maximum)
    if type(value) is not int or value < 1 or value > maximum:
        raise DocForgeError("invalid_limit", "Result limit is outside the configured range")
    return value


def re_tokenize(query: str) -> tuple[str, ...]:
    token = ""
    tokens: list[str] = []
    for character in query.casefold():
        if character.isalnum() or character in {"_", "-"}:
            token += character
        elif token:
            tokens.append(token)
            token = ""
    if token:
        tokens.append(token)
    return tuple(tokens)