"""
sqlite_store.py
----------------
SQLite-backed persistence layer for HIEROMEM.
Implements the MemoryStore interface for local, lightweight persistence.

Stores:
- Embedding vectors (as JSON)
- Scene graphs (as JSON of nodes and edges)
- Metadata (JSON)
- Trust values, timestamps, parent links, etc.

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional

import networkx as nx
import numpy as np

from hieromem_core.episode import Episode
from .memory_store import MemoryStore


class SQLiteStore(MemoryStore):
    """SQLite-backed persistence store for episodes."""

    def __init__(
        self,
        db_path: str = "memory.db",
        *,
        default_namespace: str = "default",
        debug: bool = False,
    ):
        super().__init__(debug=debug)
        self.db_path = Path(db_path)
        self.default_namespace = default_namespace or "default"
        self._init_db()
        self._log(f"SQLiteStore initialized at {self.db_path.resolve()}")

    # ------------------------------------------------------------------ #
    # Database initialization
    # ------------------------------------------------------------------ #
    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path)
        conn.execute("PRAGMA synchronous=NORMAL;")
        return conn

    def _init_db(self) -> None:
        with self._connect() as conn:
            # Enable WAL mode for better concurrency
            conn.execute("PRAGMA journal_mode=WAL;")
            conn.execute("PRAGMA synchronous=NORMAL;")

            # Create table if it doesn't exist
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS episodes (
                    id TEXT PRIMARY KEY,
                    e TEXT NOT NULL,
                    G TEXT NOT NULL,
                    tocc TEXT NOT NULL,
                    meta TEXT,
                    parent TEXT,
                    trust REAL NOT NULL
                );
                """
            )

            # Check if namespace column exists, add if missing (migration)
            cursor = conn.execute("PRAGMA table_info(episodes)")
            columns = [row[1] for row in cursor.fetchall()]

            if 'namespace' not in columns:
                conn.execute("ALTER TABLE episodes ADD COLUMN namespace TEXT DEFAULT 'default'")

            # Create indices
            conn.executescript(
                """
                CREATE INDEX IF NOT EXISTS idx_episodes_tocc ON episodes(tocc);
                CREATE INDEX IF NOT EXISTS idx_episodes_trust ON episodes(trust);
                CREATE INDEX IF NOT EXISTS idx_episodes_namespace ON episodes(namespace);
                CREATE INDEX IF NOT EXISTS idx_episodes_namespace_tocc ON episodes(namespace, tocc);
                CREATE INDEX IF NOT EXISTS idx_episodes_namespace_trust ON episodes(namespace, trust);
                CREATE INDEX IF NOT EXISTS idx_episodes_parent ON episodes(parent);
                """
            )
            conn.commit()

    # ------------------------------------------------------------------ #
    # Helpers
    # ------------------------------------------------------------------ #
    def _encode_episode(
        self, ep: Episode, namespace: Optional[str] = None
    ) -> tuple[str, str, str, str, str, str | None, float, str]:
        """Serialize Episode fields for DB storage."""
        self.validate_episode(ep)
        metadata = ep.meta if isinstance(ep.meta, dict) else {}
        ns = namespace or metadata.get("namespace") or self.default_namespace
        G = ep.G
        graph_data = {
            "nodes": [{"id": n, **(G.nodes[n] or {})} for n in G.nodes],
            "edges": [{"u": u, "v": v, **(G.edges[u, v] or {})} for u, v in G.edges],
        }
        return (
            ep.id,
            json.dumps(ep.e.tolist()),
            json.dumps(graph_data),
            ep.tocc.isoformat(),
            json.dumps(ep.meta),
            ep.parent,
            float(ep.trust),
            ns,
        )

    def _decode_episode(self, row: tuple[Any, ...]) -> Episode:
        """Deserialize Episode from DB row."""
        if len(row) == 8:
            eid, e_str, G_str, tocc, meta_str, parent, trust, namespace = row
        else:
            eid, e_str, G_str, tocc, meta_str, parent, trust = row
            namespace = self.default_namespace
        e = np.array(json.loads(e_str), dtype=np.float32)
        G_data = json.loads(G_str)
        G = nx.DiGraph()
        for n in G_data.get("nodes", []):
            nid = n.pop("id")
            G.add_node(nid, **n)
        for edge in G_data.get("edges", []):
            u, v = edge.pop("u"), edge.pop("v")
            G.add_edge(u, v, **edge)
        meta_raw = json.loads(meta_str or "{}")
        if not isinstance(meta_raw, dict):
            meta_raw = {}
        meta = dict(meta_raw)
        if "namespace" not in meta:
            meta["namespace"] = namespace
        return Episode(
            id=eid,
            e=e,
            G=G,
            tocc=datetime.fromisoformat(tocc),
            meta=meta,
            parent=parent,
            trust=trust,
        )

    # ------------------------------------------------------------------ #
    # CRUD operations
    # ------------------------------------------------------------------ #
    def save_episode(self, ep: Episode) -> None:
        self._log(f"Saving episode {ep.id}")
        with self._connect() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO episodes
                (id, e, G, tocc, meta, parent, trust, namespace)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                self._encode_episode(ep),
            )
            conn.commit()

    def save_many(self, episodes: Iterable[Episode]) -> None:
        episode_list = list(episodes)
        with self._connect() as conn:
            conn.executemany(
                """
                INSERT OR REPLACE INTO episodes
                (id, e, G, tocc, meta, parent, trust, namespace)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                [self._encode_episode(ep) for ep in episode_list],
            )
            conn.commit()
        self._log(f"Saved {len(episode_list)} episodes")

    def load_all(self, *, namespace: Optional[str] = None) -> list[Episode]:
        query = "SELECT * FROM episodes"
        params: tuple[Any, ...] = ()
        if namespace:
            query += " WHERE namespace = ?"
            params = (namespace,)
        with self._connect() as conn:
            rows = conn.execute(query, params).fetchall()
        eps = [self._decode_episode(r) for r in rows]
        self._log(f"Loaded {len(eps)} episodes")
        return eps

    def load_by_id(self, episode_id: str, *, namespace: Optional[str] = None) -> Optional[Episode]:
        query = "SELECT * FROM episodes WHERE id = ?"
        params: tuple[Any, ...] = (episode_id,)
        if namespace:
            query += " AND namespace = ?"
            params += (namespace,)
        with self._connect() as conn:
            row = conn.execute(query, params).fetchone()
        return self._decode_episode(row) if row else None

    def update_trust(self, episode_id: str, trust: float, *, namespace: Optional[str] = None) -> None:
        query = "UPDATE episodes SET trust = ? WHERE id = ?"
        params: tuple[Any, ...] = (float(trust), episode_id)
        if namespace:
            query += " AND namespace = ?"
            params += (namespace,)
        with self._connect() as conn:
            conn.execute(query, params)
            conn.commit()
        self._log(f"Updated trust for {episode_id} → {trust:.3f}")

    def remove(self, episode_id: str, *, namespace: Optional[str] = None) -> None:
        query = "DELETE FROM episodes WHERE id = ?"
        params: tuple[Any, ...] = (episode_id,)
        if namespace:
            query += " AND namespace = ?"
            params += (namespace,)
        with self._connect() as conn:
            conn.execute(query, params)
            conn.commit()
        self._log(f"Removed episode {episode_id}")

    def remove_many(self, episode_ids: Iterable[str], *, namespace: Optional[str] = None) -> None:
        ids = list(episode_ids)
        query = "DELETE FROM episodes WHERE id = ?"
        params = [(i,) for i in ids]
        if namespace:
            query += " AND namespace = ?"
            params = [(i, namespace) for i in ids]
        with self._connect() as conn:
            conn.executemany(query, params)
            conn.commit()
        self._log(f"Removed {len(ids)} episodes")

    def clear(self) -> None:
        with self._connect() as conn:
            conn.execute("DELETE FROM episodes")
            conn.commit()
        self._log("Cleared all episodes from store")
