"""Graph data types and operations for function call graphs."""

from __future__ import annotations

import logging
import sqlite3
from collections import deque
from pathlib import Path
from typing import TYPE_CHECKING

import attrs

if TYPE_CHECKING:
    from .._model import FunctionSource

log = logging.getLogger(__name__)


@attrs.frozen
class FunctionNode:
    """A node in the call graph identifying a function by file and name."""

    file_path: Path = attrs.field(converter=Path)
    qualified_name: str


@attrs.frozen
class CalleeMetadata:
    """Metadata about a callee function resolved from Jedi analysis."""

    fully_qualified_name: str
    only_function_name: str
    definition_type: str
    source_line: str


@attrs.frozen
class CallEdge:
    """A directed edge in the call graph from caller to callee."""

    caller: FunctionNode
    callee: FunctionNode
    is_cross_file: bool
    call_count: int | None = None
    total_time_ns: int | None = None
    callee_metadata: CalleeMetadata | None = None


@attrs.define
class CallGraph:
    """A directed graph of function calls with forward/reverse indexes."""

    edges: list[CallEdge]
    _forward: dict[FunctionNode, list[CallEdge]] = attrs.field(
        factory=dict, init=False, repr=False
    )
    _reverse: dict[FunctionNode, list[CallEdge]] = attrs.field(
        factory=dict, init=False, repr=False
    )
    _nodes: set[FunctionNode] = attrs.field(
        factory=set, init=False, repr=False
    )

    def __attrs_post_init__(self) -> None:
        """Build forward/reverse/nodes indexes from edges."""
        fwd: dict[FunctionNode, list[CallEdge]] = {}
        rev: dict[FunctionNode, list[CallEdge]] = {}
        nodes: set[FunctionNode] = set()
        for edge in self.edges:
            fwd.setdefault(edge.caller, []).append(edge)
            rev.setdefault(edge.callee, []).append(edge)
            nodes.add(edge.caller)
            nodes.add(edge.callee)
        self._forward = fwd
        self._reverse = rev
        self._nodes = nodes

    @property
    def forward(self) -> dict[FunctionNode, list[CallEdge]]:
        """Forward adjacency: caller -> list of outgoing edges."""
        return self._forward

    @property
    def reverse(self) -> dict[FunctionNode, list[CallEdge]]:
        """Reverse adjacency: callee -> list of incoming edges."""
        return self._reverse

    @property
    def nodes(self) -> set[FunctionNode]:
        """All nodes appearing in the graph."""
        return self._nodes

    def callees_of(self, node: FunctionNode) -> list[CallEdge]:
        """Return edges where *node* is the caller."""
        return self.forward.get(node, [])

    def callers_of(self, node: FunctionNode) -> list[CallEdge]:
        """Return edges where *node* is the callee."""
        return self.reverse.get(node, [])

    def descendants(
        self,
        node: FunctionNode,
        max_depth: int | None = None,
    ) -> set[FunctionNode]:
        """Return all transitive callees of *node*."""
        visited: set[FunctionNode] = set()
        forward_map = self._forward
        if max_depth is None:
            queue: deque[FunctionNode] = deque([node])
            while queue:
                current = queue.popleft()
                for edge in forward_map.get(current, []):
                    if edge.callee not in visited:
                        visited.add(edge.callee)
                        queue.append(edge.callee)
        else:
            depth_queue: deque[tuple[FunctionNode, int]] = deque([(node, 0)])
            while depth_queue:
                current, depth = depth_queue.popleft()
                if depth >= max_depth:
                    continue
                for edge in forward_map.get(current, []):
                    if edge.callee not in visited:
                        visited.add(edge.callee)
                        depth_queue.append((edge.callee, depth + 1))
        return visited

    def ancestors(
        self,
        node: FunctionNode,
        max_depth: int | None = None,
    ) -> set[FunctionNode]:
        """Return all transitive callers of *node*."""
        visited: set[FunctionNode] = set()
        reverse_map = self._reverse
        if max_depth is None:
            queue: list[FunctionNode] = [node]
            while queue:
                current = queue.pop()
                for edge in reverse_map.get(current, []):
                    if edge.caller not in visited:
                        visited.add(edge.caller)
                        queue.append(edge.caller)
        else:
            depth_queue: list[tuple[FunctionNode, int]] = [(node, 0)]
            while depth_queue:
                current, depth = depth_queue.pop()
                if depth >= max_depth:
                    continue
                for edge in reverse_map.get(current, []):
                    if edge.caller not in visited:
                        visited.add(edge.caller)
                        depth_queue.append((edge.caller, depth + 1))
        return visited

    def subgraph(self, nodes: set[FunctionNode]) -> CallGraph:
        """Return a new graph containing only edges between *nodes*."""
        filtered = [
            e for e in self.edges if e.caller in nodes and e.callee in nodes
        ]
        return CallGraph(edges=filtered)

    def leaf_functions(self) -> set[FunctionNode]:
        """Return nodes with no outgoing edges."""
        return self.nodes - set(self.forward.keys())

    def root_functions(self) -> set[FunctionNode]:
        """Return nodes with no incoming edges."""
        return self.nodes - set(self.reverse.keys())

    def topological_order(self) -> list[FunctionNode]:
        """Return nodes in reverse topological order (leaves first)."""
        in_degree: dict[FunctionNode, int] = {}
        all_nodes = self._nodes
        for node in all_nodes:
            in_degree.setdefault(node, 0)
        for edge in self.edges:
            in_degree[edge.callee] = in_degree.get(edge.callee, 0) + 1

        forward_map = self._forward
        queue = deque(node for node, deg in in_degree.items() if deg == 0)
        result: list[FunctionNode] = []
        while queue:
            node = queue.popleft()
            result.append(node)
            for edge in forward_map.get(node, []):
                in_degree[edge.callee] -= 1
                if in_degree[edge.callee] == 0:
                    queue.append(edge.callee)

        if len(result) < len(all_nodes):
            log.warning(
                "Call graph contains cycles: %d of %d nodes "
                "excluded from topological order",
                len(all_nodes) - len(result),
                len(all_nodes),
            )

        result.reverse()
        return result


@attrs.frozen
class IndexResult:
    """Result of indexing a single file for call edges."""

    file_path: Path = attrs.field(converter=Path)
    cached: bool
    num_edges: int
    edges: tuple[tuple[str, str, bool], ...]
    cross_file_edges: int
    error: bool


def augment_with_trace(
    graph: CallGraph,
    trace_db_path: Path,
) -> CallGraph:
    """Augment a call graph with profiling data from a trace database."""
    conn = sqlite3.connect(str(trace_db_path))
    try:
        rows = conn.execute(
            "SELECT filename, function, class_name, "
            "call_count_nonrecursive, total_time_ns FROM pstats"
        ).fetchall()
    except sqlite3.OperationalError:
        conn.close()
        return graph
    conn.close()

    lookup: dict[tuple[str, str], tuple[int, int]] = {}
    for filename, function, class_name, call_count, total_time in rows:
        qn = f"{class_name}.{function}" if class_name else function
        lookup[(filename, qn)] = (call_count, total_time)

    augmented_edges: list[CallEdge] = []
    for edge in graph.edges:
        callee_file = str(edge.callee.file_path)
        callee_qn = edge.callee.qualified_name
        stats = lookup.get((callee_file, callee_qn))
        if stats is not None:
            call_count, total_time = stats
            augmented_edges.append(
                CallEdge(
                    caller=edge.caller,
                    callee=edge.callee,
                    is_cross_file=edge.is_cross_file,
                    call_count=call_count,
                    total_time_ns=total_time,
                    callee_metadata=edge.callee_metadata,
                )
            )
        else:
            augmented_edges.append(edge)

    return CallGraph(edges=augmented_edges)


def callees_from_graph(
    graph: CallGraph,
) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]:
    """Extract callee FunctionSource objects from a call graph."""
    from collections import defaultdict  # noqa: PLC0415

    from .._model import FunctionSource  # noqa: PLC0415

    file_path_to_function_source: dict[Path, set[FunctionSource]] = (
        defaultdict(set)
    )
    function_source_list: list[FunctionSource] = []

    for edge in graph.edges:
        meta = edge.callee_metadata
        if meta is None:
            continue
        callee_path = edge.callee.file_path
        fs = FunctionSource(
            file_path=callee_path,
            qualified_name=edge.callee.qualified_name,
            fully_qualified_name=meta.fully_qualified_name,
            source_code=meta.source_line,
            only_function_name=meta.only_function_name,
            definition_type=meta.definition_type,
        )
        file_path_to_function_source[callee_path].add(fs)
        function_source_list.append(fs)

    return dict(file_path_to_function_source), function_source_list
