# SPDX-FileCopyrightText: 2026 Ashlesha T
# SPDX-License-Identifier: MIT
#
# This file is part of CogniRepo — https://github.com/ashlesh-t/cognirepo
# Licensed under MIT. See LICENSE file in repository root.

"""
Knowledge graph for CogniRepo — a directed NetworkX graph tracking relationships
between files, symbols, concepts, queries, and user actions.

Node types  : FILE, FUNCTION, CLASS, CONCEPT, QUERY, SESSION, USER_ACTION
Edge types  : RELATES_TO, DEFINED_IN, CALLED_BY, QUERIED_WITH, CO_OCCURS

Persistence : pickle to .cognirepo/graph/graph.pkl
              Falls back to an empty graph on load failure (version drift etc.)
"""
import os
import pickle
import sys
import warnings
from typing import Any

import networkx as nx


from config.paths import get_path
from config.lock import store_lock

# Python builtins + common dunder names that must never dominate the concept
# space.  Exported so mcp_server.py can reuse the same set.
PYTHON_BUILTINS: frozenset[str] = frozenset({
    # built-in functions
    "len", "str", "int", "float", "bool", "list", "dict", "set", "tuple",
    "bytes", "bytearray", "memoryview", "complex", "type", "object",
    "range", "enumerate", "zip", "map", "filter", "reversed", "sorted",
    "sum", "min", "max", "abs", "round", "pow", "divmod",
    "open", "print", "input", "repr", "hash", "id", "iter", "next",
    "any", "all", "isinstance", "issubclass", "hasattr", "getattr",
    "setattr", "delattr", "callable", "vars", "dir", "locals", "globals",
    "super", "property", "classmethod", "staticmethod",
    "format", "chr", "ord", "hex", "oct", "bin", "eval", "exec",
    # common list/dict/str methods (added as call nodes by the indexer)
    "append", "extend", "insert", "remove", "pop", "clear", "copy",
    "update", "keys", "values", "items", "get", "setdefault",
    "join", "split", "rsplit", "splitlines", "strip", "lstrip", "rstrip",
    "replace", "startswith", "endswith", "find", "rfind", "index",
    "encode", "decode", "upper", "lower", "capitalize", "title",
    "read", "readline", "readlines", "write", "writelines", "close",
    "seek", "tell", "flush", "fileno",
    # exceptions
    "Exception", "BaseException", "ValueError", "TypeError", "KeyError",
    "IndexError", "AttributeError", "RuntimeError", "StopIteration",
    "GeneratorExit", "OSError", "IOError", "FileNotFoundError",
    "NotImplementedError", "ImportError", "ModuleNotFoundError",
    "OverflowError", "ZeroDivisionError", "MemoryError", "RecursionError",
    "NameError", "UnboundLocalError", "PermissionError", "TimeoutError",
    # dunders
    "__init__", "__new__", "__del__", "__str__", "__repr__", "__bytes__",
    "__len__", "__eq__", "__ne__", "__lt__", "__le__", "__gt__", "__ge__",
    "__hash__", "__bool__", "__iter__", "__next__", "__reversed__",
    "__contains__", "__getitem__", "__setitem__", "__delitem__",
    "__enter__", "__exit__", "__call__", "__get__", "__set__",
    "__add__", "__radd__", "__iadd__", "__sub__", "__mul__", "__truediv__",
    "__class__", "__dict__", "__doc__", "__module__", "__slots__",
    "__all__", "__name__", "__file__", "__spec__", "__path__",
})


def _graph_file() -> str:
    return get_path("graph/graph.pkl")


class NodeType:  # pylint: disable=too-few-public-methods
    """Node types for the knowledge graph."""
    FILE = "FILE"
    FUNCTION = "FUNCTION"
    CLASS = "CLASS"
    CONCEPT = "CONCEPT"
    QUERY = "QUERY"
    SESSION = "SESSION"
    USER_ACTION = "USER_ACTION"
    MEMORY = "MEMORY"          # cross-agent memory nodes (synced from Claude/Gemini/etc.)
    ERROR = "ERROR"            # error pattern nodes for tracking recurring mistakes
    ENDPOINT = "ENDPOINT"      # exposed HTTP/gRPC endpoint (method + path pattern)


class EdgeType:  # pylint: disable=too-few-public-methods
    """Relationship types between nodes."""
    RELATES_TO = "RELATES_TO"
    DEFINED_IN = "DEFINED_IN"
    CALLED_BY = "CALLED_BY"   # caller → callee (forward call direction)
    CALLS = "CALLS"            # callee → caller (reverse; enables BFS to find callers without predecessors())
    QUERIED_WITH = "QUERIED_WITH"
    CO_OCCURS = "CO_OCCURS"
    IMPORTS = "IMPORTS"        # file A imports module/file B
    INHERITS = "INHERITS"      # class A inherits from class B
    EXPOSES = "EXPOSES"        # function → ENDPOINT node (this function handles this route)
    CALLS_ENDPOINT = "CALLS_ENDPOINT"  # caller function → remote ENDPOINT stub (cross-service)


class KnowledgeGraph:
    """Thin wrapper around a networkx DiGraph with CogniRepo-specific conventions."""

    def __init__(self) -> None:
        self.G: nx.DiGraph = nx.DiGraph()  # pylint: disable=invalid-name
        self._load()

    # ── persistence ───────────────────────────────────────────────────────────

    def _load(self) -> None:
        """Load graph from disk; decrypt if needed."""
        if not os.path.exists(_graph_file()):
            return
        try:
            with open(_graph_file(), "rb") as f:
                raw = f.read()
            from security import get_storage_config  # pylint: disable=import-outside-toplevel
            encrypt, project_id = get_storage_config()
            if encrypt:
                from security.encryption import get_or_create_key, decrypt_bytes  # pylint: disable=import-outside-toplevel
                try:
                    raw = decrypt_bytes(raw, get_or_create_key(project_id))
                except Exception:  # pylint: disable=broad-except
                    # The file may have been written unencrypted (e.g. by a
                    # process that resolved the wrong config context before
                    # the _CTX_DIR fix). Fall through and try plaintext —
                    # it will be encrypted on the next save(). A genuinely
                    # encrypted-with-wrong-key file still fails the pickle
                    # load below and is handled by the outer except.
                    pass
            self.G = pickle.loads(raw)  # nosec B301
        except Exception as exc:  # pylint: disable=broad-except
            warnings.warn(
                f"KnowledgeGraph: could not load {_graph_file()} ({exc}). "
                "Starting with an empty graph. Re-run `cognirepo index-repo` to rebuild.",
                stacklevel=2,
            )
            self.G = nx.DiGraph()

    def load(self) -> None:
        """Public alias for reloading the graph from disk."""
        self._load()

    def save(self) -> None:
        """Serialize the graph to a pickle file; encrypt if needed.
        Acquires a cross-process file lock to prevent concurrent writes
        from multiple MCP server processes (e.g. Claude + Gemini) from
        corrupting the pickle.
        """
        from memory.circuit_breaker import get_breaker  # pylint: disable=import-outside-toplevel
        breaker = get_breaker()
        breaker.check()
        os.makedirs(os.path.dirname(_graph_file()), exist_ok=True)
        from security import get_storage_config  # pylint: disable=import-outside-toplevel
        encrypt, project_id = get_storage_config()
        raw = pickle.dumps(self.G, protocol=pickle.HIGHEST_PROTOCOL)
        if encrypt:
            from security.encryption import get_or_create_key, encrypt_bytes  # pylint: disable=import-outside-toplevel
            raw = encrypt_bytes(raw, get_or_create_key(project_id))
        with store_lock():
            with open(_graph_file(), "wb") as f:
                f.write(raw)
        breaker.record_success()

    # ── mutation ──────────────────────────────────────────────────────────────

    def add_node(self, node_id: str, node_type: str, **attrs: Any) -> None:
        """Idempotent add — merges attrs if node already exists."""
        if self.G.has_node(node_id):
            self.G.nodes[node_id].update(attrs)
        else:
            self.G.add_node(node_id, type=node_type, **attrs)

    def add_edge(
        self,
        src: str,
        dst: str,
        edge_type: str,
        weight: float = 1.0,
        **attrs: Any,
    ) -> None:
        """Add a directed edge; if it exists, update its weight."""
        if self.G.has_edge(src, dst):
            self.G[src][dst]["weight"] = weight
            self.G[src][dst].update(attrs)
        else:
            self.G.add_edge(src, dst, rel=edge_type, weight=weight, **attrs)

    def remove_node_edges(self, node_id: str) -> None:
        """Remove all edges incident to node_id (but keep the node)."""
        edges = list(self.G.in_edges(node_id)) + list(self.G.out_edges(node_id))
        self.G.remove_edges_from(edges)

    def nodes_for_file(self, file_path: str) -> list[str]:
        """Return all node IDs whose stored 'file' attr matches file_path."""
        return [
            n for n, d in self.G.nodes(data=True) if d.get("file") == file_path
        ]

    def remove_file_nodes(self, file_path: str) -> list[str]:
        """Remove all nodes (and their incident edges) associated with file_path.

        Removes:
        - Symbol/function/class nodes whose 'file' attribute == file_path
        - The FILE node whose node_id == file_path (convention from make_node_id)

        NetworkX automatically removes all incident edges when a node is removed.
        Returns the list of removed node IDs.
        """
        removed: list[str] = []
        for nid in self.nodes_for_file(file_path):
            if self.G.has_node(nid):
                self.G.remove_node(nid)
                removed.append(nid)
        # FILE node's node_id == rel_path (see make_node_id("FILE", name) → name)
        if self.G.has_node(file_path):
            self.G.remove_node(file_path)
            removed.append(file_path)
        return removed

    # ── queries ───────────────────────────────────────────────────────────────

    def node_exists(self, node_id: str) -> bool:
        """Return True if node_id is present in the graph."""
        return self.G.has_node(node_id)

    def get_neighbours(
        self,
        node_id: str,
        depth: int = 1,
        edge_type: str | None = None,
    ) -> list[dict]:
        """
        BFS up to `depth` hops from node_id.
        Optionally filter by edge rel type.
        Returns list of {"node_id", "type", "hops", ...node_attrs}.
        """
        if not self.G.has_node(node_id):
            return []

        from collections import deque  # pylint: disable=import-outside-toplevel
        visited: dict[str, int] = {node_id: 0}
        queue: deque[str] = deque([node_id])
        results: list[dict] = []

        while queue:
            current = queue.popleft()
            current_hops = visited[current]
            if current_hops >= depth:
                continue
            for neighbor in self.G.successors(current):
                if neighbor in visited:
                    continue
                edge_data = self.G[current][neighbor]
                if edge_type and edge_data.get("rel") != edge_type:
                    continue
                hops = current_hops + 1
                visited[neighbor] = hops
                node_data = dict(self.G.nodes[neighbor])
                node_data["node_id"] = neighbor
                node_data["hops"] = hops
                results.append(node_data)
                queue.append(neighbor)

        return results

    def hop_distance(self, src: str, dst: str) -> int:
        """Shortest hop distance; sys.maxsize if no path or either node missing."""
        if not self.G.has_node(src) or not self.G.has_node(dst):
            return sys.maxsize
        try:
            return nx.shortest_path_length(self.G, src, dst)
        except nx.NetworkXNoPath:
            return sys.maxsize

    def shortest_path(self, src: str, dst: str) -> list[str] | None:
        """Returns node list of shortest path, or None if no path."""
        try:
            return nx.shortest_path(self.G, src, dst)
        except (nx.NetworkXNoPath, nx.NodeNotFound):
            return None

    def subgraph_around(
        self,
        node_id: str,
        radius: int = 2,
        max_nodes: int = 200,
        max_edges: int = 500,
        hub_degree_limit: int = 500,
    ) -> dict:
        """
        Returns {"nodes": [...], "edges": [...], "truncated": bool} of the
        neighbourhood around node_id, bounded DURING expansion.

        Previously this used nx.ego_graph(), which materializes the full
        ego graph before any cap is applied — on syscall-heavy repos a
        depth-3 call could pull 10k+ nodes and inflate RSS past the circuit
        breaker limit. The bounded BFS stops at max_nodes and skips hub
        nodes whose degree exceeds hub_degree_limit (e.g. errnoErr/uintptr
        fan-out), which are connectivity noise rather than architecture.
        """
        if not self.G.has_node(node_id):
            return {"nodes": [], "edges": []}

        visited: set[str] = {node_id}
        frontier: list[str] = [node_id]
        truncated = False
        for _ in range(max(radius, 0)):
            if truncated or not frontier:
                break
            next_frontier: list[str] = []
            for n in frontier:
                for nb in nx.all_neighbors(self.G, n):
                    if nb in visited:
                        continue
                    if self.G.degree(nb) > hub_degree_limit:
                        continue  # hub node — skip, never expand
                    visited.add(nb)
                    next_frontier.append(nb)
                    if len(visited) >= max_nodes:
                        truncated = True
                        break
                if truncated:
                    break
            frontier = next_frontier

        nodes = []
        for n in visited:
            d = self.G.nodes[n]
            # Skip builtin names — they add noise, not signal, to neighbourhoods
            bare = n.split("::")[-1]
            if d.get("type") == "CONCEPT" and bare in PYTHON_BUILTINS:
                continue
            entry = dict(d)
            entry["node_id"] = n
            nodes.append(entry)

        edges = []
        for u, v, d in self.G.edges(visited, data=True):
            if v not in visited:
                continue
            if len(edges) >= max_edges:
                truncated = True
                break
            edge: dict = {
                "src": u, "dst": v,
                "rel": d.get("rel", "?"),
                "weight": d.get("weight", 1.0),
            }
            if "purpose" in d:
                edge["purpose"] = d["purpose"]
            edges.append(edge)

        return {"nodes": nodes, "edges": edges, "truncated": truncated}

    # ── stats ─────────────────────────────────────────────────────────────────

    def stats(self) -> dict:
        """Return basic node and edge counts."""
        return {
            "nodes": self.G.number_of_nodes(),
            "edges": self.G.number_of_edges(),
        }
