"""SQLite-backed persistent index for symbols, references, and dependencies.

Stores parsed AST data so that ASTService can skip re-parsing unchanged
files on restart.  Modeled after ``VectorStore`` (WAL mode, mtime tracking).

Database location: ``.attocode/index/symbols.db``
"""

from __future__ import annotations

import logging
import os
import sqlite3
import threading
import time
from dataclasses import dataclass, field
from itertools import islice
from typing import Any

from filelock import FileLock

logger = logging.getLogger(__name__)

SCHEMA_VERSION = "7"
# Bump the suffix when a change alters the saved graph edges, for example import resolution.
# Stores keep their tables, and the next start rebuilds the index once.
SNAPSHOT_KEY = "structural_snapshot_v4"
AST_SNAPSHOT_KEY = "structural_ast_summaries_v1"
SQLITE_BUSY_TIMEOUT_MS = 30_000
INDEX_INIT_LOCK_TIMEOUT_SECONDS = 30


def _insert_analysis_rows(conn, statement: str, rows, columns: int) -> None:
    """Insert bounded groups without one SQLite step per evidence row.

    At most 64 rows / 576 parameters, below SQLite's legacy 999-variable limit.
    The caller owns the transaction; iteration and SQL errors propagate so all
    replacements roll back together. Input order (and generated IDs) is retained.
    """
    rows = iter(rows)
    placeholder = "(" + ",".join(["?"] * columns) + ")"
    while batch := list(islice(rows, 64)):
        conn.execute(
            statement + " VALUES " + ",".join([placeholder] * len(batch)),
            tuple(value for row in batch for value in row),
        )


class IndexSchemaVersionError(RuntimeError):
    """The persisted structural index cannot be opened by this client version."""


@dataclass(slots=True)
class StoredFile:
    """Cached file metadata."""

    path: str
    mtime: float
    size: int
    language: str
    line_count: int
    content_hash: str


@dataclass(slots=True)
class StoredSymbol:
    """A persisted symbol definition."""

    id: int
    file_path: str
    name: str
    qualified_name: str
    kind: str
    line: int
    end_line: int
    source: str  # "tree-sitter" | "lsp"


@dataclass(slots=True)
class StoredReference:
    """A persisted symbol reference."""

    id: int
    file_path: str
    symbol_name: str
    ref_kind: str
    line: int
    column: int
    source: str  # "tree-sitter" | "lsp"
    syntax_name: str = ""  # identifier/member expression as written in source
    caller_qualified_name: str = ""  # enclosing function/method (call-graph edge)
    definition_key: str = ""


@dataclass(slots=True)
class IndexStore:
    """SQLite-backed persistent index store.

    Usage::

        store = IndexStore(db_path=".attocode/index/symbols.db")
        store.save_file(StoredFile(...))
        store.save_symbols("src/main.py", [StoredSymbol(...)])
        stale = store.get_stale_files({"src/main.py": 1234567.0})
    """

    db_path: str
    _conn: sqlite3.Connection | None = field(default=None, repr=False)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
    # When the on-disk schema differed from ``SCHEMA_VERSION`` at startup,
    # we wipe and rebuild — this records the prior version so callers
    # (e.g. ``readiness_report``) can surface a "schema migrated, full
    # re-index pending" hint instead of leaving users guessing why the
    # first request after upgrade is slow.
    _schema_rebuilt_from: str | None = field(default=None, repr=False)

    def __post_init__(self) -> None:
        os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
        try:
            if self.db_path == ":memory:":
                self._initialize_connection()
            else:
                # Multiple MCP clients commonly share one workspace. Serialize
                # schema inspection/rebuild so simultaneous restarts do not race
                # over DDL while hydration is opening the same SQLite store.
                with FileLock(f"{self.db_path}.init.lock", timeout=INDEX_INIT_LOCK_TIMEOUT_SECONDS):
                    self._initialize_connection()
        except Exception:
            if self._conn is not None:
                self._conn.close()
                self._conn = None
            raise

    def _initialize_connection(self) -> None:
        self._conn = sqlite3.connect(
            self.db_path,
            check_same_thread=False,
            timeout=SQLITE_BUSY_TIMEOUT_MS / 1000,
        )
        self._conn.execute(f"PRAGMA busy_timeout={SQLITE_BUSY_TIMEOUT_MS}")
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.execute("PRAGMA foreign_keys=ON")
        self._migrate_stale_schema()
        self._create_tables()
        self._check_schema_version()

    @staticmethod
    def _compare_schema_version(stored: str) -> int:
        try:
            stored_version = int(stored)
            current_version = int(SCHEMA_VERSION)
        except ValueError as error:
            raise IndexSchemaVersionError(
                f"Unrecognized structural index schema {stored!r}; "
                "upgrade Attocode or rebuild the workspace index."
            ) from error
        return (stored_version > current_version) - (stored_version < current_version)

    def _get_conn(self) -> sqlite3.Connection:
        if self._conn is None:
            raise RuntimeError("IndexStore connection is closed")
        return self._conn

    def _migrate_stale_schema(self) -> None:
        """Drop stale data tables BEFORE (re)creating them when the on-disk
        schema predates the current version.

        ``_create_tables`` builds the v3 indexes (e.g. ``ix_refs_caller`` on
        ``refs(caller_qualified_name)``). On a pre-v3 DB whose ``refs`` table
        lacks that column, creating the index raises ``OperationalError`` and
        the store can't even open — so the wipe in ``_check_schema_version``
        never runs. Dropping the stale tables here keeps the upgrade path
        working; data is rebuilt from source on the next index. Children are
        dropped before ``files`` to respect FK ordering.
        """
        conn = self._get_conn()
        try:
            row = conn.execute(
                "SELECT value FROM metadata WHERE key = 'schema_version'"
            ).fetchone()
        except sqlite3.OperationalError:
            return  # no metadata table => brand-new DB, nothing to migrate
        if not row or row[0] == SCHEMA_VERSION:
            return
        if self._compare_schema_version(row[0]) > 0:
            raise IndexSchemaVersionError(
                f"Structural index schema {row[0]} is newer than this client's "
                f"schema {SCHEMA_VERSION}; restart using the upgraded Attocode executable."
            )
        logger.warning(
            "Index schema is stale (%s, expected %s); dropping data tables in "
            "%r so the new schema can be created — a full re-index will run on "
            "next access.",
            row[0], SCHEMA_VERSION, self.db_path,
        )
        conn.executescript(
            "DROP TABLE IF EXISTS refs;"
            "DROP TABLE IF EXISTS symbols;"
            "DROP TABLE IF EXISTS dependencies;"
            "DROP TABLE IF EXISTS files;"
        )
        conn.commit()

    def _create_tables(self) -> None:
        conn = self._get_conn()
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS metadata (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS files (
                path TEXT PRIMARY KEY,
                mtime REAL NOT NULL,
                size INTEGER NOT NULL,
                language TEXT NOT NULL DEFAULT '',
                line_count INTEGER NOT NULL DEFAULT 0,
                content_hash TEXT NOT NULL DEFAULT ''
            );

            CREATE TABLE IF NOT EXISTS symbols (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                file_path TEXT NOT NULL REFERENCES files(path) ON DELETE CASCADE,
                name TEXT NOT NULL,
                qualified_name TEXT NOT NULL,
                kind TEXT NOT NULL,
                line INTEGER NOT NULL DEFAULT 0,
                end_line INTEGER NOT NULL DEFAULT 0,
                source TEXT NOT NULL DEFAULT 'tree-sitter'
            );

            CREATE TABLE IF NOT EXISTS refs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                file_path TEXT NOT NULL REFERENCES files(path) ON DELETE CASCADE,
                symbol_name TEXT NOT NULL,
                ref_kind TEXT NOT NULL DEFAULT 'call',
                line INTEGER NOT NULL DEFAULT 0,
                col INTEGER NOT NULL DEFAULT 0,
                source TEXT NOT NULL DEFAULT 'tree-sitter',
                syntax_name TEXT NOT NULL DEFAULT '',
                caller_qualified_name TEXT NOT NULL DEFAULT '',
                definition_key TEXT NOT NULL DEFAULT ''
            );

            CREATE TABLE IF NOT EXISTS dependencies (
                source_path TEXT NOT NULL,
                target_path TEXT NOT NULL,
                PRIMARY KEY (source_path, target_path)
            );

            CREATE INDEX IF NOT EXISTS ix_symbols_name ON symbols(name);
            CREATE INDEX IF NOT EXISTS ix_symbols_qname ON symbols(qualified_name);
            CREATE INDEX IF NOT EXISTS ix_symbols_file ON symbols(file_path);
            CREATE INDEX IF NOT EXISTS ix_refs_symbol ON refs(symbol_name);
            CREATE INDEX IF NOT EXISTS ix_refs_file ON refs(file_path);
            CREATE INDEX IF NOT EXISTS ix_refs_caller ON refs(caller_qualified_name);
            CREATE INDEX IF NOT EXISTS ix_deps_target ON dependencies(target_path);
        """)
        conn.commit()

    def _check_schema_version(self) -> None:
        conn = self._get_conn()
        row = conn.execute(
            "SELECT value FROM metadata WHERE key = 'schema_version'"
        ).fetchone()
        if row:
            if row[0] != SCHEMA_VERSION:
                if self._compare_schema_version(row[0]) > 0:
                    raise IndexSchemaVersionError(
                        f"Structural index schema {row[0]} is newer than this client's "
                        f"schema {SCHEMA_VERSION}; restart using the upgraded Attocode executable."
                    )
                # Schema migrations clear all persisted index data — the
                # first request after upgrade will trigger a full
                # re-index that may take a while on big repos. We log
                # this prominently and stash the prior version on the
                # store instance + in the SQLite metadata table so a
                # future readiness probe can attribute a slow first
                # request to the rebuild rather than to indexing bugs.
                logger.warning(
                    "Index schema changed (%s -> %s); clearing %r and "
                    "rebuilding from source on next access — full "
                    "re-index will run on demand. Subsequent requests "
                    "may be slow until indexing catches up.",
                    row[0], SCHEMA_VERSION, self.db_path,
                )
                self.clear_all()
                conn.execute(
                    "INSERT OR REPLACE INTO metadata (key, value) VALUES ('schema_version', ?)",
                    (SCHEMA_VERSION,),
                )
                conn.execute(
                    "INSERT OR REPLACE INTO metadata (key, value) "
                    "VALUES ('schema_rebuilt_from_version', ?)",
                    (row[0],),
                )
                conn.commit()
                self._schema_rebuilt_from = row[0]
        else:
            conn.execute(
                "INSERT OR REPLACE INTO metadata (key, value) VALUES ('schema_version', ?)",
                (SCHEMA_VERSION,),
            )
            conn.commit()

    # ------------------------------------------------------------------
    # Metadata
    # ------------------------------------------------------------------

    def get_meta(self, key: str, *, max_bytes: int | None = None) -> str | None:
        conn = self._get_conn()
        with self._lock:
            row = conn.execute(
                "SELECT value FROM metadata WHERE key = ? AND (? IS NULL OR length(CAST(value AS BLOB)) <= ?)",
                (key, max_bytes, max_bytes)
            ).fetchone()
        return row[0] if row else None

    def set_meta(self, key: str, value: str) -> None:
        conn = self._get_conn()
        with self._lock:
            conn.execute(
                "INSERT OR REPLACE INTO metadata (key, value) VALUES (?, ?)",
                (key, value),
            )
            conn.commit()

    # ------------------------------------------------------------------
    # Files
    # ------------------------------------------------------------------

    def save_file(self, f: StoredFile) -> None:
        conn = self._get_conn()
        with self._lock:
            conn.execute(
                """INSERT OR REPLACE INTO files
                   (path, mtime, size, language, line_count, content_hash)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (f.path, f.mtime, f.size, f.language, f.line_count, f.content_hash),
            )
            conn.commit()

    def save_files_batch(self, files: list[StoredFile]) -> None:
        conn = self._get_conn()
        rows = [
            (f.path, f.mtime, f.size, f.language, f.line_count, f.content_hash)
            for f in files
        ]
        with self._lock:
            conn.executemany(
                """INSERT OR REPLACE INTO files
                   (path, mtime, size, language, line_count, content_hash)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                rows,
            )
            conn.commit()

    def get_file(self, path: str) -> StoredFile | None:
        conn = self._get_conn()
        with self._lock:
            row = conn.execute(
                "SELECT path, mtime, size, language, line_count, content_hash "
                "FROM files WHERE path = ?",
                (path,),
            ).fetchone()
        if row:
            return StoredFile(*row)
        return None

    def get_all_files(self) -> list[StoredFile]:
        conn = self._get_conn()
        with self._lock:
            rows = conn.execute(
                "SELECT path, mtime, size, language, line_count, content_hash FROM files"
            ).fetchall()
        return [StoredFile(*r) for r in rows]

    def get_stale_files(self, file_mtimes: dict[str, float]) -> list[str]:
        """Return file paths where current mtime > stored mtime (or not indexed)."""
        conn = self._get_conn()
        with self._lock:
            rows = conn.execute("SELECT path, mtime FROM files").fetchall()
        indexed = {row[0]: row[1] for row in rows}
        stale: list[str] = []
        for fpath, current_mtime in file_mtimes.items():
            stored_mtime = indexed.get(fpath)
            if stored_mtime is None or current_mtime > stored_mtime:
                stale.append(fpath)
        return stale

    def get_deleted_files(self, current_files: set[str]) -> list[str]:
        """Return indexed files no longer present on disk."""
        conn = self._get_conn()
        with self._lock:
            rows = conn.execute("SELECT path FROM files").fetchall()
        return [r[0] for r in rows if r[0] not in current_files]

    def remove_file(self, path: str) -> None:
        """Remove a file and all its symbols/refs (CASCADE)."""
        conn = self._get_conn()
        with self._lock:
            conn.execute("DELETE FROM files WHERE path = ?", (path,))
            conn.execute(
                "DELETE FROM dependencies WHERE source_path = ? OR target_path = ?",
                (path, path),
            )
            conn.commit()

    # ------------------------------------------------------------------
    # Symbols
    # ------------------------------------------------------------------

    def save_symbols(self, file_path: str, symbols: list[dict[str, Any]]) -> None:
        """Save symbol definitions without replacing the file's references."""
        self.save_analysis_batch([(file_path, symbols, None)])

    def save_analysis_batch(self, files) -> None:
        """Atomically replace selected files' symbols/references; None preserves a table.

        Hydration supplies at most 100 files per batch. A failed insertion rolls
        back every replacement in the batch, including earlier files.
        """
        conn = self._get_conn()
        with self._lock, conn:
            for file_path, symbols, refs in files:
                if symbols is not None:
                    conn.execute("DELETE FROM symbols WHERE file_path = ?", (file_path,))
                    _insert_analysis_rows(
                        conn,
                        """INSERT INTO symbols
                           (file_path, name, qualified_name, kind, line, end_line, source)""",
                        ((file_path, s["name"], s["qualified_name"], s["kind"],
                          s.get("line", 0), s.get("end_line", 0), s.get("source", "tree-sitter"))
                         for s in symbols),
                        7,
                    )
                if refs is not None:
                    conn.execute("DELETE FROM refs WHERE file_path = ?", (file_path,))
                    _insert_analysis_rows(
                        conn,
                        """INSERT INTO refs
                           (file_path, symbol_name, ref_kind, line, col, source,
                            syntax_name, caller_qualified_name, definition_key)""",
                        ((file_path, r["symbol_name"], r.get("ref_kind", "call"),
                          r.get("line", 0), r.get("column", 0), r.get("source", "tree-sitter"),
                          r.get("syntax_name", ""), r.get("caller_qualified_name", ""),
                          r.get("definition_key", "")) for r in refs),
                        9,
                    )

    def load_symbols(self, file_path: str | None = None) -> list[StoredSymbol]:
        """Load symbols, optionally filtered by file."""
        conn = self._get_conn()
        with self._lock:
            if file_path:
                rows = conn.execute(
                    "SELECT id, file_path, name, qualified_name, kind, line, end_line, source "
                    "FROM symbols WHERE file_path = ?",
                    (file_path,),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT id, file_path, name, qualified_name, kind, line, end_line, source "
                    "FROM symbols"
                ).fetchall()
        return [
            StoredSymbol(
                id=r[0], file_path=r[1], name=r[2], qualified_name=r[3],
                kind=r[4], line=r[5], end_line=r[6], source=r[7],
            )
            for r in rows
        ]

    # ------------------------------------------------------------------
    # References
    # ------------------------------------------------------------------

    def save_references(self, file_path: str, refs: list[dict[str, Any]]) -> None:
        """Save references without replacing the file's definitions."""
        self.save_analysis_batch([(file_path, None, refs)])

    def load_references(self, file_path: str | None = None) -> list[StoredReference]:
        """Load references, optionally filtered by file."""
        conn = self._get_conn()
        with self._lock:
            if file_path:
                rows = conn.execute(
                    "SELECT id, file_path, symbol_name, ref_kind, line, col, "
                    "source, syntax_name, caller_qualified_name, definition_key FROM refs WHERE file_path = ?",
                    (file_path,),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT id, file_path, symbol_name, ref_kind, line, col, "
                    "source, syntax_name, caller_qualified_name, definition_key FROM refs"
                ).fetchall()
        return [
            StoredReference(
                id=r[0], file_path=r[1], symbol_name=r[2], ref_kind=r[3],
                line=r[4], column=r[5], source=r[6], syntax_name=r[7],
                caller_qualified_name=r[8], definition_key=r[9],
            )
            for r in rows
        ]

    # ------------------------------------------------------------------
    # Dependencies
    # ------------------------------------------------------------------

    def save_dependencies(self, source_path: str, targets: set[str]) -> None:
        """Save dependency edges from source to targets (replaces existing)."""
        conn = self._get_conn()
        with self._lock:
            conn.execute(
                "DELETE FROM dependencies WHERE source_path = ?", (source_path,)
            )
            if targets:
                conn.executemany(
                    "INSERT OR IGNORE INTO dependencies (source_path, target_path) VALUES (?, ?)",
                    [(source_path, t) for t in targets],
                )
            conn.commit()

    def save_dependencies_batch(self, edges: list[tuple[str, str]]) -> None:
        """Bulk save dependency edges."""
        conn = self._get_conn()
        with self._lock:
            conn.executemany(
                "INSERT OR IGNORE INTO dependencies (source_path, target_path) VALUES (?, ?)",
                edges,
            )
            conn.commit()

    def load_dependencies(self) -> dict[str, set[str]]:
        """Load all dependency edges as forward map."""
        conn = self._get_conn()
        with self._lock:
            rows = conn.execute(
                "SELECT source_path, target_path FROM dependencies"
            ).fetchall()
        forward: dict[str, set[str]] = {}
        for src, tgt in rows:
            forward.setdefault(src, set()).add(tgt)
        return forward

    def clear_dependencies(self) -> None:
        conn = self._get_conn()
        with self._lock:
            conn.execute("DELETE FROM dependencies")
            conn.commit()

    # ------------------------------------------------------------------
    # Bulk operations
    # ------------------------------------------------------------------

    def clear_all(self) -> None:
        """Drop all data (keeps schema)."""
        conn = self._get_conn()
        with self._lock:
            conn.executescript(f"""
                DELETE FROM refs;
                DELETE FROM symbols;
                DELETE FROM dependencies;
                DELETE FROM files;
                DELETE FROM metadata WHERE key IN ('{SNAPSHOT_KEY}', '{AST_SNAPSHOT_KEY}');
            """)
            conn.commit()

    def stats(self) -> dict[str, int]:
        """Return counts of files, symbols, refs, and dependencies."""
        conn = self._get_conn()
        with self._lock:
            files = conn.execute("SELECT COUNT(*) FROM files").fetchone()[0]
            symbols = conn.execute("SELECT COUNT(*) FROM symbols").fetchone()[0]
            refs = conn.execute("SELECT COUNT(*) FROM refs").fetchone()[0]
            deps = conn.execute("SELECT COUNT(*) FROM dependencies").fetchone()[0]
        return {
            "files": files,
            "symbols": symbols,
            "references": refs,
            "dependencies": deps,
        }

    def record_scan_time(self) -> None:
        """Record current time as last full scan timestamp."""
        self.set_meta("last_full_scan", str(time.time()))

    def get_last_scan_time(self) -> float | None:
        """Return last full scan timestamp, or None if never scanned."""
        val = self.get_meta("last_full_scan")
        return float(val) if val else None

    def close(self) -> None:
        # Background hydration can be inside a write. Closing under it crashes sqlite3.
        with self._lock:
            if self._conn:
                self._conn.close()
                self._conn = None

    def __del__(self) -> None:
        self.close()
