"""ConceptGraph — SQLite-backed concept co-occurrence graph with spreading activation."""

import sqlite3
from collections import deque

import networkx as nx


_GRAPH_SCHEMA = """
CREATE TABLE IF NOT EXISTS concept_nodes (
    concept TEXT PRIMARY KEY,
    mentions INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS concept_edges (
    concept_a TEXT,
    concept_b TEXT,
    weight INTEGER DEFAULT 1,
    PRIMARY KEY (concept_a, concept_b)
);
CREATE TABLE IF NOT EXISTS memory_concepts (
    memory_id TEXT,
    concept TEXT,
    PRIMARY KEY (memory_id, concept)
);
CREATE INDEX IF NOT EXISTS idx_mc_concept ON memory_concepts(concept);
"""


class ConceptGraph:
    """Graph of concepts linked by co-occurrence in memories."""

    def __init__(self, conn: sqlite3.Connection):
        self._conn = conn
        self._conn.executescript(_GRAPH_SCHEMA)
        self._graph = nx.Graph()
        self._memory_concepts: dict[str, list[str]] = {}
        self._load()

    def add_memory(self, memory_id: str, concepts: list[str]) -> None:
        """Add concept nodes + co-occurrence edges for a memory."""
        if not concepts:
            return

        self._memory_concepts[memory_id] = list(concepts)

        for c in concepts:
            if not self._graph.has_node(c):
                self._graph.add_node(c, mentions=0, memories=[])
            self._graph.nodes[c]["mentions"] += 1
            if memory_id not in self._graph.nodes[c]["memories"]:
                self._graph.nodes[c]["memories"].append(memory_id)

        # Co-occurrence edges
        for i in range(len(concepts)):
            for j in range(i + 1, len(concepts)):
                a, b = concepts[i], concepts[j]
                if self._graph.has_edge(a, b):
                    self._graph[a][b]["weight"] += 1
                else:
                    self._graph.add_edge(a, b, weight=1)

        self.save()

    def remove_memory(self, memory_id: str, concepts: list[str]) -> None:
        """Decrement weights for a removed memory."""
        for c in concepts:
            if self._graph.has_node(c):
                self._graph.nodes[c]["mentions"] = max(0, self._graph.nodes[c]["mentions"] - 1)
                mems = self._graph.nodes[c]["memories"]
                if memory_id in mems:
                    mems.remove(memory_id)

        for i in range(len(concepts)):
            for j in range(i + 1, len(concepts)):
                a, b = concepts[i], concepts[j]
                if self._graph.has_edge(a, b):
                    self._graph[a][b]["weight"] -= 1
                    if self._graph[a][b]["weight"] <= 0:
                        self._graph.remove_edge(a, b)

        self._memory_concepts.pop(memory_id, None)
        self.save()

    def spreading_activation(
        self, concepts: list[str], hops: int = 2, decay: float = 0.5
    ) -> dict[str, float]:
        """BFS from seed concepts with decaying activation per hop."""
        if not concepts:
            return {}

        activation: dict[str, float] = {}
        queue: deque[tuple[str, float, int]] = deque()

        for c in concepts:
            if self._graph.has_node(c):
                activation[c] = 1.0
                queue.append((c, 1.0, 0))

        while queue:
            node, energy, hop = queue.popleft()
            if hop >= hops:
                continue
            for neighbor in self._graph.neighbors(node):
                weight = self._graph[node][neighbor].get("weight", 1)
                new_energy = energy * decay * min(weight, 3) / 3
                if neighbor not in activation or new_energy > activation[neighbor]:
                    activation[neighbor] = new_energy
                    queue.append((neighbor, new_energy, hop + 1))

        return activation

    def get_related_memories(self, concepts: list[str], n: int = 5) -> list[str]:
        """Use spreading activation to find related memory IDs."""
        activated = self.spreading_activation(concepts)
        if not activated:
            return []

        # Collect memory IDs weighted by activation
        memory_scores: dict[str, float] = {}
        for concept, score in activated.items():
            if self._graph.has_node(concept):
                for mid in self._graph.nodes[concept]["memories"]:
                    memory_scores[mid] = memory_scores.get(mid, 0) + score

        sorted_memories = sorted(memory_scores, key=memory_scores.get, reverse=True)
        return sorted_memories[:n]

    def has_node(self, concept: str) -> bool:
        return self._graph.has_node(concept)

    def has_edge(self, a: str, b: str) -> bool:
        return self._graph.has_edge(a, b)

    def edge_weight(self, a: str, b: str) -> int:
        if not self._graph.has_edge(a, b):
            return 0
        return self._graph[a][b]["weight"]

    def mentions(self, concept: str) -> int:
        if not self._graph.has_node(concept):
            return 0
        return self._graph.nodes[concept]["mentions"]

    def memories_for(self, concept: str) -> list[str]:
        if not self._graph.has_node(concept):
            return []
        return list(self._graph.nodes[concept]["memories"])

    def clear(self) -> int:
        """Remove all nodes, edges, and memory mappings."""
        n = self._graph.number_of_nodes()
        self._graph.clear()
        self._memory_concepts.clear()
        self._save_to_db()
        return n

    def node_count(self) -> int:
        return self._graph.number_of_nodes()

    def save(self) -> None:
        """Persist graph state to SQLite."""
        self._save_to_db()

    def _save_to_db(self) -> None:
        """Write full graph state to SQLite tables."""
        self._conn.execute("DELETE FROM concept_nodes")
        self._conn.execute("DELETE FROM concept_edges")
        self._conn.execute("DELETE FROM memory_concepts")

        for name, data in self._graph.nodes(data=True):
            self._conn.execute(
                "INSERT INTO concept_nodes (concept, mentions) VALUES (?, ?)",
                (name, data.get("mentions", 0)),
            )

        for u, v, data in self._graph.edges(data=True):
            a, b = sorted([u, v])
            self._conn.execute(
                "INSERT INTO concept_edges (concept_a, concept_b, weight) VALUES (?, ?, ?)",
                (a, b, data.get("weight", 1)),
            )

        for mid, concepts in self._memory_concepts.items():
            for c in concepts:
                self._conn.execute(
                    "INSERT INTO memory_concepts (memory_id, concept) VALUES (?, ?)",
                    (mid, c),
                )

        self._conn.commit()

    def _load(self) -> None:
        """Load graph from SQLite tables."""
        try:
            rows = self._conn.execute("SELECT concept, mentions FROM concept_nodes").fetchall()
            for row in rows:
                concept = row[0] if isinstance(row, tuple) else row["concept"]
                mentions = row[1] if isinstance(row, tuple) else row["mentions"]
                self._graph.add_node(concept, mentions=mentions, memories=[])

            rows = self._conn.execute("SELECT concept_a, concept_b, weight FROM concept_edges").fetchall()
            for row in rows:
                a = row[0] if isinstance(row, tuple) else row["concept_a"]
                b = row[1] if isinstance(row, tuple) else row["concept_b"]
                w = row[2] if isinstance(row, tuple) else row["weight"]
                self._graph.add_edge(a, b, weight=w)

            rows = self._conn.execute("SELECT memory_id, concept FROM memory_concepts").fetchall()
            for row in rows:
                mid = row[0] if isinstance(row, tuple) else row["memory_id"]
                concept = row[1] if isinstance(row, tuple) else row["concept"]
                self._memory_concepts.setdefault(mid, []).append(concept)
                if self._graph.has_node(concept):
                    if mid not in self._graph.nodes[concept]["memories"]:
                        self._graph.nodes[concept]["memories"].append(mid)
        except (sqlite3.OperationalError, KeyError):
            pass
