store/call_graph.py

"""
call_graph.py
-------------
Builds a function-level call graph from parsed codebase (FileInfo objects)
and persists it as call_graph.json for cross-module dependency analysis.

Graph structure (stored in JSON):
{
    "module.FileName.ClassName.function_name": {
        "id":         "module.FileName.ClassName.function_name",
        "name":       "function_name",
        "class_name": "ClassName",
        "module":     "module",
        "file":       "relative/path/to/file.py",
        "calls":      ["other_func", "AnotherClass.method"],
        "called_by":  ["parent_func", "caller_func"]
    },
    ...
}

Depends on:
    - ast_parser.parse_codebase() → list[FileInfo]
    - rich
"""

import json
from pathlib import Path
from dataclasses import dataclass, field
from collections import defaultdict

from rich.console import Console
from rich.tree import Tree
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich import box

from ingest.parse_ast import parse_codebase, FileInfo

console = Console()

# ── Constants ─────────────────────────────────────────────────────────────────

CALL_GRAPH_PATH = "./call_graph.json"


# ── Node Model ────────────────────────────────────────────────────────────────

@dataclass
class CallNode:
    """Represents a single function/method in the call graph."""
    id:         str                    # unique fully-qualified ID
    name:       str                    # function name
    class_name: str                    # class name or "" for top-level
    module:     str                    # top-level module/package
    file:       str                    # relative file path
    calls:      list[str] = field(default_factory=list)      # what it calls
    called_by:  list[str] = field(default_factory=list)      # what calls it

    def to_dict(self) -> dict:
        return {
            "id":         self.id,
            "name":       self.name,
            "class_name": self.class_name,
            "module":     self.module,
            "file":       self.file,
            "calls":      self.calls,
            "called_by":  self.called_by,
        }

    @staticmethod
    def from_dict(data: dict) -> "CallNode":
        return CallNode(
            id=data["id"],
            name=data["name"],
            class_name=data.get("class_name", ""),
            module=data["module"],
            file=data["file"],
            calls=data.get("calls", []),
            called_by=data.get("called_by", []),
        )


# ── CallGraph ─────────────────────────────────────────────────────────────────

class CallGraph:
    """
    Builds, stores, queries, and persists the function call graph
    for a monolithic codebase.
    """

    def __init__(self):
        # node_id → CallNode
        self._graph: dict[str, CallNode] = {}
        # name → list of node_ids (for fuzzy lookup by function name alone)
        self._name_index: dict[str, list[str]] = defaultdict(list)

    # ── Build ─────────────────────────────────────────────────────────────────

    def build(self, parsed_files: list[FileInfo]) -> None:
        """
        Build the call graph from a list of parsed FileInfo objects.
        Two passes:
            1. Register all functions as nodes
            2. Resolve call relationships and populate called_by

        Args:
            parsed_files: Output of ast_parser.parse_codebase().
        """
        self._graph.clear()
        self._name_index.clear()

        # Pass 1 — register all nodes
        for file_info in parsed_files:
            # Top-level functions
            for func in file_info.functions:
                node = CallNode(
                    id=self._make_id(file_info, func.name, class_name=""),
                    name=func.name,
                    class_name="",
                    module=file_info.module,
                    file=file_info.relative,
                    calls=func.calls,
                )
                self._register(node)

            # Class methods
            for cls in file_info.classes:
                for method in cls.methods:
                    node = CallNode(
                        id=self._make_id(file_info, method.name, class_name=cls.name),
                        name=method.name,
                        class_name=cls.name,
                        module=file_info.module,
                        file=file_info.relative,
                        calls=method.calls,
                    )
                    self._register(node)

        # Pass 2 — resolve called_by relationships
        for node in self._graph.values():
            for raw_call in node.calls:
                # Strip object prefix: "self.calculate" → "calculate"
                callee_name = raw_call.split(".")[-1]
                candidate_ids = self._name_index.get(callee_name, [])
                for callee_id in candidate_ids:
                    callee = self._graph.get(callee_id)
                    if callee and node.id not in callee.called_by:
                        callee.called_by.append(node.id)

        total_nodes = len(self._graph)
        total_edges = sum(len(n.calls) for n in self._graph.values())
        console.print(
            f"[green]✔[/green] Call graph built: "
            f"[cyan]{total_nodes}[/cyan] nodes · "
            f"[magenta]{total_edges}[/magenta] edges\n"
        )

    def _make_id(self, file_info: FileInfo,
                 func_name: str,
                 class_name: str) -> str:
        """Generate a unique fully-qualified node ID."""
        stem = Path(file_info.relative).stem   # filename without extension
        parts = [file_info.module, stem]
        if class_name:
            parts.append(class_name)
        parts.append(func_name)
        return ".".join(parts)

    def _register(self, node: CallNode) -> None:
        """Register a node in the graph and name index."""
        self._graph[node.id] = node
        self._name_index[node.name].append(node.id)

    # ── Persist ───────────────────────────────────────────────────────────────

    def save(self, path: str = CALL_GRAPH_PATH) -> None:
        """
        Persist the call graph to a JSON file.

        Args:
            path: File path to write call_graph.json.
        """
        data = {node_id: node.to_dict() for node_id, node in self._graph.items()}

        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

        size_kb = round(Path(path).stat().st_size / 1024, 2)
        console.print(
            f"[green]✔[/green] Call graph saved to "
            f"[cyan]{path}[/cyan] ({size_kb} KB)\n"
        )

    def load(self, path: str = CALL_GRAPH_PATH) -> None:
        """
        Load a persisted call graph from JSON.

        Args:
            path: File path to read call_graph.json from.

        Raises:
            FileNotFoundError: If the file does not exist.
        """
        graph_path = Path(path)
        if not graph_path.exists():
            raise FileNotFoundError(
                f"Call graph not found at '{path}'. "
                "Run build_and_save() first."
            )

        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)

        self._graph.clear()
        self._name_index.clear()

        for node_id, node_data in data.items():
            node = CallNode.from_dict(node_data)
            self._graph[node_id] = node
            self._name_index[node.name].append(node_id)

        console.print(
            f"[green]✔[/green] Call graph loaded: "
            f"[cyan]{len(self._graph)}[/cyan] nodes from "
            f"[cyan]{path}[/cyan]\n"
        )

    # ── Query API ─────────────────────────────────────────────────────────────

    def get_node(self, function_name: str,
                 class_name: str | None = None) -> list[CallNode]:
        """
        Look up nodes by function name, optionally filtered by class.

        Args:
            function_name: Name of the function/method to look up.
            class_name:    Optional class name to narrow results.

        Returns:
            List of matching CallNode objects.
        """
        candidate_ids = self._name_index.get(function_name, [])
        nodes = [self._graph[nid] for nid in candidate_ids if nid in self._graph]

        if class_name:
            nodes = [n for n in nodes if n.class_name == class_name]

        return nodes

    def get_calls(self, function_name: str,
                  class_name: str | None = None) -> list[CallNode]:
        """
        Get all functions that the given function calls (outgoing edges).

        Args:
            function_name: Source function name.
            class_name:    Optional class scope.

        Returns:
            List of CallNode objects this function calls.
        """
        source_nodes = self.get_node(function_name, class_name)
        if not source_nodes:
            return []

        results = []
        seen = set()
        for source in source_nodes:
            for raw_call in source.calls:
                callee_name = raw_call.split(".")[-1]
                for callee_node in self.get_node(callee_name):
                    if callee_node.id not in seen:
                        seen.add(callee_node.id)
                        results.append(callee_node)

        return results

    def get_called_by(self, function_name: str,
                      class_name: str | None = None) -> list[CallNode]:
        """
        Get all functions that call the given function (incoming edges).

        Args:
            function_name: Target function name.
            class_name:    Optional class scope.

        Returns:
            List of CallNode objects that call this function.
        """
        target_nodes = self.get_node(function_name, class_name)
        if not target_nodes:
            return []

        results = []
        seen = set()
        for target in target_nodes:
            for caller_id in target.called_by:
                if caller_id not in seen and caller_id in self._graph:
                    seen.add(caller_id)
                    results.append(self._graph[caller_id])

        return results

    def get_impact(self, function_name: str,
                   class_name: str | None = None,
                   depth: int = 2) -> dict[str, list[CallNode]]:
        """
        Trace the impact of changing a function — who calls it,
        who calls those callers, up to `depth` levels.

        Args:
            function_name: Function to analyze.
            class_name:    Optional class scope.
            depth:         How many levels up to trace (default 2).

        Returns:
            Dict mapping depth level string → list of affected CallNodes.
        """
        impact: dict[str, list[CallNode]] = {}
        current_level = self.get_node(function_name, class_name)
        visited = {n.id for n in current_level}

        for level in range(1, depth + 1):
            next_level = []
            for node in current_level:
                for caller_id in node.called_by:
                    if caller_id not in visited and caller_id in self._graph:
                        visited.add(caller_id)
                        next_level.append(self._graph[caller_id])
            if not next_level:
                break
            impact[f"level_{level}"] = next_level
            current_level = next_level

        return impact

    def stats(self) -> dict:
        """Return summary statistics for the call graph."""
        total_edges = sum(len(n.calls) for n in self._graph.values())
        modules = list({n.module for n in self._graph.values()})
        isolated = [n for n in self._graph.values()
                    if not n.calls and not n.called_by]
        return {
            "total_nodes":    len(self._graph),
            "total_edges":    total_edges,
            "modules":        modules,
            "isolated_nodes": len(isolated),
        }

    def is_loaded(self) -> bool:
        """Return True if graph has been built or loaded."""
        return len(self._graph) > 0

    # ── Rich Renderers ────────────────────────────────────────────────────────

    def render_stats(self) -> None:
        """Render call graph statistics as a rich panel."""
        s = self.stats()

        table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
        table.add_column(style="dim")
        table.add_column(style="bold white")

        table.add_row("Total nodes",    str(s["total_nodes"]))
        table.add_row("Total edges",    str(s["total_edges"]))
        table.add_row("Modules",        ", ".join(s["modules"]))
        table.add_row("Isolated nodes", str(s["isolated_nodes"]))
        table.add_row(
            "Status",
            "[bold green]✔ Loaded[/bold green]"
            if self.is_loaded()
            else "[bold red]✘ Empty[/bold red]"
        )

        console.print(Panel(
            table,
            title="[bold cyan]Call Graph Stats[/bold cyan]",
            border_style="cyan",
        ))

    def render_node(self, node: CallNode) -> None:
        """Render a single node with its calls and callers."""
        scope = f"{node.class_name}.{node.name}" if node.class_name else node.name

        calls_str    = "\n".join(f"  → {c}" for c in node.calls) or "  [dim]none[/dim]"
        called_by_str = "\n".join(f"  ← {c}" for c in node.called_by) or "  [dim]none[/dim]"

        content = Text()
        content.append("Module:  ", style="dim")
        content.append(f"{node.module}\n")
        content.append("File:    ", style="dim")
        content.append(f"{node.file}\n\n")
        content.append("Calls:\n", style="bold cyan")
        content.append(calls_str + "\n\n")
        content.append("Called by:\n", style="bold magenta")
        content.append(called_by_str)

        console.print(Panel(
            content,
            title=f"[bold white]{scope}[/bold white]",
            border_style="cyan",
        ))

    def render_impact(self, function_name: str,
                      impact: dict[str, list[CallNode]]) -> None:
        """Render the impact tree of a function change."""
        root = Tree(
            f"[bold yellow]⚡ Impact of changing:[/bold yellow] "
            f"[bold white]{function_name}[/bold white]",
            guide_style="dim yellow",
        )

        if not impact:
            root.add("[dim]No callers found — isolated function.[/dim]")
        else:
            for level, nodes in impact.items():
                level_num = level.split("_")[1]
                level_branch = root.add(
                    f"[bold cyan]Level {level_num} — {len(nodes)} affected[/bold cyan]"
                )
                for node in nodes:
                    scope = (
                        f"{node.class_name}.{node.name}"
                        if node.class_name else node.name
                    )
                    level_branch.add(
                        f"[magenta]{scope}[/magenta]  "
                        f"[dim]{node.file}[/dim]"
                    )

        console.print(root)

    def render_dependencies(self, function_name: str,
                            calls: list[CallNode],
                            called_by: list[CallNode]) -> None:
        """Render outgoing and incoming dependencies for a function."""
        table = Table(
            "Direction", "Function", "Class", "Module", "File",
            box=box.ROUNDED,
            header_style="bold bright_cyan",
            border_style="cyan",
            show_lines=True,
            title=f"Dependencies: {function_name}",
        )

        for node in calls:
            table.add_row(
                "[cyan]→ calls[/cyan]",
                node.name,
                node.class_name or "[dim]—[/dim]",
                node.module,
                node.file,
            )

        for node in called_by:
            table.add_row(
                "[magenta]← called by[/magenta]",
                node.name,
                node.class_name or "[dim]—[/dim]",
                node.module,
                node.file,
            )

        if not calls and not called_by:
            console.print(f"[yellow]No dependencies found for '{function_name}'.[/yellow]")
        else:
            console.print(table)


# ── Convenience Pipeline ──────────────────────────────────────────────────────

def build_and_save(root_path: str,
                   graph_path: str = CALL_GRAPH_PATH) -> CallGraph:
    """
    Full pipeline: parse codebase → build graph → save to JSON.

    Args:
        root_path:  Absolute path to the monolithic codebase root.
        graph_path: Path to write call_graph.json.

    Returns:
        Built and saved CallGraph instance.
    """
    console.rule("[bold cyan]Call Graph Builder[/bold cyan]")
    console.print(f"[bold]📂 Root:[/bold] {root_path}\n")

    parsed_files = parse_codebase(root_path)

    graph = CallGraph()
    graph.build(parsed_files)
    graph.save(graph_path)
    graph.render_stats()

    return graph


def load_graph(graph_path: str = CALL_GRAPH_PATH) -> CallGraph:
    """
    Load a persisted call graph from JSON.

    Args:
        graph_path: Path to call_graph.json.

    Returns:
        Loaded CallGraph instance.
    """
    graph = CallGraph()
    graph.load(graph_path)
    return graph


# ── Singleton Access ──────────────────────────────────────────────────────────

_graph_instance: CallGraph | None = None


def get_call_graph(graph_path: str = CALL_GRAPH_PATH) -> CallGraph:
    """
    Return a singleton CallGraph instance.
    Loads from JSON on first call, reuses on subsequent calls.

    Args:
        graph_path: Path to call_graph.json.

    Returns:
        Shared CallGraph instance.
    """
    global _graph_instance
    if _graph_instance is None:
        _graph_instance = load_graph(graph_path)
    return _graph_instance


# ── Entry Point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        console.print(
            "[red]Usage:[/red] python call_graph.py <codebase_path> "
            "[function_name]"
        )
        sys.exit(1)

    root  = sys.argv[1]
    query = sys.argv[2] if len(sys.argv) > 2 else None

    # Build and save
    graph = build_and_save(root)

    if query:
        console.rule(f"[bold cyan]Query: {query}[/bold cyan]")

        # Node detail
        nodes = graph.get_node(query)
        if not nodes:
            console.print(f"[yellow]Function '{query}' not found in graph.[/yellow]")
        else:
            for node in nodes:
                graph.render_node(node)

        # Dependencies
        calls     = graph.get_calls(query)
        called_by = graph.get_called_by(query)
        graph.render_dependencies(query, calls, called_by)

        # Impact analysis
        impact = graph.get_impact(query, depth=2)
        graph.render_impact(query, impact)
    else:
        # Show sample nodes
        console.rule("[cyan]Sample Nodes[/cyan]")
        sample = list(graph._graph.values())[:5]
        for node in sample:
            graph.render_node(node)