"""
Memory Store — project-local storage for Kodara indices.

Storage: {project_root}/.kodara/index.db  (SQLite)

Why SQLite:
  - Zero dependencies (built into Python)
  - Single file, lives alongside the project (like JSON did)
  - Handles 100k+ files without loading everything into RAM
  - Efficient partial reads: get one module, check one hash
  - Simple migration path to PostgreSQL when cloud sync is needed

Schema:
  meta        — repo metadata (key/value)
  modules     — one row per file (path + JSON blob)
  graph_edges — dependency edges (from, to, kind)
  file_hashes — SHA-256 per file for incremental updates

Backward compatibility:
  If .kodara/index.json exists, it is auto-migrated to index.db on first load.
"""

from __future__ import annotations

import json
import time
import sqlite3
import hashlib
from contextlib import contextmanager
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Optional


KODARA_DIR   = ".kodara"
INDEX_FILE   = "index.db"
LEGACY_FILE  = "index.json"
INDEX_VERSION = 3          # bumped from 2 (JSON) → 3 (SQLite)

# ── Schema ────────────────────────────────────────────────────────────────────

_DDL = """
PRAGMA journal_mode = WAL;
PRAGMA synchronous  = NORMAL;

CREATE TABLE IF NOT EXISTS meta (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS modules (
    path TEXT PRIMARY KEY,
    data TEXT NOT NULL          -- JSON blob of module metadata
);

CREATE TABLE IF NOT EXISTS graph_edges (
    from_path TEXT NOT NULL,
    to_path   TEXT NOT NULL,
    kind      TEXT NOT NULL DEFAULT 'import',
    PRIMARY KEY (from_path, to_path, kind)
);

CREATE TABLE IF NOT EXISTS file_hashes (
    path TEXT PRIMARY KEY,
    hash TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_edges_from ON graph_edges(from_path);
CREATE INDEX IF NOT EXISTS idx_edges_to   ON graph_edges(to_path);
"""


# ── KodaraIndex dataclass (public API unchanged) ──────────────────────────────

@dataclass
class KodaraIndex:
    """Top-level index — same interface as before, now backed by SQLite."""
    repo_id:         str
    display_name:    str
    source_root:     str
    indexed_at:      float
    index_version:   int
    file_count:      int
    language_counts: dict[str, int]
    modules:         dict[str, dict]   # path → module metadata
    graph_edges:     list[dict]        # [{from, to, kind}]
    file_hashes:     dict[str, str]    # path → SHA-256


# ── MemoryStore ───────────────────────────────────────────────────────────────

class MemoryStore:
    """
    Project-local memory store backed by SQLite.

    Public API is identical to the previous JSON-based store:
        store.save(index)
        store.load() → KodaraIndex | None
        store.exists() → bool
        store.delete()
        store.get_file_hashes() → dict

    Extra efficient methods (avoid loading everything into RAM):
        store.get_module(path) → dict | None
        store.get_modules_batch(paths) → dict[path, dict]
        store.update_module(path, data)
        store.get_changed_files(current_hashes) → (added, modified, deleted)
    """

    def __init__(self, project_root: str | Path):
        self.project_root = Path(project_root).resolve()
        self.kodara_dir   = self.project_root / KODARA_DIR
        self.db_path      = self.kodara_dir / INDEX_FILE
        self.legacy_path  = self.kodara_dir / LEGACY_FILE

    @classmethod
    def for_cwd(cls) -> "MemoryStore":
        return cls(Path.cwd())

    @classmethod
    def find(cls, path: str | Path) -> "MemoryStore":
        return cls(Path(path).resolve())

    # ── Internal ──────────────────────────────────────────────────────────────

    @contextmanager
    def _connect(self):
        """Context manager — yields a connection with WAL mode enabled."""
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.close()

    def _ensure_schema(self, conn: sqlite3.Connection) -> None:
        conn.executescript(_DDL)

    # ── Directory / Init ──────────────────────────────────────────────────────

    def init_dir(self) -> None:
        self.kodara_dir.mkdir(parents=True, exist_ok=True)

    def _init_db(self) -> None:
        self.init_dir()
        with self._connect() as conn:
            self._ensure_schema(conn)

    # ── Save ──────────────────────────────────────────────────────────────────

    def save(self, index: KodaraIndex) -> None:
        """Persist the full index to SQLite (replaces previous content)."""
        self._init_db()
        with self._connect() as conn:
            self._ensure_schema(conn)

            # Meta
            meta = {
                "repo_id":         index.repo_id,
                "display_name":    index.display_name,
                "source_root":     index.source_root,
                "indexed_at":      str(index.indexed_at),
                "index_version":   str(INDEX_VERSION),
                "file_count":      str(index.file_count),
                "language_counts": json.dumps(index.language_counts),
            }
            conn.executemany(
                "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)",
                meta.items(),
            )

            # Modules — replace all
            conn.execute("DELETE FROM modules")
            conn.executemany(
                "INSERT INTO modules (path, data) VALUES (?, ?)",
                [(path, json.dumps(data)) for path, data in index.modules.items()],
            )

            # Graph edges — replace all
            conn.execute("DELETE FROM graph_edges")
            conn.executemany(
                "INSERT OR IGNORE INTO graph_edges (from_path, to_path, kind) VALUES (?, ?, ?)",
                [(e["from"], e["to"], e.get("kind", "import")) for e in index.graph_edges],
            )

            # File hashes — replace all
            conn.execute("DELETE FROM file_hashes")
            conn.executemany(
                "INSERT INTO file_hashes (path, hash) VALUES (?, ?)",
                index.file_hashes.items(),
            )

    # ── Load ──────────────────────────────────────────────────────────────────

    def load(self) -> Optional[KodaraIndex]:
        """Load full index. Auto-migrates from JSON if needed."""
        # Auto-migrate legacy JSON → SQLite
        if not self.db_path.exists() and self.legacy_path.exists():
            self._migrate_from_json()

        if not self.db_path.exists():
            return None

        try:
            with self._connect() as conn:
                self._ensure_schema(conn)

                # Meta
                rows = conn.execute("SELECT key, value FROM meta").fetchall()
                if not rows:
                    return None
                meta = {r["key"]: r["value"] for r in rows}

                version = int(meta.get("index_version", 0))
                if version < 2:   # accept both v2 (migrated) and v3+
                    return None

                # Modules
                modules = {
                    row["path"]: json.loads(row["data"])
                    for row in conn.execute("SELECT path, data FROM modules")
                }

                # Graph edges
                graph_edges = [
                    {"from": row["from_path"], "to": row["to_path"], "kind": row["kind"]}
                    for row in conn.execute("SELECT from_path, to_path, kind FROM graph_edges")
                ]

                # File hashes
                file_hashes = {
                    row["path"]: row["hash"]
                    for row in conn.execute("SELECT path, hash FROM file_hashes")
                }

                return KodaraIndex(
                    repo_id         = meta.get("repo_id", ""),
                    display_name    = meta.get("display_name", ""),
                    source_root     = meta.get("source_root", ""),
                    indexed_at      = float(meta.get("indexed_at", 0)),
                    index_version   = version,
                    file_count      = int(meta.get("file_count", 0)),
                    language_counts = json.loads(meta.get("language_counts", "{}")),
                    modules         = modules,
                    graph_edges     = graph_edges,
                    file_hashes     = file_hashes,
                )
        except Exception:
            return None

    # ── Exists / Delete ───────────────────────────────────────────────────────

    def exists(self) -> bool:
        return self.db_path.exists() or self.legacy_path.exists()

    def delete(self) -> bool:
        deleted = False
        if self.db_path.exists():
            self.db_path.unlink()
            deleted = True
        if self.legacy_path.exists():
            self.legacy_path.unlink()
            deleted = True
        return deleted

    # ── Efficient partial reads (new — avoids full load) ──────────────────────

    def get_module(self, path: str) -> Optional[dict]:
        """Get a single module without loading the entire index."""
        if not self.db_path.exists():
            return None
        try:
            with self._connect() as conn:
                row = conn.execute(
                    "SELECT data FROM modules WHERE path = ?", (path,)
                ).fetchone()
                return json.loads(row["data"]) if row else None
        except Exception:
            return None

    def get_modules_batch(self, paths: list[str]) -> dict[str, dict]:
        """Get multiple modules by path without loading everything."""
        if not self.db_path.exists() or not paths:
            return {}
        try:
            placeholders = ",".join("?" * len(paths))
            with self._connect() as conn:
                rows = conn.execute(
                    f"SELECT path, data FROM modules WHERE path IN ({placeholders})",
                    paths,
                ).fetchall()
                return {row["path"]: json.loads(row["data"]) for row in rows}
        except Exception:
            return {}

    def update_module(self, path: str, data: dict) -> None:
        """Update a single module (for incremental updates)."""
        self._init_db()
        with self._connect() as conn:
            self._ensure_schema(conn)
            conn.execute(
                "INSERT OR REPLACE INTO modules (path, data) VALUES (?, ?)",
                (path, json.dumps(data)),
            )

    def get_file_hashes(self) -> dict[str, str]:
        """Get all file hashes — used by incremental updater."""
        if not self.db_path.exists():
            # Try legacy JSON
            idx = self.load()
            return idx.file_hashes if idx else {}
        try:
            with self._connect() as conn:
                return {
                    row["path"]: row["hash"]
                    for row in conn.execute("SELECT path, hash FROM file_hashes")
                }
        except Exception:
            return {}

    def get_changed_files(
        self, current_hashes: dict[str, str]
    ) -> tuple[set[str], set[str], set[str]]:
        """
        Compare current file hashes against stored ones.
        Returns: (added, modified, deleted)
        """
        stored = self.get_file_hashes()
        stored_paths  = set(stored.keys())
        current_paths = set(current_hashes.keys())

        added    = current_paths - stored_paths
        deleted  = stored_paths - current_paths
        modified = {
            p for p in current_paths & stored_paths
            if current_hashes[p] != stored[p]
        }
        return added, modified, deleted

    # ── Migration: JSON → SQLite ───────────────────────────────────────────────

    def _migrate_from_json(self) -> None:
        """Silently migrate legacy index.json to index.db."""
        try:
            data = json.loads(self.legacy_path.read_text(encoding="utf-8"))
            index = KodaraIndex(
                repo_id         = data.get("repo_id", ""),
                display_name    = data.get("display_name", ""),
                source_root     = data.get("source_root", ""),
                indexed_at      = float(data.get("indexed_at", time.time())),
                index_version   = INDEX_VERSION,
                file_count      = int(data.get("file_count", 0)),
                language_counts = data.get("language_counts", {}),
                modules         = data.get("modules", {}),
                graph_edges     = data.get("graph_edges", []),
                file_hashes     = data.get("file_hashes", {}),
            )
            self.save(index)
            # Rename old file so it doesn't trigger migration again
            self.legacy_path.rename(self.legacy_path.with_suffix(".json.bak"))
        except Exception:
            pass  # If migration fails, next full init will rebuild

    # ── Repo ID ───────────────────────────────────────────────────────────────

    @staticmethod
    def repo_id_from_path(root: str | Path) -> str:
        root = str(Path(root).resolve())
        h    = hashlib.sha256(root.encode()).hexdigest()[:8]
        name = Path(root).name
        return f"local/{name}-{h}"
