core/graph_mapper.py

"""
core/graph_mapper.py
====================

AST-based Python codebase architecture mapper.

Builds a directed graph (``networkx.DiGraph``) where:
- **Nodes** = top-level functions, methods, and classes discovered in the project.
- **Edges** = relationships between them:
    * ``calls``       — function-to-function / method invocations within the project
    * ``imports``     — module-level ``import`` / ``from ... import ...`` statements
    * ``inherits``    — class inheritance
    * ``contains``    — class-to-method containment (lightweight, optional)

Nodes carry metadata used by the Cytoscape.js front-end:
- ``id``         : fully-qualified node id (e.g. ``core/agents.py::AgentRouter.run``)
- ``label``      : short display name
- ``kind``       : ``function`` | ``method`` | ``class``
- ``module``     : POSIX-style file path relative to the project root
- ``cluster``    : top-level folder (e.g. ``core``, ``ui``, ``plugins``) — used for visual clustering
- ``file``       : same as ``module``
- ``lineno``     : line number in the source file
- ``degree``     : populated after graph build (for sizing nodes in the UI)

Public API
----------
- ``build_graph(project_root, ignore_dirs=None) -> nx.DiGraph``
- ``export_cytoscape_json(graph) -> dict``
- ``export_d3_json(graph) -> dict``
- ``dump_cytoscape_to_file(project_root, out_path) -> dict``

The module has **no third-party dependency except NetworkX**, and degrades
gracefully if NetworkX is missing (a tiny shim is provided so the file is
still importable in minimal environments — the heavy lifting requires NX).
"""

from __future__ import annotations

import ast
import json
import os
import re
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

try:
    import networkx as nx
    _HAS_NX = True
except Exception:  # pragma: no cover — soft fallback so import never crashes
    _HAS_NX = False

    class _MiniDiGraph:
        """Minimal DiGraph shim — just enough for build_graph + export."""
        def __init__(self):
            self._nodes: Dict[str, Dict[str, Any]] = {}
            self._edges: List[Tuple[str, str, Dict[str, Any]]] = []

        def add_node(self, n, **attrs):
            self._nodes.setdefault(n, {}).update(attrs)

        def add_edge(self, u, v, **attrs):
            self.add_node(u)
            self.add_node(v)
            self._edges.append((u, v, dict(attrs)))

        @property
        def nodes(self):
            return self._nodes

        def edges(self, data=False):
            if data:
                return list(self._edges)
            return [(u, v) for u, v, _ in self._edges]

        def number_of_nodes(self):
            return len(self._nodes)

        def number_of_edges(self):
            return len(self._edges)

        def degree(self, n):
            d = 0
            for u, v, _ in self._edges:
                if u == n or v == n:
                    d += 1
            return d

    class _NXShim:
        DiGraph = _MiniDiGraph

    nx = _NXShim()  # type: ignore


# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

DEFAULT_IGNORE_DIRS: Set[str] = {
    "__pycache__",
    ".git",
    ".venv",
    "venv",
    "env",
    "node_modules",
    ".mypy_cache",
    ".pytest_cache",
    "build",
    "dist",
    ".tox",
}


# ---------------------------------------------------------------------------
# AST visitor
# ---------------------------------------------------------------------------

class _SymbolCollector(ast.NodeVisitor):
    """
    First pass: collect every function/class definition the project exposes,
    keyed by ``module::QualName`` and by short name for fuzzy call resolution.
    """

    def __init__(self, module_path: str):
        self.module_path = module_path  # POSIX-style relative path
        self.symbols: Dict[str, Dict[str, Any]] = {}
        self._stack: List[str] = []  # qualified-name stack

    # ------------------------------------------------------------------
    def _qual(self, name: str) -> str:
        return ".".join(self._stack + [name])

    def _node_id(self, qual: str) -> str:
        return f"{self.module_path}::{qual}"

    # ------------------------------------------------------------------
    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        qual = self._qual(node.name)
        nid = self._node_id(qual)
        self.symbols[nid] = {
            "id": nid,
            "label": node.name,
            "qualname": qual,
            "kind": "class",
            "module": self.module_path,
            "file": self.module_path,
            "lineno": node.lineno,
            "bases": [_safe_dump(b) for b in node.bases],
        }
        self._stack.append(node.name)
        self.generic_visit(node)
        self._stack.pop()

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

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

    def _handle_func(self, node: ast.AST) -> None:
        name = getattr(node, "name", "<lambda>")
        qual = self._qual(name)
        nid = self._node_id(qual)
        kind = "method" if self._stack else "function"
        self.symbols[nid] = {
            "id": nid,
            "label": name,
            "qualname": qual,
            "kind": kind,
            "module": self.module_path,
            "file": self.module_path,
            "lineno": getattr(node, "lineno", 0),
        }
        self._stack.append(name)
        self.generic_visit(node)
        self._stack.pop()


class _EdgeCollector(ast.NodeVisitor):
    """
    Second pass: walk each function body to record call edges, plus capture
    module-level imports and class inheritance edges.
    """

    def __init__(
        self,
        module_path: str,
        symbols_by_short: Dict[str, List[str]],
        modules_by_dotted: Dict[str, str],
    ):
        self.module_path = module_path
        self.symbols_by_short = symbols_by_short
        self.modules_by_dotted = modules_by_dotted
        self.edges: List[Tuple[str, str, Dict[str, Any]]] = []
        self.containment: List[Tuple[str, str]] = []
        self._stack: List[str] = []
        # local alias table: name-in-this-file -> module path it points to
        self._import_aliases: Dict[str, str] = {}

    # ------------------------------------------------------------------
    def _qual(self, name: str) -> str:
        return ".".join(self._stack + [name])

    def _current_node_id(self) -> Optional[str]:
        if not self._stack:
            return None
        return f"{self.module_path}::{'.'.join(self._stack)}"

    # ------------------------------------------------------------------
    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            mod = self.modules_by_dotted.get(alias.name)
            self._import_aliases[alias.asname or alias.name.split(".")[0]] = mod or alias.name
            if mod:
                self.edges.append(
                    (
                        f"{self.module_path}::<module>",
                        f"{mod}::<module>",
                        {"type": "imports"},
                    )
                )
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        if node.module:
            mod = self.modules_by_dotted.get(node.module)
            if mod:
                self.edges.append(
                    (
                        f"{self.module_path}::<module>",
                        f"{mod}::<module>",
                        {"type": "imports"},
                    )
                )
                for alias in node.names:
                    self._import_aliases[alias.asname or alias.name] = mod
        self.generic_visit(node)

    # ------------------------------------------------------------------
    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        qual = self._qual(node.name)
        nid = f"{self.module_path}::{qual}"

        for base in node.bases:
            base_name = _safe_dump(base)
            target = self._resolve_short(base_name)
            if target:
                self.edges.append((nid, target, {"type": "inherits"}))

        self._stack.append(node.name)
        # containment: class -> its methods (added on method visit)
        self.generic_visit(node)
        self._stack.pop()

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

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

    def _handle_func(self, node: ast.AST) -> None:
        name = getattr(node, "name", "<lambda>")
        qual = self._qual(name)
        nid = f"{self.module_path}::{qual}"

        # Containment: enclosing class -> this method
        if self._stack:
            parent = f"{self.module_path}::{'.'.join(self._stack)}"
            self.containment.append((parent, nid))

        self._stack.append(name)
        self.generic_visit(node)
        self._stack.pop()

    # ------------------------------------------------------------------
    def visit_Call(self, node: ast.Call) -> None:
        caller = self._current_node_id()
        if caller:
            callee_short = _call_target_name(node.func)
            if callee_short:
                target = self._resolve_short(callee_short)
                if target and target != caller:
                    self.edges.append((caller, target, {"type": "calls"}))
        self.generic_visit(node)

    # ------------------------------------------------------------------
    def _resolve_short(self, short: str) -> Optional[str]:
        """
        Best-effort resolution of a textual reference to a project node id.
        Handles ``foo``, ``Foo.bar``, ``module.func``.
        """
        if not short:
            return None

        # Drop everything after the last dot if multi-part, but also try the
        # last component on its own (e.g. ``self.helper`` -> ``helper``).
        candidates = [short]
        if "." in short:
            candidates.append(short.split(".")[-1])
            candidates.append(short.split(".")[0])

        for cand in candidates:
            ids = self.symbols_by_short.get(cand)
            if not ids:
                continue
            # Prefer a match inside the same module if available
            same_mod = [i for i in ids if i.startswith(f"{self.module_path}::")]
            if same_mod:
                return same_mod[0]
            return ids[0]
        return None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _safe_dump(node: ast.AST) -> str:
    """Render an AST expression node to dotted text, best-effort."""
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        return f"{_safe_dump(node.value)}.{node.attr}"
    if isinstance(node, ast.Call):
        return _safe_dump(node.func)
    try:
        return ast.unparse(node)  # type: ignore[attr-defined]
    except Exception:
        return ""


def _call_target_name(func: ast.AST) -> str:
    """Return the textual call target (e.g. ``module.fn`` or ``obj.method``)."""
    return _safe_dump(func)


def _module_dotted(rel_path: str) -> str:
    """``core/agents.py`` -> ``core.agents``."""
    p = rel_path.replace("\\", "/")
    if p.endswith(".py"):
        p = p[:-3]
    if p.endswith("/__init__"):
        p = p[: -len("/__init__")]
    return p.replace("/", ".")


def _cluster_for(rel_path: str) -> str:
    """First path segment, used for visual clustering. ``app.py`` -> ``(root)``."""
    rel_path = rel_path.replace("\\", "/")
    head, _, tail = rel_path.partition("/")
    return head if tail else "(root)"


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def build_graph(
    project_root: str,
    ignore_dirs: Optional[Iterable[str]] = None,
) -> "nx.DiGraph":
    """
    Walk ``project_root`` for ``*.py`` files, parse each with ``ast``, and
    return a populated ``networkx.DiGraph``.
    """
    if not _HAS_NX:
        # We still return our shim, but warn the caller via graph attribute.
        graph = nx.DiGraph()
    else:
        graph = nx.DiGraph()

    root = Path(project_root).resolve()
    if not root.exists():
        raise FileNotFoundError(f"project_root does not exist: {project_root}")

    ignored = set(DEFAULT_IGNORE_DIRS) | set(ignore_dirs or [])

    # ---- Discover .py files ------------------------------------------------
    py_files: List[Path] = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in ignored and not d.startswith(".")]
        for fn in filenames:
            if fn.endswith(".py"):
                py_files.append(Path(dirpath) / fn)

    # ---- Build module dotted-name lookup -----------------------------------
    modules_by_dotted: Dict[str, str] = {}
    rel_paths: Dict[Path, str] = {}
    for f in py_files:
        rel = f.relative_to(root).as_posix()
        rel_paths[f] = rel
        modules_by_dotted[_module_dotted(rel)] = rel

    # ---- Pass 1: collect all symbols ---------------------------------------
    all_symbols: Dict[str, Dict[str, Any]] = {}
    per_file_tree: Dict[Path, ast.AST] = {}
    parse_errors: List[Dict[str, str]] = []

    for f in py_files:
        rel = rel_paths[f]
        try:
            src = f.read_text(encoding="utf-8", errors="replace")
            tree = ast.parse(src, filename=str(f))
        except SyntaxError as exc:
            parse_errors.append({"file": rel, "error": f"SyntaxError: {exc}"})
            continue
        per_file_tree[f] = tree

        collector = _SymbolCollector(module_path=rel)
        collector.visit(tree)
        all_symbols.update(collector.symbols)

        # virtual <module> node — useful for import edges + clustering
        mod_id = f"{rel}::<module>"
        all_symbols[mod_id] = {
            "id": mod_id,
            "label": rel.split("/")[-1],
            "qualname": "<module>",
            "kind": "module",
            "module": rel,
            "file": rel,
            "lineno": 1,
        }

    # short-name index for call resolution
    symbols_by_short: Dict[str, List[str]] = {}
    for nid, meta in all_symbols.items():
        symbols_by_short.setdefault(meta["label"], []).append(nid)
        qn = meta.get("qualname", "")
        if qn and qn != meta["label"]:
            symbols_by_short.setdefault(qn, []).append(nid)

    # ---- Materialize nodes -------------------------------------------------
    for nid, meta in all_symbols.items():
        graph.add_node(
            nid,
            label=meta["label"],
            kind=meta["kind"],
            module=meta["module"],
            file=meta["file"],
            cluster=_cluster_for(meta["module"]),
            qualname=meta.get("qualname", ""),
            lineno=meta.get("lineno", 0),
        )

    # ---- Pass 2: edges -----------------------------------------------------
    for f, tree in per_file_tree.items():
        rel = rel_paths[f]
        ec = _EdgeCollector(
            module_path=rel,
            symbols_by_short=symbols_by_short,
            modules_by_dotted=modules_by_dotted,
        )
        ec.visit(tree)

        for u, v, attrs in ec.edges:
            if u in graph.nodes and v in graph.nodes:
                graph.add_edge(u, v, **attrs)

        for parent, child in ec.containment:
            if parent in graph.nodes and child in graph.nodes:
                graph.add_edge(parent, child, type="contains")

    # ---- Annotate degree (for UI node sizing) ------------------------------
    for nid in list(graph.nodes):
        try:
            deg = graph.degree(nid)
        except Exception:
            deg = 0
        graph.nodes[nid]["degree"] = int(deg if not isinstance(deg, tuple) else deg[1])

    # Stash metadata on the graph itself
    if hasattr(graph, "graph"):
        graph.graph["project_root"] = str(root)
        graph.graph["python_files_scanned"] = len(py_files)
        graph.graph["parse_errors"] = parse_errors

    return graph


# ---------------------------------------------------------------------------
# Exporters
# ---------------------------------------------------------------------------

def export_cytoscape_json(graph: "nx.DiGraph") -> Dict[str, Any]:
    """
    Serialize ``graph`` into Cytoscape.js ``elements`` format::

        {
          "elements": {
            "nodes": [{"data": {...}}, ...],
            "edges": [{"data": {...}}, ...]
          },
          "meta": {...}
        }
    """
    nodes_out: List[Dict[str, Any]] = []
    edges_out: List[Dict[str, Any]] = []

    if hasattr(graph, "nodes"):
        node_iter = graph.nodes(data=True) if not isinstance(graph.nodes, dict) else graph.nodes.items()
    else:
        node_iter = []

    for nid, attrs in node_iter:
        data = {"id": nid, **attrs}
        nodes_out.append({"data": data})

    if hasattr(graph, "edges"):
        edge_iter = graph.edges(data=True)
    else:
        edge_iter = []

    for i, edge in enumerate(edge_iter):
        if len(edge) == 3:
            u, v, attrs = edge
        else:
            u, v = edge[0], edge[1]
            attrs = {}
        edges_out.append(
            {
                "data": {
                    "id": f"e{i}",
                    "source": u,
                    "target": v,
                    **attrs,
                }
            }
        )

    meta = {
        "node_count": len(nodes_out),
        "edge_count": len(edges_out),
        "clusters": sorted({n["data"].get("cluster", "(root)") for n in nodes_out}),
    }
    if hasattr(graph, "graph") and isinstance(graph.graph, dict):
        meta.update(
            {
                "project_root": graph.graph.get("project_root"),
                "python_files_scanned": graph.graph.get("python_files_scanned"),
                "parse_errors": graph.graph.get("parse_errors", []),
            }
        )

    return {
        "elements": {"nodes": nodes_out, "edges": edges_out},
        "meta": meta,
    }


def export_d3_json(graph: "nx.DiGraph") -> Dict[str, Any]:
    """Optional D3.js-style ``{nodes, links}`` export."""
    cyto = export_cytoscape_json(graph)
    nodes = [n["data"] for n in cyto["elements"]["nodes"]]
    links = [
        {"source": e["data"]["source"], "target": e["data"]["target"], **{k: v for k, v in e["data"].items() if k not in ("source", "target", "id")}}
        for e in cyto["elements"]["edges"]
    ]
    return {"nodes": nodes, "links": links, "meta": cyto["meta"]}


def dump_cytoscape_to_file(
    project_root: str,
    out_path: str,
    ignore_dirs: Optional[Iterable[str]] = None,
) -> Dict[str, Any]:
    """Build the graph and write Cytoscape JSON to ``out_path``. Returns the dict."""
    g = build_graph(project_root, ignore_dirs=ignore_dirs)
    payload = export_cytoscape_json(g)
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    Path(out_path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return payload


def detect_relevant_scope(
    project_root: str,
    seed_files: Optional[Iterable[str]] = None,
) -> Dict[str, Any]:
    """Expand a set of seed files into a graph-backed safe edit scope.

    The returned ``allowed_files`` set includes the seed files themselves plus
    directly connected Python files discovered by GraphMapper. This is used by
    the Diff Guard to block patches that wander outside the planner's intended
    blast radius.
    """
    root = Path(project_root).resolve()
    normalized_seed: List[str] = []
    for rel in list(seed_files or []):
        clean = str(rel or "").replace("\\", "/").strip().lstrip("/")
        clean = re.sub(r"/+", "/", clean)
        if clean and clean not in normalized_seed:
            normalized_seed.append(clean)

    allowed = set(normalized_seed)
    related = set()
    graph_nodes_considered = 0

    if not root.exists() or not root.is_dir():
        return {
            "seed_files": sorted(allowed),
            "allowed_files": sorted(allowed),
            "related_files": [],
            "graph_nodes_considered": 0,
        }

    try:
        graph = build_graph(str(root))
    except Exception:
        graph = None

    if graph is not None:
        edges = []
        try:
            edges = list(graph.edges(data=True))
        except Exception:
            try:
                edges = list(graph.edges())
            except Exception:
                edges = []

        seed_nodes = set()
        node_items = []
        try:
            node_items = list(graph.nodes.items())
        except Exception:
            node_items = []
        for node_id, attrs in node_items:
            file_rel = str((attrs or {}).get("file") or "")
            if file_rel in allowed:
                seed_nodes.add(node_id)
        graph_nodes_considered = len(seed_nodes)

        touched_nodes = set(seed_nodes)
        for edge in edges:
            try:
                src, dst = edge[0], edge[1]
            except Exception:
                continue
            if src in seed_nodes or dst in seed_nodes:
                touched_nodes.add(src)
                touched_nodes.add(dst)

        for node_id in touched_nodes:
            attrs = {}
            try:
                attrs = graph.nodes[node_id]
            except Exception:
                attrs = {}
            file_rel = str((attrs or {}).get("file") or "")
            if file_rel and file_rel not in allowed:
                related.add(file_rel)
                allowed.add(file_rel)

    return {
        "seed_files": sorted(set(normalized_seed)),
        "allowed_files": sorted(allowed),
        "related_files": sorted(related),
        "graph_nodes_considered": graph_nodes_considered,
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":  # pragma: no cover
    import argparse

    ap = argparse.ArgumentParser(description="Map a Python codebase into a Cytoscape graph.")
    ap.add_argument("project_root", help="Path to the project root to scan")
    ap.add_argument(
        "-o", "--out",
        default="graph.json",
        help="Output JSON file (Cytoscape format). Default: graph.json",
    )
    ap.add_argument(
        "--d3",
        action="store_true",
        help="Also dump a D3-style {nodes, links} file alongside.",
    )
    args = ap.parse_args()

    g = build_graph(args.project_root)
    payload = export_cytoscape_json(g)
    Path(args.out).write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(f"[graph_mapper] wrote {args.out}  "
          f"({payload['meta']['node_count']} nodes, {payload['meta']['edge_count']} edges)")

    if args.d3:
        d3_path = Path(args.out).with_suffix(".d3.json")
        d3_path.write_text(json.dumps(export_d3_json(g), indent=2), encoding="utf-8")
        print(f"[graph_mapper] wrote {d3_path}")