"""NetworkX-based in-memory graph cache for fast traversal queries."""

from __future__ import annotations

import json
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import networkx as nx


@dataclass(slots=True)
class GraphNode:
    """Lightweight node record returned by queries."""
    name: str
    node_type: str  # "code_entity", "doc_entity", "service", etc.
    subsystem: str  # "code", "docs", "otel"
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class GraphEdge:
    """Lightweight edge record returned by queries."""
    source: str
    target: str
    edge_type: str  # "calls", "imports", "mentions", "monitors", etc.
    weight: float = 1.0
    metadata: dict[str, Any] = field(default_factory=dict)


class GraphStore:
    """Thread-safe in-memory graph with persistence support.

    Wraps NetworkX DiGraph with locking, snapshot persistence, and
    common query patterns (neighborhood, shortest path, PageRank, communities).
    """

    def __init__(self, persist_path: Path | None = None) -> None:
        self._graph = nx.DiGraph()
        self._lock = threading.RLock()
        self._persist_path = persist_path

    # -- Persistence ----------------------------------------------------------

    def save_snapshot(self, path: Path | None = None) -> None:
        """Persist graph to JSON (node-link format)."""
        target = path or self._persist_path
        if not target:
            return
        with self._lock:
            data = nx.node_link_data(self._graph)
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(json.dumps(data, default=str), encoding="utf-8")

    def load_snapshot(self, path: Path | None = None) -> bool:
        """Load graph from JSON snapshot. Returns True on success."""
        target = path or self._persist_path
        if not target or not target.exists():
            return False
        try:
            raw = json.loads(target.read_text(encoding="utf-8"))
            graph = nx.node_link_graph(raw, directed=True)
            with self._lock:
                self._graph = graph
            return True
        except (json.JSONDecodeError, KeyError, nx.NetworkXError):
            return False

    # -- Build / Mutate -------------------------------------------------------

    def clear(self) -> None:
        with self._lock:
            self._graph.clear()

    def add_node(self, name: str, *, node_type: str = "", subsystem: str = "", **attrs: Any) -> None:
        with self._lock:
            self._graph.add_node(name, node_type=node_type, subsystem=subsystem, **attrs)

    def add_edge(self, source: str, target: str, *, edge_type: str = "", weight: float = 1.0, **attrs: Any) -> None:
        with self._lock:
            self._graph.add_edge(source, target, edge_type=edge_type, weight=weight, **attrs)

    def remove_node(self, name: str) -> None:
        with self._lock:
            if name in self._graph:
                self._graph.remove_node(name)

    def build_from_code_repo(self, code_repo: Any) -> int:
        """Populate graph from a CodeRepository's entities and dependencies."""
        from codectx.db.models import Entity, Dependency, EntityRelation
        from common.db import session_scope

        count = 0
        with self._lock:
            with session_scope(code_repo.session_factory) as session:
                for e in session.query(Entity).all():
                    self._graph.add_node(
                        e.qualified_name or e.name,
                        node_type=e.type,
                        subsystem="code",
                        file_id=e.file_id,
                        kind=e.kind or "",
                    )
                    count += 1
                for d in session.query(Dependency).all():
                    self._graph.add_edge(
                        f"file:{d.source_file_id}",
                        d.target_path,
                        edge_type="imports",
                        dep_type=d.dep_type or "",
                    )
                for r in session.query(EntityRelation).all():
                    source = session.get(Entity, r.source_id)
                    target = session.get(Entity, r.target_id)
                    if source and target:
                        self._graph.add_edge(
                            source.qualified_name or source.name,
                            target.qualified_name or target.name,
                            edge_type=r.relation_type or "relates",
                        )
        return count

    def build_from_docs_repo(self, docs_repo: Any) -> int:
        """Populate graph from a DocsRepository's entities and relationships."""
        from docsctx.db.models import DocumentEntity, DocumentRelationship
        from common.db import session_scope

        count = 0
        with self._lock:
            with session_scope(docs_repo.session_factory) as session:
                for e in session.query(DocumentEntity).all():
                    self._graph.add_node(
                        e.name,
                        node_type=e.entity_type or "doc_entity",
                        subsystem="docs",
                        entity_id=e.id,
                    )
                    count += 1
                for r in session.query(DocumentRelationship).all():
                    self._graph.add_edge(
                        r.source_name,
                        r.target_name,
                        edge_type=r.relation_type or "relates",
                        weight=r.weight if hasattr(r, 'weight') else 1.0,
                    )
        return count

    # -- Query ----------------------------------------------------------------

    @property
    def node_count(self) -> int:
        with self._lock:
            return self._graph.number_of_nodes()

    @property
    def edge_count(self) -> int:
        with self._lock:
            return self._graph.number_of_edges()

    def has_node(self, name: str) -> bool:
        with self._lock:
            return name in self._graph

    def get_node(self, name: str) -> dict[str, Any] | None:
        with self._lock:
            if name in self._graph:
                return dict(self._graph.nodes[name])
            return None

    def neighborhood(self, name: str, depth: int = 1, max_nodes: int = 100) -> dict[str, Any]:
        """Return the ego-graph neighbourhood of a node."""
        with self._lock:
            if name not in self._graph:
                return {"center": name, "nodes": [], "edges": []}

            # BFS to collect nodes within depth
            visited: set[str] = set()
            frontier = {name}
            for _ in range(depth):
                next_frontier: set[str] = set()
                for n in frontier:
                    visited.add(n)
                    if len(visited) >= max_nodes:
                        break
                    next_frontier.update(self._graph.successors(n))
                    next_frontier.update(self._graph.predecessors(n))
                frontier = next_frontier - visited
                if len(visited) >= max_nodes:
                    break
            visited.update(frontier)
            nodes_list = list(visited)[:max_nodes]

            sub = self._graph.subgraph(nodes_list)
            nodes = [
                {"name": n, **dict(sub.nodes[n])}
                for n in sub.nodes
            ]
            edges = [
                {"source": u, "target": v, **dict(d)}
                for u, v, d in sub.edges(data=True)
            ]
            return {"center": name, "nodes": nodes, "edges": edges}

    def shortest_path(self, source: str, target: str) -> list[str] | None:
        """Return shortest path between two nodes, or None if unreachable."""
        with self._lock:
            try:
                return nx.shortest_path(self._graph, source, target)
            except (nx.NodeNotFound, nx.NetworkXNoPath):
                return None

    def page_rank(self, top_k: int = 20) -> list[tuple[str, float]]:
        """Return top-k nodes by PageRank score."""
        with self._lock:
            if self._graph.number_of_nodes() == 0:
                return []
            pr = nx.pagerank(self._graph, max_iter=100)
        ranked = sorted(pr.items(), key=lambda x: x[1], reverse=True)
        return ranked[:top_k]

    def communities(self, algorithm: str = "louvain") -> list[set[str]]:
        """Detect communities. Returns list of node-name sets."""
        with self._lock:
            undirected = self._graph.to_undirected()
            if undirected.number_of_nodes() == 0:
                return []
            if algorithm == "louvain":
                try:
                    from networkx.algorithms.community import louvain_communities
                    return [set(c) for c in louvain_communities(undirected)]
                except ImportError:
                    pass
            # Fallback: greedy modularity
            from networkx.algorithms.community import greedy_modularity_communities
            return [set(c) for c in greedy_modularity_communities(undirected)]

    def subgraph(self, entity_names: list[str]) -> dict[str, Any]:
        """Return subgraph induced by the given node names."""
        with self._lock:
            present = [n for n in entity_names if n in self._graph]
            sub = self._graph.subgraph(present)
            nodes = [{"name": n, **dict(sub.nodes[n])} for n in sub.nodes]
            edges = [{"source": u, "target": v, **dict(d)} for u, v, d in sub.edges(data=True)]
            return {"nodes": nodes, "edges": edges}

    def cross_reference(self, code_entity: str, doc_entity: str) -> dict[str, Any]:
        """Find paths and shared neighbours between a code and doc entity."""
        result: dict[str, Any] = {
            "code_entity": code_entity,
            "doc_entity": doc_entity,
            "path": None,
            "shared_neighbors": [],
        }
        with self._lock:
            code_exists = code_entity in self._graph
            doc_exists = doc_entity in self._graph
            if not code_exists or not doc_exists:
                return result
            try:
                result["path"] = nx.shortest_path(self._graph, code_entity, doc_entity)
            except (nx.NetworkXNoPath, nx.NodeNotFound):
                try:
                    result["path"] = nx.shortest_path(self._graph, doc_entity, code_entity)
                except (nx.NetworkXNoPath, nx.NodeNotFound):
                    pass
            code_nb = set(self._graph.successors(code_entity)) | set(self._graph.predecessors(code_entity))
            doc_nb = set(self._graph.successors(doc_entity)) | set(self._graph.predecessors(doc_entity))
            result["shared_neighbors"] = list(code_nb & doc_nb)
        return result

    def to_dict(self) -> dict[str, Any]:
        """Serialize full graph for API responses."""
        with self._lock:
            return {
                "node_count": self._graph.number_of_nodes(),
                "edge_count": self._graph.number_of_edges(),
                "nodes": [{"name": n, **dict(self._graph.nodes[n])} for n in self._graph.nodes],
                "edges": [
                    {"source": u, "target": v, **dict(d)}
                    for u, v, d in self._graph.edges(data=True)
                ],
            }
