database.py

"""Persistent storage layer: SQLite + sqlite-vec for epistemic memory records."""

from __future__ import annotations

import sqlite3
import subprocess
from datetime import datetime, timezone
from typing import Any

import sqlite_vec
from sqlite_vec import serialize_float32

# EmbeddingGemma default (MRL can truncate to 512/256/128).
# This constant is the fallback default; the actual dimension is
# stored in codegrave_meta and controlled by config [storage].embedding_dim.
EMBEDDING_DIM = 768

# A handful of oversized fields would otherwise blow the judge model's context
# window on every future query_memory call — cap per-field text length.
MAX_FIELD_LENGTH = 20_000


def _check_field_lengths(fields: dict[str, str | None]) -> None:
    for name, value in fields.items():
        if value is not None and len(value) > MAX_FIELD_LENGTH:
            raise ValueError(
                f"{name} is {len(value)} chars, exceeds the {MAX_FIELD_LENGTH}-char limit"
            )


def _resolve_commit_hash(provided: str | None) -> str | None:
    """Validate and resolve a commit hash, with auto-detection of HEAD.

    Returns a full 40-char hash, or ``None`` if there's no git repo and no
    hash was provided.  Raises ``ValueError`` when *provided* is non-empty
    but ``git rev-parse --verify`` can't find it in this repository.
    """

    def _git(args: list[str]) -> subprocess.CompletedProcess | None:
        try:
            return subprocess.run(
                ["git", *args],
                capture_output=True,
                text=True,
                timeout=5,
            )
        except FileNotFoundError:
            return None

    if provided:
        result = _git(["rev-parse", "--verify", provided])
        if result is None:
            # git not installed — trust the provided hash as-is.
            return provided
        if result.returncode != 0:
            raise ValueError(
                f"Commit hash {provided!r} not found in this repository"
            )
        return result.stdout.strip()

    # No hash provided — try to auto-detect HEAD.
    result = _git(["rev-parse", "HEAD"])
    if result and result.returncode == 0:
        return result.stdout.strip()
    return None

ENTRY_TYPES = frozenset(
    {
        "BUG",
        "REGRESSION",
        "UPSTREAM",
        "PERF",
        "DECISION",
        "LIMITATION",
        "GOTCHA",
        "INVARIANT",
    }
)

MEMORY_SELECT_COLUMNS = """
    memory_records.id,
    memory_records.entry_type,
    memory_records.title,
    memory_records.symptoms,
    memory_records.root_cause,
    memory_records.wrong_fixes,
    memory_records.resolution,
    memory_records.lesson,
    memory_records.file_path,
    memory_records.commit_hash,
    memory_records.created_at,
    memory_records.status,
    memory_records.namespace,
    memory_records.approved_by,
    memory_records.issue_url,
    memory_records.pr_url,
    memory_records.validating_test,
    memory_records.confidence,
    memory_records.derived_from
"""


class DatabaseManager:
    """Handles all SQLite operations with sqlite-vec vector search."""

    def __init__(self, db_path: str = ".codegrave.db", embedding_dim: int = 768):
        self.db_path = db_path
        self.embedding_dim = embedding_dim
        self._init_db()

    def _get_connection(self) -> sqlite3.Connection:
        if self.db_path == ":memory:":
            conn = sqlite3.connect(
                "file:codegrave_mem?mode=memory&cache=shared", uri=True
            )
        else:
            conn = sqlite3.connect(self.db_path)
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA busy_timeout = 5000")
        return conn

    def _init_db(self) -> None:
        with self._get_connection() as conn:
            # ── Metadata table (dimension tracking, future migrations) ──
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS codegrave_meta (
                    key   TEXT PRIMARY KEY,
                    value TEXT NOT NULL
                )
                """
            )
            # Check for dimension mismatch with existing DB.
            existing = conn.execute(
                "SELECT value FROM codegrave_meta WHERE key = 'embedding_dim'"
            ).fetchone()
            if existing is not None:
                stored_dim = int(existing["value"])
                if stored_dim != self.embedding_dim:
                    raise ValueError(
                        f"Existing database uses {stored_dim}-dimensional embeddings "
                        f"but the current configuration specifies {self.embedding_dim}. "
                        f"To migrate: delete the record_embeddings virtual table rows "
                        f"(or the entire .codegrave.db) and re-record your memories, "
                        f"or revert embedding_dim to {stored_dim} in your config."
                    )
            else:
                conn.execute(
                    "INSERT INTO codegrave_meta (key, value) VALUES ('embedding_dim', ?)",
                    (str(self.embedding_dim),),
                )

            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS memory_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    entry_type TEXT NOT NULL,
                    title TEXT NOT NULL,
                    symptoms TEXT NOT NULL,
                    root_cause TEXT NOT NULL,
                    wrong_fixes TEXT,
                    resolution TEXT NOT NULL,
                    lesson TEXT NOT NULL,
                    file_path TEXT,
                    commit_hash TEXT,
                    created_at TEXT NOT NULL
                )
                """
            )
            # Migrate pre-existing DBs that lack the commit_hash column.
            try:
                conn.execute(
                    "ALTER TABLE memory_records ADD COLUMN commit_hash TEXT"
                )
            except sqlite3.OperationalError:
                pass  # column already exists
            # Migrate pre-existing DBs that lack the namespace column.
            try:
                conn.execute(
                    "ALTER TABLE memory_records ADD COLUMN namespace TEXT NOT NULL DEFAULT '.'"
                )
            except sqlite3.OperationalError:
                pass  # column already exists
            # Migrate pre-existing DBs that lack the status column.
            try:
                conn.execute(
                    "ALTER TABLE memory_records ADD COLUMN status TEXT NOT NULL DEFAULT 'approved'"
                )
            except sqlite3.OperationalError:
                pass  # column already exists
            # Migrate pre-existing DBs that lack provenance columns.
            for col in ("approved_by", "issue_url", "pr_url", "validating_test", "confidence"):
                try:
                    conn.execute(f"ALTER TABLE memory_records ADD COLUMN {col} TEXT")
                except sqlite3.OperationalError:
                    pass  # column already exists
            # Migrate pre-existing DBs that lack the derived_from column
            # (comma-separated source record ids, set when a record is the
            # canonical result of a consolidate/merge action).
            try:
                conn.execute("ALTER TABLE memory_records ADD COLUMN derived_from TEXT")
            except sqlite3.OperationalError:
                pass  # column already exists
            conn.execute(
                f"""
                CREATE VIRTUAL TABLE IF NOT EXISTS record_embeddings USING vec0(
                    embedding float[{self.embedding_dim}]
                )
                """
            )
            # Holding table for consolidation proposals: an agent stages a
            # merge here after a human approves the grouping conversationally,
            # but nothing in memory_records changes until a human commits (or
            # discards) it offline via `codegrave --review`. Source records
            # stay untouched and searchable in the meantime.
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS consolidated (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    source_ids TEXT NOT NULL,
                    entry_type TEXT NOT NULL,
                    title TEXT NOT NULL,
                    symptoms TEXT NOT NULL,
                    root_cause TEXT NOT NULL,
                    wrong_fixes TEXT,
                    resolution TEXT NOT NULL,
                    lesson TEXT NOT NULL,
                    rationale TEXT NOT NULL,
                    created_at TEXT NOT NULL
                )
                """
            )
            # WAL mode allows concurrent readers and writers — essential
            # when multiple IDE MCP transports each have their own process
            # hitting the same .codegrave.db file.
            conn.execute("PRAGMA journal_mode=WAL")
            conn.commit()

    def insert_record(
        self,
        entry_type: str,
        title: str,
        symptoms: str,
        root_cause: str,
        wrong_fixes: str | None,
        resolution: str,
        lesson: str,
        file_path: str | None,
        embedding: list[float],
        commit_hash: str | None = None,
        status: str = "approved",
        namespace: str = ".",
        approved_by: str | None = None,
        issue_url: str | None = None,
        pr_url: str | None = None,
        validating_test: str | None = None,
        confidence: str | None = None,
        derived_from: str | None = None,
    ) -> int:
        entry_type = entry_type.strip().upper()
        if entry_type not in ENTRY_TYPES:
            raise ValueError(
                f"Invalid entry_type {entry_type!r}. "
                f"Expected one of: {', '.join(sorted(ENTRY_TYPES))}"
            )
        if len(embedding) != self.embedding_dim:
            raise ValueError(
                f"Expected embedding of length {self.embedding_dim}, got {len(embedding)}"
            )
        _check_field_lengths(
            {
                "title": title,
                "symptoms": symptoms,
                "root_cause": root_cause,
                "wrong_fixes": wrong_fixes,
                "resolution": resolution,
                "lesson": lesson,
                "file_path": file_path,
            }
        )

        commit_hash = _resolve_commit_hash(commit_hash or None)

        created_at = datetime.now(timezone.utc).isoformat()
        with self._get_connection() as conn:
            cur = conn.execute(
                """
                INSERT INTO memory_records (
                    entry_type, title, symptoms, root_cause, wrong_fixes,
                    resolution, lesson, file_path, commit_hash, created_at,
                    status, namespace, approved_by, issue_url, pr_url,
                    validating_test, confidence, derived_from
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    entry_type,
                    title,
                    symptoms,
                    root_cause,
                    wrong_fixes,
                    resolution,
                    lesson,
                    file_path,
                    commit_hash,
                    created_at,
                    status,
                    namespace,
                    approved_by,
                    issue_url,
                    pr_url,
                    validating_test,
                    confidence,
                    derived_from,
                ),
            )
            record_id = cur.lastrowid
            assert record_id is not None
            conn.execute(
                """
                INSERT INTO record_embeddings (rowid, embedding)
                VALUES (?, ?)
                """,
                (record_id, serialize_float32(embedding)),
            )
            conn.commit()
            return record_id

    def search_similar(
        self, query_embedding: list[float], limit: int = 15,
        namespace: str | None = None,
    ) -> list[dict[str, Any]]:
        if len(query_embedding) != self.embedding_dim:
            raise ValueError(
                f"Expected embedding of length {self.embedding_dim}, got {len(query_embedding)}"
            )

        ns = namespace if namespace is not None else "."
        ns_clause, ns_params = _namespace_clause(ns)
        with self._get_connection() as conn:
            rows = conn.execute(
                f"""
                SELECT
                    {MEMORY_SELECT_COLUMNS},
                    record_embeddings.distance AS distance
                FROM record_embeddings
                JOIN memory_records
                  ON memory_records.id = record_embeddings.rowid
                WHERE record_embeddings.embedding MATCH ?
                  AND k = ?
                  AND memory_records.status = 'approved'
                  AND {ns_clause}
                ORDER BY distance
                """,
                (serialize_float32(query_embedding), limit, *ns_params),
            ).fetchall()
            return [dict(row) for row in rows]

    def get_record(self, record_id: int) -> dict[str, Any] | None:
        with self._get_connection() as conn:
            row = conn.execute(
                "SELECT * FROM memory_records WHERE id = ?", (record_id,)
            ).fetchone()
            return dict(row) if row else None

    def update_record(
        self,
        record_id: int,
        *,
        entry_type: str | None = None,
        title: str | None = None,
        symptoms: str | None = None,
        root_cause: str | None = None,
        wrong_fixes: str | None = None,
        resolution: str | None = None,
        lesson: str | None = None,
        file_path: str | None = None,
        embedding: list[float] | None = None,
        status: str | None = None,
        approved_by: str | None = None,
        issue_url: str | None = None,
        pr_url: str | None = None,
        validating_test: str | None = None,
        confidence: str | None = None,
    ) -> bool:
        """Update the given fields on a record; a None argument leaves that field unchanged.

        Returns False if no record with *record_id* exists.
        """
        if entry_type is not None:
            entry_type = entry_type.strip().upper()
            if entry_type not in ENTRY_TYPES:
                raise ValueError(
                    f"Invalid entry_type {entry_type!r}. "
                    f"Expected one of: {', '.join(sorted(ENTRY_TYPES))}"
                )
        if embedding is not None and len(embedding) != self.embedding_dim:
            raise ValueError(
                f"Expected embedding of length {self.embedding_dim}, got {len(embedding)}"
            )
        _check_field_lengths(
            {
                "title": title,
                "symptoms": symptoms,
                "root_cause": root_cause,
                "wrong_fixes": wrong_fixes,
                "resolution": resolution,
                "lesson": lesson,
                "file_path": file_path,
            }
        )

        fields = {
            "entry_type": entry_type,
            "title": title,
            "symptoms": symptoms,
            "root_cause": root_cause,
            "wrong_fixes": wrong_fixes,
            "resolution": resolution,
            "lesson": lesson,
            "file_path": file_path,
            "status": status,
            "approved_by": approved_by,
            "issue_url": issue_url,
            "pr_url": pr_url,
            "validating_test": validating_test,
            "confidence": confidence,
        }
        fields = {k: v for k, v in fields.items() if v is not None}

        with self._get_connection() as conn:
            if fields:
                set_clause = ", ".join(f"{k} = ?" for k in fields)
                cur = conn.execute(
                    f"UPDATE memory_records SET {set_clause} WHERE id = ?",
                    (*fields.values(), record_id),
                )
                if cur.rowcount == 0:
                    return False
            else:
                row = conn.execute(
                    "SELECT 1 FROM memory_records WHERE id = ?", (record_id,)
                ).fetchone()
                if row is None:
                    return False

            if embedding is not None:
                conn.execute(
                    "DELETE FROM record_embeddings WHERE rowid = ?", (record_id,)
                )
                conn.execute(
                    "INSERT INTO record_embeddings (rowid, embedding) VALUES (?, ?)",
                    (record_id, serialize_float32(embedding)),
                )
            conn.commit()
            return True

    def delete_record(self, record_id: int) -> bool:
        """Delete a record and its embedding. Returns False if it didn't exist."""
        with self._get_connection() as conn:
            cur = conn.execute("DELETE FROM memory_records WHERE id = ?", (record_id,))
            conn.execute("DELETE FROM record_embeddings WHERE rowid = ?", (record_id,))
            conn.commit()
            return cur.rowcount > 0

    def keyword_search(
        self, query: str, limit: int = 15, namespace: str | None = None
    ) -> list[dict[str, Any]]:
        """Simple keyword (LIKE) search across memory text fields.

        Used by the ``agent`` embedding provider when no vector model is
        available.  Splits the query on whitespace and OR-s each term
        across title, symptoms, root_cause, resolution, and lesson.
        """
        terms = [t.strip() for t in query.split() if len(t.strip()) >= 2]
        ns = namespace if namespace is not None else "."
        ns_clause, ns_params = _namespace_clause(ns)
        if not terms:
            with self._get_connection() as conn:
                rows = conn.execute(
                    f"""
                    SELECT {MEMORY_SELECT_COLUMNS}, NULL AS distance
                    FROM memory_records
                    WHERE status = 'approved' AND {ns_clause}
                    ORDER BY created_at DESC
                    LIMIT ?
                    """,
                    (*ns_params, limit),
                ).fetchall()
                return [dict(row) for row in rows]

        clauses: list[str] = []
        params: list[str] = []
        for term in terms:
            pattern = f"%{term}%"
            clauses.append(
                "(title LIKE ? OR symptoms LIKE ? OR root_cause LIKE ? "
                "OR resolution LIKE ? OR lesson LIKE ?)"
            )
            params.extend([pattern] * 5)

        where = "(" + " OR ".join(clauses) + ")"
        with self._get_connection() as conn:
            rows = conn.execute(
                f"""
                SELECT {MEMORY_SELECT_COLUMNS}, NULL AS distance
                FROM memory_records
                WHERE {where} AND status = 'approved' AND {ns_clause}
                ORDER BY created_at DESC
                LIMIT ?
                """,
                (*params, *ns_params, limit),
            ).fetchall()
            return [dict(row) for row in rows]

    def get_all_records(self) -> list[dict[str, Any]]:
        with self._get_connection() as conn:
            rows = conn.execute(
                """
                SELECT * FROM memory_records
                ORDER BY entry_type ASC, created_at DESC
                """
            ).fetchall()
            return [dict(row) for row in rows]

    def get_pending_records(self) -> list[dict[str, Any]]:
        """Return records awaiting human review (status='pending')."""
        with self._get_connection() as conn:
            rows = conn.execute(
                """
                SELECT * FROM memory_records
                WHERE status = 'pending'
                ORDER BY created_at ASC
                """
            ).fetchall()
            return [dict(row) for row in rows]

    def insert_consolidation(
        self,
        source_ids: str,
        entry_type: str,
        title: str,
        symptoms: str,
        root_cause: str,
        wrong_fixes: str | None,
        resolution: str,
        lesson: str,
        rationale: str,
    ) -> int:
        """Stage a proposed merge for offline review. Does not touch memory_records."""
        entry_type = entry_type.strip().upper()
        if entry_type not in ENTRY_TYPES:
            raise ValueError(
                f"Invalid entry_type {entry_type!r}. "
                f"Expected one of: {', '.join(sorted(ENTRY_TYPES))}"
            )
        ids = [s.strip() for s in source_ids.split(",") if s.strip()]
        if len(ids) < 2:
            raise ValueError(
                f"source_ids must list at least 2 record ids to merge, got {source_ids!r}"
            )
        _check_field_lengths(
            {
                "title": title,
                "symptoms": symptoms,
                "root_cause": root_cause,
                "wrong_fixes": wrong_fixes,
                "resolution": resolution,
                "lesson": lesson,
                "rationale": rationale,
            }
        )

        created_at = datetime.now(timezone.utc).isoformat()
        with self._get_connection() as conn:
            cur = conn.execute(
                """
                INSERT INTO consolidated (
                    source_ids, entry_type, title, symptoms, root_cause,
                    wrong_fixes, resolution, lesson, rationale, created_at
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    ",".join(ids),
                    entry_type,
                    title,
                    symptoms,
                    root_cause,
                    wrong_fixes,
                    resolution,
                    lesson,
                    rationale,
                    created_at,
                ),
            )
            proposal_id = cur.lastrowid
            assert proposal_id is not None
            conn.commit()
            return proposal_id

    def get_consolidation_proposals(self) -> list[dict[str, Any]]:
        with self._get_connection() as conn:
            rows = conn.execute(
                "SELECT * FROM consolidated ORDER BY created_at ASC"
            ).fetchall()
            return [dict(row) for row in rows]

    def delete_consolidation_proposal(self, proposal_id: int) -> bool:
        with self._get_connection() as conn:
            cur = conn.execute("DELETE FROM consolidated WHERE id = ?", (proposal_id,))
            conn.commit()
            return cur.rowcount > 0


def _namespace_clause(namespace: str) -> tuple[str, list[str]]:
    """Build a SQL clause that matches *namespace* and all subdirectories.

    ``"."`` matches everything (no filter).  ``"frontend"`` matches
    ``"frontend"`` exactly and ``"frontend/%"`` (e.g. ``"frontend/components"``).
    """
    if namespace == ".":
        return ("1 = 1", [])
    return (
        "(memory_records.namespace = ? OR memory_records.namespace LIKE ?)",
        [namespace, namespace + "/%"],
    )


_db: DatabaseManager | None = None


def init_db(db_path: str = ".codegrave.db", embedding_dim: int = EMBEDDING_DIM) -> DatabaseManager:
    global _db
    _db = DatabaseManager(db_path, embedding_dim=embedding_dim)
    return _db


def get_db() -> DatabaseManager:
    if _db is None:
        return init_db()
    return _db


def insert_record(
    entry_type: str,
    title: str,
    symptoms: str,
    root_cause: str,
    wrong_fixes: str | None,
    resolution: str,
    lesson: str,
    file_path: str | None,
    embedding: list[float],
    commit_hash: str | None = None,
    status: str = "approved",
    namespace: str = ".",
    approved_by: str | None = None,
    issue_url: str | None = None,
    pr_url: str | None = None,
    validating_test: str | None = None,
    confidence: str | None = None,
    derived_from: str | None = None,
) -> int:
    return get_db().insert_record(
        entry_type=entry_type,
        title=title,
        symptoms=symptoms,
        root_cause=root_cause,
        wrong_fixes=wrong_fixes,
        resolution=resolution,
        lesson=lesson,
        file_path=file_path,
        embedding=embedding,
        commit_hash=commit_hash,
        status=status,
        namespace=namespace,
        approved_by=approved_by,
        issue_url=issue_url,
        pr_url=pr_url,
        validating_test=validating_test,
        confidence=confidence,
        derived_from=derived_from,
    )


def search_vectors(query_embedding: list[float], limit: int = 15) -> list[dict[str, Any]]:
    return get_db().search_similar(query_embedding, limit=limit)


def get_record(record_id: int) -> dict[str, Any] | None:
    return get_db().get_record(record_id)


def get_all_records() -> list[dict[str, Any]]:
    return get_db().get_all_records()


def get_pending_records() -> list[dict[str, Any]]:
    return get_db().get_pending_records()


def update_record(record_id: int, **fields: Any) -> bool:
    return get_db().update_record(record_id, **fields)


def delete_record(record_id: int) -> bool:
    return get_db().delete_record(record_id)


def insert_consolidation(
    source_ids: str,
    entry_type: str,
    title: str,
    symptoms: str,
    root_cause: str,
    wrong_fixes: str | None,
    resolution: str,
    lesson: str,
    rationale: str,
) -> int:
    return get_db().insert_consolidation(
        source_ids=source_ids,
        entry_type=entry_type,
        title=title,
        symptoms=symptoms,
        root_cause=root_cause,
        wrong_fixes=wrong_fixes,
        resolution=resolution,
        lesson=lesson,
        rationale=rationale,
    )


def get_consolidation_proposals() -> list[dict[str, Any]]:
    return get_db().get_consolidation_proposals()


def delete_consolidation_proposal(proposal_id: int) -> bool:
    return get_db().delete_consolidation_proposal(proposal_id)


def embedding_text(
    title: str,
    symptoms: str,
    root_cause: str,
    resolution: str,
    lesson: str,
) -> str:
    """Concatenate structured fields used to build the memory embedding."""
    return "\n".join(
        part.strip()
        for part in (title, symptoms, root_cause, resolution, lesson)
        if part and part.strip()
    )