"""Knowledge graph construction for data flow analysis.

Graph model:
  Nodes: entry_point, sanitizer, sink, trust_boundary
  Edges: data_flow (directed, from source toward sink)

The LLM reasons over this graph to assess whether sanitization
actually covers identified threat paths.
"""

from __future__ import annotations

import ast
import warnings
from collections import defaultdict
from pathlib import Path
from typing import Any

import networkx as nx

from sanicode.scanner.data_flow import (
    detect_entry_points,
    detect_sanitizers,
    detect_sinks,
)


class KnowledgeGraph:
    """In-memory knowledge graph backed by NetworkX DiGraph."""

    def __init__(self) -> None:
        self._graph: nx.DiGraph = nx.DiGraph()
        self._node_counter: int = 0

    def _next_id(self, prefix: str) -> str:
        self._node_counter += 1
        return f"{prefix}_{self._node_counter}"

    def add_entry_point(
        self,
        label: str,
        file: Path | None = None,
        line: int | None = None,
        **attrs: Any,
    ) -> str:
        """Add an entry point node (e.g., an HTTP handler, CLI argument reader).

        Returns:
            The node ID assigned to this entry point.
        """
        node_id = self._next_id("entry")
        self._graph.add_node(
            node_id,
            kind="entry_point",
            label=label,
            file=str(file) if file else None,
            line=line,
            **attrs,
        )
        return node_id

    def add_sink(
        self,
        label: str,
        file: Path | None = None,
        line: int | None = None,
        cwe_id: int | None = None,
        **attrs: Any,
    ) -> str:
        """Add a sink node (e.g., SQL query, shell command, eval call).

        Returns:
            The node ID assigned to this sink.
        """
        node_id = self._next_id("sink")
        self._graph.add_node(
            node_id,
            kind="sink",
            label=label,
            file=str(file) if file else None,
            line=line,
            cwe_id=cwe_id,
            **attrs,
        )
        return node_id

    def add_sanitizer(
        self,
        label: str,
        file: Path | None = None,
        line: int | None = None,
        **attrs: Any,
    ) -> str:
        """Add a sanitizer node (e.g., html.escape, parameterized query wrapper).

        Returns:
            The node ID assigned to this sanitizer.
        """
        node_id = self._next_id("sanitizer")
        self._graph.add_node(
            node_id,
            kind="sanitizer",
            label=label,
            file=str(file) if file else None,
            line=line,
            **attrs,
        )
        return node_id

    def add_data_flow(self, from_node: str, to_node: str, **attrs: Any) -> None:
        """Add a directed data-flow edge between two graph nodes.

        Args:
            from_node: ID of the upstream node.
            to_node: ID of the downstream node.
        """
        if from_node not in self._graph:
            raise ValueError(f"Source node not in graph: {from_node!r}")
        if to_node not in self._graph:
            raise ValueError(f"Target node not in graph: {to_node!r}")
        self._graph.add_edge(from_node, to_node, **attrs)

    def to_dict(self) -> dict[str, Any]:
        """Serialize the graph to a plain dict (JSON-serializable).

        Returns a node-link format compatible with NetworkX's
        node_link_data / node_link_graph round-trip.
        """
        return nx.node_link_data(self._graph, edges="links")

    def nodes_by_kind(self, kind: str) -> list[tuple[str, dict]]:
        """Return all nodes of a given kind.

        Args:
            kind: Node kind to filter by (e.g. "entry_point", "sink", "sanitizer").

        Returns:
            List of (node_id, attrs) tuples for nodes whose ``kind`` attribute
            matches the requested value.
        """
        return [
            (node_id, attrs)
            for node_id, attrs in self._graph.nodes(data=True)
            if attrs.get("kind") == kind
        ]

    @classmethod
    def from_parsed(
        cls, file_trees: list[tuple[Path, ast.Module]]
    ) -> KnowledgeGraph:
        """Build a KnowledgeGraph from pre-parsed ASTs.

        Use this when the caller has already parsed the files (e.g., to also
        run pattern detection on the same trees without re-parsing).

        Args:
            file_trees: List of (file_path, parsed_ast) pairs.

        Returns:
            A populated KnowledgeGraph with heuristic data-flow edges.
        """
        kg = cls()
        kg._populate(file_trees)
        return kg

    @classmethod
    def from_source(cls, path: Path) -> KnowledgeGraph:
        """Build a KnowledgeGraph by scanning a source path.

        Parses one Python file or all ``.py`` files under a directory, runs the
        three data-flow detectors, adds the results as graph nodes, and then
        constructs heuristic data-flow edges within each enclosing function.

        Args:
            path: Path to a Python file or directory to scan.

        Returns:
            A populated KnowledgeGraph.

        Raises:
            FileNotFoundError: If ``path`` does not exist.
        """
        path = path.resolve()
        if not path.exists():
            raise FileNotFoundError(f"Path does not exist: {path}")

        file_trees: list[tuple[Path, ast.Module]] = []

        if path.is_file():
            tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
            file_trees.append((path, tree))
        else:
            for py_file in sorted(path.rglob("*.py")):
                try:
                    tree = ast.parse(
                        py_file.read_text(encoding="utf-8"), filename=str(py_file)
                    )
                    file_trees.append((py_file, tree))
                except SyntaxError as exc:
                    warnings.warn(
                        f"Skipping {py_file}: syntax error — {exc}",
                        stacklevel=2,
                    )

        kg = cls()
        kg._populate(file_trees)
        return kg

    def _populate(self, file_trees: list[tuple[Path, ast.Module]]) -> None:
        """Detect data-flow primitives and build heuristic edges.

        Args:
            file_trees: List of (file_path, parsed_ast) pairs.
        """
        # Nodes grouped by enclosing function name for edge construction.
        by_function: dict[str, dict[str, list[str]]] = defaultdict(
            lambda: {"entries": [], "sanitizers": [], "sinks": []}
        )

        for file_path, tree in file_trees:
            entry_points = detect_entry_points(tree, file_path)
            sinks = detect_sinks(tree, file_path)
            sanitizers = detect_sanitizers(tree, file_path)

            for ep in entry_points:
                node_id = self.add_entry_point(
                    label=ep.name,
                    file=ep.file,
                    line=ep.line,
                    entry_kind=ep.kind,
                    function=ep.function,
                )
                if ep.function is not None:
                    by_function[ep.function]["entries"].append(node_id)
                elif ep.kind == "http_handler":
                    # The function IS the entry point.  Sinks/sanitizers detected
                    # inside it will have function=ep.name, so group this entry
                    # under that same key so heuristic edges can be built.
                    by_function[ep.name]["entries"].append(node_id)

            for si in sinks:
                node_id = self.add_sink(
                    label=si.name,
                    file=si.file,
                    line=si.line,
                    cwe_id=si.cwe_id,
                    sink_kind=si.kind,
                    function=si.function,
                )
                if si.function is not None:
                    by_function[si.function]["sinks"].append(node_id)

            for san in sanitizers:
                node_id = self.add_sanitizer(
                    label=san.name,
                    file=san.file,
                    line=san.line,
                    sanitizer_kind=san.kind,
                    function=san.function,
                )
                if san.function is not None:
                    by_function[san.function]["sanitizers"].append(node_id)

        # Heuristic intra-function edge construction.
        #
        # For each function that has both entry points and sinks, create a
        # data-flow path.  When sanitizers are also present in the same function
        # they are chained in sequence: entry → san1 → san2 → … → sink.
        #
        # TODO: Cross-function edges (caller → callee data flow) are not yet
        # implemented.  This requires a call-graph pass and is deferred to a
        # future enhancement.
        for _fn_name, groups in by_function.items():
            entries = groups["entries"]
            sinks = groups["sinks"]
            sans = groups["sanitizers"]

            if not entries or not sinks:
                continue

            for entry_id in entries:
                for sink_id in sinks:
                    if sans:
                        chain = [entry_id, *sans, sink_id]
                        for upstream, downstream in zip(chain, chain[1:], strict=False):
                            self.add_data_flow(
                                upstream, downstream, confidence="heuristic"
                            )
                    else:
                        self.add_data_flow(entry_id, sink_id, confidence="heuristic")

    @property
    def node_count(self) -> int:
        """Number of nodes in the graph."""
        return self._graph.number_of_nodes()

    @property
    def edge_count(self) -> int:
        """Number of edges in the graph."""
        return self._graph.number_of_edges()

    def __repr__(self) -> str:
        return (
            f"KnowledgeGraph(nodes={self.node_count}, edges={self.edge_count})"
        )
