"""Storage backend — SQLite persistence for nodes and edges with numpy-based vector search."""

from __future__ import annotations

import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

from .edge import MemoryEdge
from .node import MemoryNode


def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two vectors."""
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a, b) / (norm_a * norm_b))


class MemoryStore:
    """SQLite-backed store for MemoryNodes and MemoryEdges.

    Embeddings are stored as JSON-encoded float arrays.  Vector similarity
    search is done in-process via numpy (suitable for v0.1 scale).
    """

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(str(self.path), check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._create_tables()

    # ------------------------------------------------------------------
    # Schema
    # ------------------------------------------------------------------

    def _create_tables(self) -> None:
        cur = self._conn.cursor()
        cur.executescript(
            """
            CREATE TABLE IF NOT EXISTS nodes (
                id            TEXT PRIMARY KEY,
                content       TEXT NOT NULL,
                embedding     TEXT NOT NULL,
                importance    REAL NOT NULL DEFAULT 0.5,
                timestamp     TEXT NOT NULL,
                entity_tags   TEXT NOT NULL DEFAULT '[]',
                context_scope TEXT NOT NULL DEFAULT 'global',
                summary_level INTEGER NOT NULL DEFAULT 0,
                access_count  INTEGER NOT NULL DEFAULT 0,
                last_accessed TEXT,
                memory_type   TEXT NOT NULL DEFAULT 'memory',
                deleted_at    TEXT
            );

            CREATE TABLE IF NOT EXISTS edges (
                id                TEXT PRIMARY KEY,
                source_id         TEXT NOT NULL,
                target_id         TEXT NOT NULL,
                relationship_type TEXT NOT NULL,
                strength          REAL NOT NULL DEFAULT 1.0,
                timestamp         TEXT NOT NULL,
                FOREIGN KEY (source_id) REFERENCES nodes(id),
                FOREIGN KEY (target_id) REFERENCES nodes(id)
            );

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

            CREATE TABLE IF NOT EXISTS traces (
                id            TEXT PRIMARY KEY,
                timestamp     TEXT NOT NULL,
                query         TEXT NOT NULL,
                node_ids_used TEXT NOT NULL DEFAULT '[]',
                response      TEXT NOT NULL DEFAULT '',
                scope         TEXT NOT NULL DEFAULT 'global',
                metadata      TEXT NOT NULL DEFAULT '{}'
            );

            CREATE INDEX IF NOT EXISTS idx_nodes_scope ON nodes(context_scope);
            CREATE INDEX IF NOT EXISTS idx_nodes_importance ON nodes(importance);
            CREATE INDEX IF NOT EXISTS idx_nodes_timestamp ON nodes(timestamp);
            CREATE INDEX IF NOT EXISTS idx_nodes_type ON nodes(memory_type);
            CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source_id);
            CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target_id);
            CREATE INDEX IF NOT EXISTS idx_edges_type ON edges(relationship_type);
            CREATE INDEX IF NOT EXISTS idx_traces_timestamp ON traces(timestamp);
            CREATE INDEX IF NOT EXISTS idx_traces_scope ON traces(scope);
            """
        )
        self._conn.commit()
        self._migrate()

    # ------------------------------------------------------------------
    # Migration
    # ------------------------------------------------------------------

    def _migrate(self) -> None:
        """Add columns that may be missing from older databases."""
        cur = self._conn.cursor()
        cur.execute("PRAGMA table_info(nodes)")
        cols = {row[1] for row in cur.fetchall()}
        if "memory_type" not in cols:
            cur.execute("ALTER TABLE nodes ADD COLUMN memory_type TEXT NOT NULL DEFAULT 'memory'")
        if "deleted_at" not in cols:
            cur.execute("ALTER TABLE nodes ADD COLUMN deleted_at TEXT")
        self._conn.commit()

    # ------------------------------------------------------------------
    # Node CRUD
    # ------------------------------------------------------------------

    def save_node(self, node: MemoryNode) -> None:
        """Insert or replace a MemoryNode."""
        cur = self._conn.cursor()
        cur.execute(
            """
            INSERT OR REPLACE INTO nodes
                (id, content, embedding, importance, timestamp,
                 entity_tags, context_scope, summary_level,
                 access_count, last_accessed, memory_type, deleted_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                node.id,
                node.content,
                json.dumps(node.embedding),
                node.importance,
                node.timestamp.isoformat(),
                json.dumps(node.entity_tags),
                node.context_scope,
                node.summary_level,
                node.access_count,
                node.last_accessed.isoformat() if node.last_accessed else None,
                node.memory_type,
                node.deleted_at.isoformat() if node.deleted_at else None,
            ),
        )
        self._conn.commit()

    def get_node(self, node_id: str) -> Optional[MemoryNode]:
        """Retrieve a single node by ID."""
        cur = self._conn.cursor()
        cur.execute("SELECT * FROM nodes WHERE id = ?", (node_id,))
        row = cur.fetchone()
        if row is None:
            return None
        return self._row_to_node(row)

    def delete_node(self, node_id: str) -> None:
        """Delete a node and all edges that reference it."""
        cur = self._conn.cursor()
        cur.execute("DELETE FROM edges WHERE source_id = ? OR target_id = ?", (node_id, node_id))
        cur.execute("DELETE FROM nodes WHERE id = ?", (node_id,))
        self._conn.commit()

    def update_node(self, node: MemoryNode) -> None:
        """Alias for save_node (upsert semantics)."""
        self.save_node(node)

    def list_nodes(
        self,
        scope: Optional[str] = None,
        min_importance: float = 0.0,
        summary_level: Optional[int] = None,
        limit: int = 100,
        include_deleted: bool = False,
        memory_type: Optional[str] = None,
    ) -> List[MemoryNode]:
        """List nodes with optional filters."""
        query = "SELECT * FROM nodes WHERE importance >= ?"
        params: list = [min_importance]
        if not include_deleted:
            query += " AND deleted_at IS NULL"
        if scope is not None:
            query += " AND context_scope = ?"
            params.append(scope)
        if summary_level is not None:
            query += " AND summary_level = ?"
            params.append(summary_level)
        if memory_type is not None:
            query += " AND memory_type = ?"
            params.append(memory_type)
        query += " ORDER BY timestamp DESC LIMIT ?"
        params.append(limit)

        cur = self._conn.cursor()
        cur.execute(query, params)
        return [self._row_to_node(r) for r in cur.fetchall()]

    def count_nodes(self) -> int:
        cur = self._conn.cursor()
        cur.execute("SELECT COUNT(*) FROM nodes")
        return cur.fetchone()[0]

    # ------------------------------------------------------------------
    # Edge CRUD
    # ------------------------------------------------------------------

    def save_edge(self, edge: MemoryEdge) -> None:
        cur = self._conn.cursor()
        cur.execute(
            """
            INSERT OR REPLACE INTO edges
                (id, source_id, target_id, relationship_type, strength, timestamp)
            VALUES (?, ?, ?, ?, ?, ?)
            """,
            (
                edge.id,
                edge.source_id,
                edge.target_id,
                edge.relationship_type,
                edge.strength,
                edge.timestamp.isoformat(),
            ),
        )
        self._conn.commit()

    def get_edges(
        self,
        node_id: str,
        relationship_types: Optional[List[str]] = None,
        direction: str = "outgoing",
    ) -> List[MemoryEdge]:
        """Get edges connected to a node."""
        if direction == "outgoing":
            query = "SELECT * FROM edges WHERE source_id = ?"
        elif direction == "incoming":
            query = "SELECT * FROM edges WHERE target_id = ?"
        else:
            query = "SELECT * FROM edges WHERE source_id = ? OR target_id = ?"

        params: list = [node_id] if direction in ("outgoing", "incoming") else [node_id, node_id]

        if relationship_types:
            placeholders = ",".join("?" for _ in relationship_types)
            query += f" AND relationship_type IN ({placeholders})"
            params.extend(relationship_types)

        cur = self._conn.cursor()
        cur.execute(query, params)
        return [self._row_to_edge(r) for r in cur.fetchall()]

    def delete_edge(self, edge_id: str) -> None:
        cur = self._conn.cursor()
        cur.execute("DELETE FROM edges WHERE id = ?", (edge_id,))
        self._conn.commit()

    def count_edges(self) -> int:
        cur = self._conn.cursor()
        cur.execute("SELECT COUNT(*) FROM edges")
        return cur.fetchone()[0]

    # ------------------------------------------------------------------
    # Vector search
    # ------------------------------------------------------------------

    def vector_search(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        scope: Optional[str] = None,
        min_importance: float = 0.0,
        entity_tags: Optional[List[str]] = None,
        after: Optional[datetime] = None,
        before: Optional[datetime] = None,
        include_deleted: bool = False,
    ) -> List[Tuple[MemoryNode, float]]:
        """Brute-force cosine-similarity search over stored embeddings.

        Returns list of (node, similarity_score) tuples, highest first.
        """
        query = "SELECT * FROM nodes WHERE importance >= ?"
        params: list = [min_importance]

        if not include_deleted:
            query += " AND deleted_at IS NULL"
        if scope is not None:
            query += " AND context_scope = ?"
            params.append(scope)
        if after is not None:
            query += " AND timestamp >= ?"
            params.append(after.isoformat())
        if before is not None:
            query += " AND timestamp <= ?"
            params.append(before.isoformat())

        cur = self._conn.cursor()
        cur.execute(query, params)
        rows = cur.fetchall()

        q_vec = np.asarray(query_embedding, dtype=np.float32)
        scored: List[Tuple[MemoryNode, float]] = []

        entity_set = set(entity_tags) if entity_tags else None

        for row in rows:
            node = self._row_to_node(row)
            if not node.embedding:
                continue
            # Optional entity filter
            if entity_set and not entity_set.intersection(node.entity_tags):
                continue
            n_vec = np.asarray(node.embedding, dtype=np.float32)
            sim = _cosine_similarity(q_vec, n_vec)
            scored.append((node, sim))

        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

    # ------------------------------------------------------------------
    # Graph traversal
    # ------------------------------------------------------------------

    def traverse(
        self,
        start_entity: str,
        relationship_types: Optional[List[str]] = None,
        max_depth: int = 3,
    ) -> List[Dict]:
        """BFS traversal from nodes matching an entity tag.

        Returns a list of dicts: {node, depth, edge, path}.
        """
        # Find seed nodes that contain the start entity
        cur = self._conn.cursor()
        cur.execute("SELECT * FROM nodes")
        all_rows = cur.fetchall()

        seed_ids: List[str] = []
        node_cache: Dict[str, MemoryNode] = {}
        for row in all_rows:
            node = self._row_to_node(row)
            node_cache[node.id] = node
            if start_entity in node.entity_tags:
                seed_ids.append(node.id)

        visited: set = set()
        results: List[Dict] = []
        queue: List[Tuple[str, int, List[str]]] = [(sid, 0, [sid]) for sid in seed_ids]

        while queue:
            current_id, depth, path = queue.pop(0)
            if current_id in visited:
                continue
            visited.add(current_id)

            node = node_cache.get(current_id) or self.get_node(current_id)
            if node is None:
                continue
            node_cache[node.id] = node

            edges = self.get_edges(current_id, relationship_types, direction="outgoing")
            results.append({"node": node, "depth": depth, "path": path})

            if depth < max_depth:
                for edge in edges:
                    if edge.target_id not in visited:
                        queue.append(
                            (edge.target_id, depth + 1, path + [edge.target_id])
                        )

        return results

    # ------------------------------------------------------------------
    # Stats
    # ------------------------------------------------------------------

    def stats(self) -> Dict:
        """Return summary statistics about the store."""
        cur = self._conn.cursor()
        cur.execute("SELECT COUNT(*) FROM nodes")
        node_count = cur.fetchone()[0]
        cur.execute("SELECT COUNT(*) FROM edges")
        edge_count = cur.fetchone()[0]
        cur.execute("SELECT DISTINCT context_scope FROM nodes")
        scopes = [r[0] for r in cur.fetchall()]
        cur.execute("SELECT AVG(importance) FROM nodes")
        avg_imp = cur.fetchone()[0] or 0.0
        return {
            "node_count": node_count,
            "edge_count": edge_count,
            "scopes": scopes,
            "avg_importance": round(avg_imp, 4),
        }

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    @staticmethod
    def _row_to_node(row: sqlite3.Row) -> MemoryNode:
        last_acc = row["last_accessed"]
        deleted = row["deleted_at"] if "deleted_at" in row.keys() else None
        mtype = row["memory_type"] if "memory_type" in row.keys() else "memory"
        return MemoryNode(
            id=row["id"],
            content=row["content"],
            embedding=json.loads(row["embedding"]),
            importance=row["importance"],
            timestamp=datetime.fromisoformat(row["timestamp"]),
            entity_tags=json.loads(row["entity_tags"]),
            context_scope=row["context_scope"],
            summary_level=row["summary_level"],
            access_count=row["access_count"],
            last_accessed=datetime.fromisoformat(last_acc) if last_acc else None,
            memory_type=mtype,
            deleted_at=datetime.fromisoformat(deleted) if deleted else None,
        )

    @staticmethod
    def _row_to_edge(row: sqlite3.Row) -> MemoryEdge:
        return MemoryEdge(
            id=row["id"],
            source_id=row["source_id"],
            target_id=row["target_id"],
            relationship_type=row["relationship_type"],
            strength=row["strength"],
            timestamp=datetime.fromisoformat(row["timestamp"]),
        )

    # ------------------------------------------------------------------
    # Persona CRUD
    # ------------------------------------------------------------------

    def set_persona(self, key: str, value: str) -> None:
        cur = self._conn.cursor()
        cur.execute(
            "INSERT OR REPLACE INTO personas (key, value) VALUES (?, ?)",
            (key, value),
        )
        self._conn.commit()

    def get_persona(self, key: str) -> Optional[str]:
        cur = self._conn.cursor()
        cur.execute("SELECT value FROM personas WHERE key = ?", (key,))
        row = cur.fetchone()
        return row[0] if row else None

    def get_all_persona(self) -> Dict[str, str]:
        cur = self._conn.cursor()
        cur.execute("SELECT key, value FROM personas")
        return {r[0]: r[1] for r in cur.fetchall()}

    def delete_persona(self, key: Optional[str] = None) -> None:
        cur = self._conn.cursor()
        if key:
            cur.execute("DELETE FROM personas WHERE key = ?", (key,))
        else:
            cur.execute("DELETE FROM personas")
        self._conn.commit()

    # ------------------------------------------------------------------
    # Trace CRUD
    # ------------------------------------------------------------------

    def save_trace(
        self,
        trace_id: str,
        timestamp: datetime,
        query: str,
        node_ids_used: List[str],
        response: str = "",
        scope: str = "global",
        metadata: Optional[Dict] = None,
    ) -> None:
        cur = self._conn.cursor()
        cur.execute(
            """
            INSERT OR REPLACE INTO traces
                (id, timestamp, query, node_ids_used, response, scope, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                trace_id,
                timestamp.isoformat(),
                query,
                json.dumps(node_ids_used),
                response,
                scope,
                json.dumps(metadata or {}),
            ),
        )
        self._conn.commit()

    def list_traces(
        self,
        scope: Optional[str] = None,
        limit: int = 50,
    ) -> List[Dict]:
        query = "SELECT * FROM traces"
        params: list = []
        if scope is not None:
            query += " WHERE scope = ?"
            params.append(scope)
        query += " ORDER BY timestamp DESC LIMIT ?"
        params.append(limit)
        cur = self._conn.cursor()
        cur.execute(query, params)
        return [self._row_to_trace(r) for r in cur.fetchall()]

    def get_trace(self, trace_id: str) -> Optional[Dict]:
        cur = self._conn.cursor()
        cur.execute("SELECT * FROM traces WHERE id = ?", (trace_id,))
        row = cur.fetchone()
        return self._row_to_trace(row) if row else None

    def traces_using_node(self, node_id: str) -> List[Dict]:
        cur = self._conn.cursor()
        cur.execute("SELECT * FROM traces ORDER BY timestamp DESC")
        results = []
        for row in cur.fetchall():
            ids = json.loads(row["node_ids_used"])
            if node_id in ids:
                results.append(self._row_to_trace(row))
        return results

    @staticmethod
    def _row_to_trace(row: sqlite3.Row) -> Dict:
        return {
            "id": row["id"],
            "timestamp": datetime.fromisoformat(row["timestamp"]),
            "query": row["query"],
            "node_ids_used": json.loads(row["node_ids_used"]),
            "response": row["response"],
            "scope": row["scope"],
            "metadata": json.loads(row["metadata"]),
        }

    # ------------------------------------------------------------------
    # Soft-delete
    # ------------------------------------------------------------------

    def soft_delete_node(self, node_id: str) -> None:
        """Mark a node as deleted without removing it."""
        cur = self._conn.cursor()
        cur.execute(
            "UPDATE nodes SET deleted_at = ? WHERE id = ?",
            (datetime.now(timezone.utc).isoformat(), node_id),
        )
        self._conn.commit()

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