"""Call graph construction and querying using NetworkX.

Builds a call graph from Python source files to understand
function relationships and trace data paths.
"""

from __future__ import annotations

import ast
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set

# NetworkX is optional - provide fallback
try:
    import networkx as nx

    NETWORKX_AVAILABLE = True
except ImportError:
    NETWORKX_AVAILABLE = False
    nx = None


@dataclass
class FunctionInfo:
    """Information about a function in the call graph."""

    name: str
    file_path: str
    start_line: int
    end_line: int
    is_async: bool = False
    is_method: bool = False
    class_name: str = ""
    calls: List[str] = field(default_factory=list)
    called_by: List[str] = field(default_factory=list)
    parameters: List[str] = field(default_factory=list)

    @property
    def full_name(self) -> str:
        """Get fully qualified name."""
        if self.class_name:
            return f"{self.file_path}:{self.class_name}.{self.name}"
        return f"{self.file_path}:{self.name}"

    def __str__(self) -> str:
        prefix = "async " if self.is_async else ""
        return f"{prefix}{self.name}() at {self.file_path}:{self.start_line}"


class CallGraph:
    """Build and query call graph from Python source files.

    Uses NetworkX for graph operations when available, with
    fallback to simple dict-based implementation.
    """

    def __init__(self, source_files: Dict[str, str]):
        """Initialize the call graph.

        Args:
            source_files: Dict mapping file paths to source code
        """
        self.files = source_files
        self.functions: Dict[str, FunctionInfo] = {}
        self._edges: Dict[str, Set[str]] = {}  # caller -> set of callees
        self._reverse_edges: Dict[str, Set[str]] = {}  # callee -> set of callers
        self._unresolved_calls: List[tuple[str, str]] = []  # (caller, callee_attr)

        if NETWORKX_AVAILABLE:
            self.graph = nx.DiGraph()
        else:
            self.graph = None

    def build(self) -> None:
        """Build call graph from source files."""
        # First pass: collect all function definitions
        for filepath, source in self.files.items():
            try:
                tree = ast.parse(source)
                self._collect_functions(tree, filepath)
            except SyntaxError:
                continue

        # Second pass: collect function calls
        for filepath, source in self.files.items():
            try:
                tree = ast.parse(source)
                self._collect_calls(tree, filepath)
            except SyntaxError:
                continue

    def _collect_functions(self, tree: ast.AST, filepath: str) -> None:
        """Collect all function definitions from AST.

        Args:
            tree: Parsed AST
            filepath: Path to source file
        """
        class FunctionCollector(ast.NodeVisitor):
            def __init__(self, graph: CallGraph, file_path: str):
                self.graph = graph
                self.file_path = file_path
                self.current_class: Optional[str] = None

            def visit_ClassDef(self, node: ast.ClassDef) -> None:
                old_class = self.current_class
                self.current_class = node.name
                self.generic_visit(node)
                self.current_class = old_class

            def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
                self._add_function(node, is_async=False)
                self.generic_visit(node)

            def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
                self._add_function(node, is_async=True)
                self.generic_visit(node)

            def _add_function(
                self,
                node: ast.FunctionDef | ast.AsyncFunctionDef,
                is_async: bool,
            ) -> None:
                # Build full name
                if self.current_class:
                    full_name = f"{self.file_path}:{self.current_class}.{node.name}"
                else:
                    full_name = f"{self.file_path}:{node.name}"

                # Extract parameter names
                params = []
                for arg in node.args.args:
                    params.append(arg.arg)

                info = FunctionInfo(
                    name=node.name,
                    file_path=self.file_path,
                    start_line=node.lineno,
                    end_line=node.end_lineno or node.lineno,
                    is_async=is_async,
                    is_method=self.current_class is not None,
                    class_name=self.current_class or "",
                    parameters=params,
                )

                self.graph.functions[full_name] = info
                self.graph._edges[full_name] = set()
                self.graph._reverse_edges[full_name] = set()

                if self.graph.graph is not None:
                    self.graph.graph.add_node(full_name, info=info)

        FunctionCollector(self, filepath).visit(tree)

    def _collect_calls(self, tree: ast.AST, filepath: str) -> None:
        """Collect function calls from AST.

        Args:
            tree: Parsed AST
            filepath: Path to source file
        """
        class CallCollector(ast.NodeVisitor):
            def __init__(self, graph: CallGraph, file_path: str):
                self.graph = graph
                self.file_path = file_path
                self.current_function: Optional[str] = None
                self.current_class: Optional[str] = None

            def visit_ClassDef(self, node: ast.ClassDef) -> None:
                old_class = self.current_class
                self.current_class = node.name
                self.generic_visit(node)
                self.current_class = old_class

            def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
                self._visit_function(node)

            def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
                self._visit_function(node)

            def _visit_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
                old_func = self.current_function
                if self.current_class:
                    self.current_function = f"{self.file_path}:{self.current_class}.{node.name}"
                else:
                    self.current_function = f"{self.file_path}:{node.name}"
                self.generic_visit(node)
                self.current_function = old_func

            def visit_Call(self, node: ast.Call) -> None:
                if self.current_function:
                    callee = self._get_callee_name(node)
                    if callee:
                        self.graph._add_edge(self.current_function, callee)
                    elif isinstance(node.func, ast.Attribute):
                        # Track unresolved receiver calls for diagnostics
                        self.graph._unresolved_calls.append(
                            (self.current_function, node.func.attr)
                        )

                    # H-119: Track functions passed as arguments (higher-order calls)
                    # e.g., map(process, items) -> edge from caller to process
                    for arg in node.args:
                        if isinstance(arg, ast.Name):
                            arg_name = arg.id
                            # Check if this name refers to a known function
                            for fn, info in self.graph.functions.items():
                                if info.name == arg_name:
                                    self.graph._add_edge(self.current_function, arg_name)
                                    break
                self.generic_visit(node)

            def _get_callee_name(self, node: ast.Call) -> Optional[str]:
                if isinstance(node.func, ast.Name):
                    # Direct function call: foo()
                    name = node.func.id
                    # Class instantiation: resolve ClassName() -> ClassName.__init__
                    if name and name[0].isupper():
                        # Likely a class — try to resolve to Class.__init__
                        for fn, info in self.graph.functions.items():
                            if info.name == "__init__" and info.class_name == name:
                                return f"{info.class_name}.__init__"
                    return name
                elif isinstance(node.func, ast.Attribute):
                    # Method call: obj.method()
                    # Try to resolve receiver type via assignment tracking
                    receiver_type = self._resolve_receiver_type(node.func.value)
                    if receiver_type:
                        return f"{receiver_type}.{node.func.attr}"
                    # Receiver unresolved — skip rather than bare-name matching
                    # which would create false edges to every function with the
                    # same name (C-38).
                    return None
                return None

            def _resolve_receiver_type(self, node: ast.expr) -> Optional[str]:
                """Try to determine the type of a receiver expression.

                Handles: ClassName(), self, cls, super(), and parameter
                type annotations within the enclosing function.
                """
                if isinstance(node, ast.Name):
                    if node.id == "self" and self.current_class:
                        return self.current_class
                    if node.id in ("super", "cls"):
                        return self.current_class
                    # Look up parameter type annotations in the enclosing function
                    return self._resolve_from_param_annotation(node.id)
                elif isinstance(node, ast.Call):
                    # ClassName() -> ClassName
                    if isinstance(node.func, ast.Name) and node.func.id[0:1].isupper():
                        return node.func.id
                    if isinstance(node.func, ast.Name) and node.func.id == "super":
                        return self.current_class
                return None

            def _resolve_from_param_annotation(self, name: str) -> Optional[str]:
                """Resolve a variable name to a type via function parameter annotations.

                Walks the enclosing function's args list to find a matching
                parameter with a type annotation.

                Args:
                    name: Variable/parameter name to look up.

                Returns:
                    The annotation type name, or None if not found.
                """
                if not self.current_function:
                    return None
                # Find the enclosing function node
                for fn_full, info in self.graph.functions.items():
                    if fn_full == self.current_function:
                        # Walk the tree to find the actual AST node
                        # so we can inspect its args annotations.
                        break
                else:
                    return None

                # Search the file's AST for the function node
                source = self.graph.files.get(self.file_path)
                if not source:
                    return None
                try:
                    tree = ast.parse(source)
                except SyntaxError:
                    return None
                for node in ast.walk(tree):
                    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        if node.lineno == info.start_line:
                            for arg in node.args.args:
                                if arg.arg == name and arg.annotation:
                                    try:
                                        return ast.unparse(arg.annotation)
                                    except Exception:
                                        return None
                return None

        CallCollector(self, filepath).visit(tree)

    def _add_edge(self, caller: str, callee_name: str) -> None:
        """Add edge from caller to callee.

        Args:
            caller: Full name of calling function
            callee_name: Simple or qualified name (e.g., "foo" or "ClassName.method")
        """
        # Check for qualified name (ClassName.method or ClassName.__init__)
        if "." in callee_name:
            class_name, method_name = callee_name.rsplit(".", 1)
            # Find by class_name + method_name
            candidates = [
                fn for fn, info in self.functions.items()
                if info.name == method_name and info.class_name == class_name
            ]
        else:
            # Find all matching functions by bare name
            candidates = [fn for fn, info in self.functions.items() if info.name == callee_name]

        if not candidates:
            return

        # Prefer same-file match to avoid cross-module false positives
        caller_file = self.functions[caller].file_path if caller in self.functions else ""
        same_file = [fn for fn in candidates if self.functions[fn].file_path == caller_file]

        if same_file:
            target = same_file[0]
        elif len(candidates) == 1:
            target = candidates[0]
        else:
            # Multiple cross-file candidates — skip to avoid false positives
            return

        # Add edge (deduplicate in lists)
        self._edges[caller].add(target)
        self._reverse_edges[target].add(caller)

        # Update function info (deduplicate)
        if caller not in self.functions[target].called_by:
            self.functions[target].called_by.append(caller)
        if caller in self.functions and target not in self.functions[caller].calls:
            self.functions[caller].calls.append(target)

        # Add to NetworkX graph
        if self.graph is not None:
            self.graph.add_edge(caller, target)

    def get_callers(self, function: str) -> List[str]:
        """Get functions that call this function.

        Args:
            function: Full function name

        Returns:
            List of caller function names
        """
        if self.graph is not None and function in self.graph:
            return list(self.graph.predecessors(function))
        return list(self._reverse_edges.get(function, set()))

    def get_unresolved_calls(self) -> List[tuple[str, str]]:
        """Return unresolved (caller, callee_attr) pairs."""
        return list(self._unresolved_calls)

    def get_callees(self, function: str) -> List[str]:
        """Get functions called by this function.

        Args:
            function: Full function name

        Returns:
            List of callee function names
        """
        if self.graph is not None and function in self.graph:
            return list(self.graph.successors(function))
        return list(self._edges.get(function, set()))

    def trace_data_path(self, start_func: str, end_func: str) -> List[str]:
        """Find shortest path from start to end function.

        Args:
            start_func: Starting function name
            end_func: Ending function name

        Returns:
            List of function names in path, or empty if no path
        """
        if self.graph is not None:
            try:
                return list(nx.shortest_path(self.graph, start_func, end_func))
            except (nx.NetworkXNoPath, nx.NodeNotFound):
                return []

        # Fallback: BFS
        return self._bfs_path(start_func, end_func)

    def _bfs_path(self, start: str, end: str) -> List[str]:
        """BFS to find path between functions.

        Args:
            start: Starting function
            end: Ending function

        Returns:
            Path as list of function names
        """
        if start not in self._edges or end not in self._edges:
            return []

        visited = {start}
        queue = [(start, [start])]

        while queue:
            node, path = queue.pop(0)
            if node == end:
                return path

            for neighbor in self._edges.get(node, set()):
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append((neighbor, path + [neighbor]))

        return []

    def get_transitive_callers(
        self,
        function: str,
        max_depth: int = 5,
    ) -> Set[str]:
        """Get all functions that transitively call this function.

        Args:
            function: Target function name
            max_depth: Maximum depth to search

        Returns:
            Set of caller function names
        """
        callers: Set[str] = set()

        def dfs(node: str, depth: int) -> None:
            if depth > max_depth:
                return
            for caller in self.get_callers(node):
                if caller not in callers:
                    callers.add(caller)
                    dfs(caller, depth + 1)

        dfs(function, 0)
        return callers

    def get_transitive_callees(
        self,
        function: str,
        max_depth: int = 5,
    ) -> Set[str]:
        """Get all functions transitively called by this function.

        Args:
            function: Starting function name
            max_depth: Maximum depth to search

        Returns:
            Set of callee function names
        """
        callees: Set[str] = set()

        def dfs(node: str, depth: int) -> None:
            if depth > max_depth:
                return
            for callee in self.get_callees(node):
                if callee not in callees:
                    callees.add(callee)
                    dfs(callee, depth + 1)

        dfs(function, 0)
        return callees

    def get_subgraph_around(
        self,
        function: str,
        depth: int = 2,
    ) -> Dict[str, Set[str]]:
        """Get subgraph around a function.

        Args:
            function: Center function
            depth: Depth of neighborhood

        Returns:
            Dict of node -> neighbors in subgraph
        """
        nodes: Set[str] = {function}

        # BFS to find nearby nodes
        current_level = {function}
        for _ in range(depth):
            next_level: Set[str] = set()
            for node in current_level:
                next_level.update(self._edges.get(node, set()))
                next_level.update(self._reverse_edges.get(node, set()))
            nodes.update(next_level)
            current_level = next_level

        # Build subgraph
        subgraph: Dict[str, Set[str]] = {}
        for node in nodes:
            subgraph[node] = self._edges.get(node, set()) & nodes

        return subgraph

    def find_function_at_line(
        self,
        filepath: str,
        line: int,
    ) -> Optional[FunctionInfo]:
        """Find the function containing a specific line.

        Args:
            filepath: Path to source file
            line: Line number

        Returns:
            FunctionInfo or None
        """
        candidates = []
        for func in self.functions.values():
            if func.file_path == filepath:
                if func.start_line <= line <= func.end_line:
                    candidates.append(func)

        if not candidates:
            return None

        # Return innermost function (smallest range)
        return min(candidates, key=lambda f: f.end_line - f.start_line)

    def format_dot(self) -> str:
        """Format call graph as DOT for visualization.

        Returns:
            DOT format string
        """
        lines = ["digraph CallGraph {", "    rankdir=LR;"]

        for func_name in self.functions:
            # Escape special characters
            label = func_name.replace(":", "\\n").replace(".", "\\n")
            lines.append(f'    "{func_name}" [label="{label}"];')

        for caller, callees in self._edges.items():
            for callee in callees:
                lines.append(f'    "{caller}" -> "{callee}";')

        lines.append("}")
        return "\n".join(lines)


# Convenience functions


def build_call_graph(source_files: Dict[str, str]) -> CallGraph:
    """Build a call graph from source files.

    Args:
        source_files: Dict mapping file paths to source code

    Returns:
        Built CallGraph
    """
    graph = CallGraph(source_files)
    graph.build()
    return graph


def find_callers(
    source_code: str,
    function_name: str,
    file_path: str = "source.py",
) -> List[str]:
    """Find functions that call a given function.

    Args:
        source_code: Python source code
        function_name: Name of function to find callers for
        file_path: Virtual file path

    Returns:
        List of caller function names
    """
    graph = CallGraph({file_path: source_code})
    graph.build()

    # Find the function
    for full_name in graph.functions:
        if graph.functions[full_name].name == function_name:
            return graph.get_callers(full_name)

    return []
