"""
MCP server for querying the security graph.

Loads a pre-built graph JSON file into a NetworkX MultiDiGraph and exposes
query tools over stdio so an AI agent can explore code-to-data relationships,
find orphaned resources, trace data flows, and analyze access patterns.

Usage:
    PYTHONPATH=/Users/jacobburley/Desktop/data_security python scripts/graph_mcp_server.py

The graph JSON file is generated by scripts/build_graph.py.
"""

import collections
import json
import logging
import os
import sys
import uuid
from datetime import datetime, timezone

import networkx as nx
from mcp.server.fastmcp import FastMCP

# All logging to stderr — stdout is reserved for MCP protocol
logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s: %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger(__name__)

GRAPH_DATA_PATH = os.environ.get(
    "GRAPH_DATA_PATH",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "graph_data.json"),
)

# Duplicated from lib/graph/loaders/_helpers.py to keep MCP server standalone
EDGE_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_URL, "dsec:graph:edge")


# ── In-memory graph store ────────────────────────────────────────────


class GraphStore:
    """Loads serialized graph JSON into a NetworkX MultiDiGraph with
    lightweight secondary indices for frequent type-based lookups."""

    def __init__(self, json_path: str):
        self.G: nx.MultiDiGraph = nx.MultiDiGraph()
        self.tables_by_name: dict[str, str] = {}         # table_name → urn
        self.nodes_by_type: dict[str, list[str]] = {}     # node_type → [urn, ...]
        self.edges_by_type: dict[str, list[tuple[str, str, str]]] = {}  # relation_type → [(from, to, key), ...]
        self.soft_links: list[dict] = []
        self.generated_at = "unknown"
        self._load(json_path)
        self._soft_links_path = os.path.join(
            os.path.dirname(json_path), "soft_links.json"
        )
        self._load_soft_links()

    def _load(self, json_path: str):
        with open(json_path) as f:
            data = json.load(f)

        self.generated_at = data.get("generated_at", "unknown")

        for node in data["nodes"]:
            urn = node["urn"]
            self.G.add_node(
                urn,
                node_type=node.get("node_type", "unknown"),
                parent_urn=node.get("parent_urn"),
                metadata=node.get("metadata", {}),
                organization_id=node.get("organization_id"),
            )
            node_type = node.get("node_type", "unknown")
            self.nodes_by_type.setdefault(node_type, []).append(urn)

            if node_type == "table":
                table_name = node.get("metadata", {}).get("table_name", "")
                if table_name:
                    self.tables_by_name[table_name] = urn

        for edge in data["edges"]:
            key = edge["uuid"]
            self.G.add_edge(
                edge["from_urn"],
                edge["to_urn"],
                key=key,
                relation_type=edge.get("relation_type", "UNKNOWN"),
                metadata=edge.get("metadata", {}),
                organization_id=edge.get("organization_id"),
            )
            rel_type = edge.get("relation_type", "UNKNOWN")
            self.edges_by_type.setdefault(rel_type, []).append(
                (edge["from_urn"], edge["to_urn"], key)
            )

        logger.info(
            "Loaded graph: %d nodes, %d edges",
            self.G.number_of_nodes(),
            self.G.number_of_edges(),
        )

    def _load_soft_links(self):
        if not os.path.exists(self._soft_links_path):
            logger.info("No soft_links.json found — skipping")
            return

        with open(self._soft_links_path) as f:
            data = json.load(f)

        loaded = 0
        for link in data.get("soft_links", []):
            from_urn = link["from_urn"]
            to_urn = link["to_urn"]
            if from_urn not in self.G:
                logger.warning("Soft link skipped — from_urn not in graph: %s", from_urn)
                continue
            if to_urn not in self.G:
                logger.warning("Soft link skipped — to_urn not in graph: %s", to_urn)
                continue

            relation_type = link.get("relation_type", "CODE_TO_DATA")
            edge_key = str(uuid.uuid5(
                EDGE_NAMESPACE, f"{from_urn}:{to_urn}:{relation_type}"
            ))
            org_id = self.G.nodes[from_urn].get("organization_id")

            self.G.add_edge(
                from_urn, to_urn, key=edge_key,
                relation_type=relation_type,
                metadata={
                    "detection_method": link.get("detection_method", "soft_link"),
                    "confidence": link.get("confidence", 0.7),
                    "note": link.get("note", ""),
                },
                organization_id=org_id,
            )
            self.edges_by_type.setdefault(relation_type, []).append(
                (from_urn, to_urn, edge_key)
            )
            self.soft_links.append(link)
            loaded += 1

        logger.info("Loaded %d soft link(s)", loaded)

    def _save_soft_links(self):
        tmp_path = self._soft_links_path + ".tmp"
        with open(tmp_path, "w") as f:
            json.dump({"soft_links": self.soft_links}, f, indent=2)
            f.write("\n")
        os.replace(tmp_path, self._soft_links_path)

    def node_dict(self, urn: str) -> dict | None:
        """Return a node as a dict with urn/node_type/parent_urn/metadata
        keys, or None if the URN doesn't exist."""
        if urn not in self.G:
            return None
        attrs = self.G.nodes[urn]
        return {
            "urn": urn,
            "node_type": attrs.get("node_type", "unknown"),
            "parent_urn": attrs.get("parent_urn"),
            "metadata": attrs.get("metadata", {}),
        }


# ── Load store at import time ────────────────────────────────────────

store = GraphStore(GRAPH_DATA_PATH)

# ── MCP server ───────────────────────────────────────────────────────

mcp = FastMCP("Security Graph Query Server")


def _node_label(node: dict) -> str:
    """Short display label for a node."""
    m = node.get("metadata", {})
    for key in ("function_name", "class_name", "column_name", "table_name",
                "schema_name", "repo_name", "database_name"):
        if key in m:
            return str(m[key])
    if "file_path" in m:
        return m["file_path"].rsplit("/", 1)[-1]
    return node["urn"].rsplit("/", 1)[-1]


def _format_node(node: dict, compact: bool = False) -> str:
    """Format a node as a readable string."""
    label = _node_label(node)
    ntype = node.get("node_type", "unknown")
    if compact:
        return f"[{ntype}] {label}  ({node['urn']})"
    lines = [
        f"[{ntype}] {label}",
        f"  URN: {node['urn']}",
    ]
    if node.get("parent_urn"):
        lines.append(f"  Parent: {node['parent_urn']}")
    for k, v in node.get("metadata", {}).items():
        lines.append(f"  {k}: {v}")
    return "\n".join(lines)


def _format_edge_data(from_urn: str, to_urn: str, data: dict) -> str:
    """Format an edge as a readable string."""
    parts = [f"{data.get('relation_type', '?')}: {from_urn} → {to_urn}"]
    for k, v in data.get("metadata", {}).items():
        parts.append(f"  {k}: {v}")
    return "\n".join(parts)


# ── Tool 1: Graph Summary ───────────────────────────────────────────


@mcp.tool()
def get_graph_summary() -> str:
    """Get overview statistics of the security graph including node counts
    by type (database, schema, table, column, codebase, file, class,
    function), edge counts by relation type, and generation timestamp."""
    node_counts = {
        ntype: len(urns) for ntype, urns in sorted(store.nodes_by_type.items())
    }
    edge_counts = {
        etype: len(triples) for etype, triples in sorted(store.edges_by_type.items())
    }

    lines = [
        f"Graph generated at: {store.generated_at}",
        f"Total nodes: {store.G.number_of_nodes()}",
        f"Total edges: {store.G.number_of_edges()}",
        "",
        "Nodes by type:",
    ]
    for ntype, count in node_counts.items():
        lines.append(f"  {ntype}: {count}")
    lines.append("")
    lines.append("Edges by relation type:")
    for etype, count in edge_counts.items():
        lines.append(f"  {etype}: {count}")
    lines.append("")
    lines.append(f"Database tables: {len(store.tables_by_name)}")
    lines.append(f"Table names: {', '.join(sorted(store.tables_by_name.keys()))}")
    return "\n".join(lines)


# ── Tool 2: Search Nodes ────────────────────────────────────────────


@mcp.tool()
def search_nodes(
    node_type: str = "",
    name_pattern: str = "",
    urn_pattern: str = "",
    metadata_key: str = "",
    metadata_value: str = "",
    limit: int = 50,
) -> str:
    """Search graph nodes by type, name, URN pattern, or metadata.

    Args:
        node_type: Filter by type (database, schema, table, column,
                   codebase, file, class, function).
        name_pattern: Substring match against common name fields
                      (table_name, class_name, function_name, etc.).
        urn_pattern: Substring match against the node URN.
        metadata_key: Filter to nodes that have this metadata key.
        metadata_value: If metadata_key is set, match this value substring.
        limit: Maximum results to return (default 50).
    """
    if node_type:
        candidate_urns = store.nodes_by_type.get(node_type, [])
    else:
        candidate_urns = list(store.G.nodes)

    results = []
    name_fields = [
        "table_name", "column_name", "schema_name", "class_name",
        "function_name", "file_path", "database_name", "repo_name",
    ]

    for urn in candidate_urns:
        attrs = store.G.nodes[urn]
        meta = attrs.get("metadata", {})

        if urn_pattern and urn_pattern.lower() not in urn.lower():
            continue
        if name_pattern:
            matched = any(
                name_pattern.lower() in str(meta.get(f, "")).lower()
                for f in name_fields
            )
            if not matched:
                continue
        if metadata_key:
            val = meta.get(metadata_key)
            if val is None:
                continue
            if metadata_value and metadata_value.lower() not in str(val).lower():
                continue

        results.append(store.node_dict(urn))
        if len(results) >= limit:
            break

    if not results:
        return "No nodes found matching the given criteria."

    lines = [f"Found {len(results)} node(s):"]
    for node in results:
        lines.append("")
        lines.append(_format_node(node, compact=False))
    return "\n".join(lines)


# ── Tool 3: Get Node Details ────────────────────────────────────────


@mcp.tool()
def get_node_details(urn: str) -> str:
    """Get full details for a specific node by URN, including all metadata
    and connected edges (both incoming and outgoing)."""
    node = store.node_dict(urn)
    if not node:
        return f"No node found with URN: {urn}"

    lines = [_format_node(node)]
    lines.append("")

    outgoing = list(store.G.out_edges(urn, data=True, keys=True))
    if outgoing:
        lines.append(f"Outgoing edges ({len(outgoing)}):")
        for _, to_urn, _key, data in outgoing:
            target = store.node_dict(to_urn)
            target_label = _node_label(target) if target else to_urn
            target_type = target["node_type"] if target else "?"
            lines.append(
                f"  --[{data.get('relation_type', '?')}]--> [{target_type}] {target_label}"
            )
            for k, v in data.get("metadata", {}).items():
                lines.append(f"      {k}: {v}")

    incoming = list(store.G.in_edges(urn, data=True, keys=True))
    if incoming:
        lines.append(f"Incoming edges ({len(incoming)}):")
        for from_urn, _, _key, data in incoming:
            source = store.node_dict(from_urn)
            source_label = _node_label(source) if source else from_urn
            source_type = source["node_type"] if source else "?"
            lines.append(
                f"  <--[{data.get('relation_type', '?')}]-- [{source_type}] {source_label}"
            )
            for k, v in data.get("metadata", {}).items():
                lines.append(f"      {k}: {v}")

    return "\n".join(lines)


# ── Tool 4: Get Neighbors ───────────────────────────────────────────


@mcp.tool()
def get_neighbors(
    urn: str,
    direction: str = "both",
    edge_type: str = "",
) -> str:
    """Get all nodes connected to a given node.

    Args:
        urn: The URN of the center node.
        direction: 'outgoing', 'incoming', or 'both' (default).
        edge_type: Optional filter by relation type
                   (CONTAINS, DATA_TO_DATA, CODE_TO_DATA, etc.).
    """
    if urn not in store.G:
        return f"No node found with URN: {urn}"

    edge_list: list[tuple[str, str, dict]] = []  # (neighbor_urn, arrow, data)
    if direction in ("outgoing", "both"):
        for _, to_urn, data in store.G.out_edges(urn, data=True):
            if edge_type and data.get("relation_type") != edge_type:
                continue
            arrow = f"--[{data.get('relation_type', '?')}]-->"
            edge_list.append((to_urn, arrow, data))
    if direction in ("incoming", "both"):
        for from_urn, _, data in store.G.in_edges(urn, data=True):
            if edge_type and data.get("relation_type") != edge_type:
                continue
            arrow = f"<--[{data.get('relation_type', '?')}]--"
            edge_list.append((from_urn, arrow, data))

    if not edge_list:
        return f"No neighbors found for {urn} (direction={direction}, edge_type={edge_type or 'any'})."

    lines = [f"Neighbors of {urn} ({len(edge_list)} edges):"]
    for neighbor_urn, arrow, _data in edge_list:
        neighbor = store.node_dict(neighbor_urn)
        if neighbor:
            lines.append(f"  {arrow} [{neighbor['node_type']}] {_node_label(neighbor)}")
            lines.append(f"          URN: {neighbor_urn}")
        else:
            lines.append(f"  {arrow} {neighbor_urn}")

    return "\n".join(lines)


# ── Tool 5: Find Code for Table ─────────────────────────────────────


@mcp.tool()
def find_code_for_table(table_name: str) -> str:
    """Given a database table name, find all code nodes (ORM classes and
    functions) that reference it via CODE_TO_DATA edges. Shows detection
    method and confidence for each link."""
    table_urn = store.tables_by_name.get(table_name)
    if not table_urn:
        available = ", ".join(sorted(store.tables_by_name.keys())[:20])
        return f"No table found with name '{table_name}'. Available tables: {available}..."

    # Find CODE_TO_DATA edges pointing INTO this table
    code_edges = [
        (from_urn, data)
        for from_urn, _, data in store.G.in_edges(table_urn, data=True)
        if data.get("relation_type") == "CODE_TO_DATA"
    ]

    if not code_edges:
        return f"No code references found for table '{table_name}'."

    orm_edges = [(u, d) for u, d in code_edges if d.get("metadata", {}).get("detection_method") == "orm_tablename"]
    func_edges = [(u, d) for u, d in code_edges if d.get("metadata", {}).get("detection_method") == "orm_reference"]

    lines = [f"Code references for table '{table_name}' ({len(code_edges)} total):"]

    if orm_edges:
        lines.append(f"\nORM Models ({len(orm_edges)}):")
        for from_urn, data in orm_edges:
            source = store.node_dict(from_urn)
            if source:
                lines.append(f"  {source['metadata'].get('class_name', '?')}")
                lines.append(f"    File: {source['metadata'].get('file_path', '?')}")
                lines.append(f"    URN: {from_urn}")
                lines.append(f"    Confidence: {data.get('metadata', {}).get('confidence', '?')}")

    if func_edges:
        lines.append(f"\nFunctions ({len(func_edges)}):")
        for from_urn, data in func_edges:
            source = store.node_dict(from_urn)
            if source:
                func_name = source["metadata"].get("function_name", "?")
                file_path = source["metadata"].get("file_path", "?")
                model = data.get("metadata", {}).get("referenced_model", "?")
                lines.append(f"  {func_name}()")
                lines.append(f"    File: {file_path}")
                lines.append(f"    Via ORM model: {model}")
                lines.append(f"    Confidence: {data.get('metadata', {}).get('confidence', '?')}")

    return "\n".join(lines)


# ── Tool 6: Find Tables for Code ────────────────────────────────────


@mcp.tool()
def find_tables_for_code(name: str) -> str:
    """Given a function or class name, find all database tables it
    references via CODE_TO_DATA edges."""
    matching = []
    for ntype in ("function", "class"):
        key = "function_name" if ntype == "function" else "class_name"
        for urn in store.nodes_by_type.get(ntype, []):
            meta = store.G.nodes[urn].get("metadata", {})
            if meta.get(key) == name:
                matching.append(store.node_dict(urn))

    if not matching:
        return f"No function or class found with name '{name}'."

    lines = []
    for node in matching:
        urn = node["urn"]
        code_edges = [
            (to_urn, data)
            for _, to_urn, data in store.G.out_edges(urn, data=True)
            if data.get("relation_type") == "CODE_TO_DATA"
        ]

        label = _node_label(node)
        ntype = node["node_type"]
        lines.append(f"[{ntype}] {label}")
        lines.append(f"  URN: {urn}")
        lines.append(f"  File: {node['metadata'].get('file_path', '?')}")

        if code_edges:
            lines.append(f"  Tables referenced ({len(code_edges)}):")
            for to_urn, data in code_edges:
                target = store.node_dict(to_urn)
                t_name = target["metadata"].get("table_name", "?") if target else "?"
                method = data.get("metadata", {}).get("detection_method", "?")
                lines.append(f"    → {t_name} (method: {method})")
        else:
            lines.append("  No table references found.")
        lines.append("")

    return "\n".join(lines)


# ── Tool 7: Find Orphaned Tables ────────────────────────────────────


@mcp.tool()
def find_orphaned_tables() -> str:
    """Find all database tables that have NO CODE_TO_DATA edges pointing
    to them. These are potential orphaned resources with no known code
    references in the scanned codebase."""
    referenced_urns: set[str] = set()
    for _from, to, _key in store.edges_by_type.get("CODE_TO_DATA", []):
        referenced_urns.add(to)

    orphaned = []
    for table_urn in store.nodes_by_type.get("table", []):
        if table_urn not in referenced_urns:
            orphaned.append(store.node_dict(table_urn))

    if not orphaned:
        return "No orphaned tables found — all tables have at least one code reference."

    lines = [f"Found {len(orphaned)} orphaned table(s) with no code references:"]
    for table in orphaned:
        name = table["metadata"].get("table_name", "?")
        schema = table["metadata"].get("schema_name", "?")
        lines.append(f"  {schema}.{name}")
        lines.append(f"    URN: {table['urn']}")

        # Check if it has DATA_TO_DATA edges (FK relationships)
        data_edge_count = 0
        for _, _, d in store.G.out_edges(table["urn"], data=True):
            if d.get("relation_type") == "DATA_TO_DATA":
                data_edge_count += 1
        for _, _, d in store.G.in_edges(table["urn"], data=True):
            if d.get("relation_type") == "DATA_TO_DATA":
                data_edge_count += 1
        if data_edge_count:
            lines.append(f"    FK relationships: {data_edge_count}")

    return "\n".join(lines)


# ── Tool 8: Table Access Map ────────────────────────────────────────


@mcp.tool()
def get_table_access_map(table_name: str = "") -> str:
    """Map tables to their code access patterns. For each table, shows
    ORM models and functions that access it.

    Args:
        table_name: Optional specific table name. If empty, shows map
                    for all tables that have code references.
    """
    # Group CODE_TO_DATA edges by target table name
    table_map: dict[str, list[tuple[str, dict]]] = {}
    for from_urn, to_urn, _key in store.edges_by_type.get("CODE_TO_DATA", []):
        target = store.node_dict(to_urn)
        if not target or target.get("node_type") != "table":
            continue
        t_name = target["metadata"].get("table_name", "?")
        if table_name and t_name != table_name:
            continue
        edge_data = store.G.edges[from_urn, to_urn, _key]
        table_map.setdefault(t_name, []).append((from_urn, edge_data))

    if not table_map:
        if table_name:
            return f"No code access patterns found for table '{table_name}'."
        return "No code access patterns found."

    lines = [f"Table access map ({len(table_map)} tables):"]
    for t_name in sorted(table_map.keys()):
        edges = table_map[t_name]
        orm_models = []
        functions = []
        for from_urn, data in edges:
            source = store.node_dict(from_urn)
            if not source:
                continue
            method = data.get("metadata", {}).get("detection_method", "")
            if method == "orm_tablename":
                orm_models.append(source)
            else:
                functions.append((source, data))

        lines.append(f"\n  {t_name}:")
        if orm_models:
            model_names = [m["metadata"].get("class_name", "?") for m in orm_models]
            lines.append(f"    ORM models: {', '.join(model_names)}")
        if functions:
            lines.append(f"    Functions ({len(functions)}):")
            for func, data in functions:
                func_name = func["metadata"].get("function_name", "?")
                file_path = func["metadata"].get("file_path", "?")
                via_model = data.get("metadata", {}).get("referenced_model", "?")
                lines.append(f"      {func_name}() in {file_path} (via {via_model})")

    return "\n".join(lines)


# ── Tool 9: Trace Data Path ─────────────────────────────────────────


@mcp.tool()
def trace_data_path(
    start_urn: str,
    edge_types: str = "DATA_TO_DATA,CODE_TO_DATA",
    max_depth: int = 10,
    direction: str = "outgoing",
) -> str:
    """BFS traversal from a starting node following specified edge types.
    Returns all reachable nodes with their distance from the start.
    Useful for data flow analysis and impact assessment.

    Args:
        start_urn: URN of the starting node.
        edge_types: Comma-separated relation types to follow
                    (e.g., 'DATA_TO_DATA,CODE_TO_DATA').
        max_depth: Maximum traversal depth (default 10).
        direction: 'outgoing' (default), 'incoming', or 'both'.
    """
    if start_urn not in store.G:
        return f"No node found with URN: {start_urn}"

    allowed = set(edge_types.split(","))
    visited: dict[str, int] = {}
    queue = collections.deque([(start_urn, 0)])

    while queue:
        current, depth = queue.popleft()
        if current in visited or depth > max_depth:
            continue
        visited[current] = depth

        if direction in ("outgoing", "both"):
            for _, to_urn, data in store.G.out_edges(current, data=True):
                if data.get("relation_type") in allowed and to_urn not in visited:
                    queue.append((to_urn, depth + 1))
        if direction in ("incoming", "both"):
            for from_urn, _, data in store.G.in_edges(current, data=True):
                if data.get("relation_type") in allowed and from_urn not in visited:
                    queue.append((from_urn, depth + 1))

    if len(visited) <= 1:
        return f"No reachable nodes from {start_urn} via edge types: {edge_types}"

    lines = [f"Reachable nodes from {start_urn} ({len(visited) - 1} found):"]
    for urn, depth in sorted(visited.items(), key=lambda x: x[1]):
        if urn == start_urn:
            continue
        node = store.node_dict(urn)
        if node:
            lines.append(f"  depth={depth}  [{node['node_type']}] {_node_label(node)}")
            lines.append(f"           URN: {urn}")
        else:
            lines.append(f"  depth={depth}  {urn}")

    return "\n".join(lines)


# ── Tool 10: Get Subgraph ───────────────────────────────────────────


@mcp.tool()
def get_subgraph(center_urn: str, hops: int = 2) -> str:
    """Extract a subgraph around a center node within N hops in all
    directions. Returns both nodes and edges. Useful for understanding
    the local neighborhood of any resource.

    Args:
        center_urn: URN of the center node.
        hops: Number of hops to expand (default 2).
    """
    if center_urn not in store.G:
        return f"No node found with URN: {center_urn}"

    # Use nx.ego_graph on undirected view to get all nodes within N hops
    undirected = store.G.to_undirected(as_view=True)
    ego_nodes = set(nx.ego_graph(undirected, center_urn, radius=hops).nodes())

    # Extract the directed subgraph for those nodes
    sub = store.G.subgraph(ego_nodes)

    # Summarize by type
    type_counts: dict[str, int] = {}
    for urn in sub.nodes():
        t = store.G.nodes[urn].get("node_type", "unknown")
        type_counts[t] = type_counts.get(t, 0) + 1

    lines = [
        f"Subgraph around {center_urn} ({hops} hops):",
        f"  Nodes: {sub.number_of_nodes()}",
        f"  Edges: {sub.number_of_edges()}",
        f"  Node types: {', '.join(f'{k}={v}' for k, v in sorted(type_counts.items()))}",
        "",
        "Nodes:",
    ]
    for urn in sub.nodes():
        node = store.node_dict(urn)
        if node:
            lines.append(f"  [{node['node_type']}] {_node_label(node)}")
            lines.append(f"    URN: {urn}")

    lines.append("")
    lines.append("Edges:")
    for from_urn, to_urn, data in sub.edges(data=True):
        from_node = store.node_dict(from_urn)
        to_node = store.node_dict(to_urn)
        from_label = _node_label(from_node) if from_node else from_urn
        to_label = _node_label(to_node) if to_node else to_urn
        lines.append(f"  {from_label} --[{data.get('relation_type', '?')}]--> {to_label}")

    return "\n".join(lines)


# ── Tool 11: Find Shortest Path ─────────────────────────────────────


@mcp.tool()
def find_shortest_path(
    from_urn: str,
    to_urn: str,
    edge_types: str = "",
) -> str:
    """Find the shortest path between two nodes in the graph, ignoring
    edge direction. Useful for understanding how two resources are
    related even if separated by many hops.

    Args:
        from_urn: URN of the starting node.
        to_urn: URN of the target node.
        edge_types: Optional comma-separated relation types to restrict
                    traversal (e.g., 'CODE_TO_DATA,CODE_TO_CODE').
                    If empty, all edge types are followed.
    """
    if from_urn not in store.G:
        return f"No node found with URN: {from_urn}"
    if to_urn not in store.G:
        return f"No node found with URN: {to_urn}"

    # Optionally filter to specific edge types (zero-copy view)
    if edge_types:
        allowed = set(edge_types.split(","))
        view = nx.subgraph_view(
            store.G,
            filter_edge=lambda u, v, k: store.G.edges[u, v, k].get("relation_type") in allowed,
        )
    else:
        view = store.G

    undirected = view.to_undirected()

    try:
        path = nx.shortest_path(undirected, from_urn, to_urn)
    except nx.NetworkXNoPath:
        return f"No path found between {from_urn} and {to_urn}."
    except nx.NodeNotFound as e:
        return str(e)

    hop_count = len(path) - 1
    lines = [f"Shortest path ({hop_count} hop{'s' if hop_count != 1 else ''}):"]

    for i, urn in enumerate(path):
        node = store.node_dict(urn)
        if node:
            lines.append(f"  [{node['node_type']}] {_node_label(node)}  ({urn})")
        else:
            lines.append(f"  {urn}")

        # Show the edge to the next hop
        if i < len(path) - 1:
            next_urn = path[i + 1]
            # Check for directed edge in either direction
            edge_label = _lookup_edge_label(urn, next_urn)
            lines.append(f"  {edge_label}")

    return "\n".join(lines)


def _lookup_edge_label(urn_a: str, urn_b: str) -> str:
    """Find the relation type between two adjacent nodes and return a
    formatted arrow showing direction."""
    # Forward edge: a → b
    if store.G.has_edge(urn_a, urn_b):
        for _key, data in store.G[urn_a][urn_b].items():
            return f"  --[{data.get('relation_type', '?')}]-->"
    # Reverse edge: b → a
    if store.G.has_edge(urn_b, urn_a):
        for _key, data in store.G[urn_b][urn_a].items():
            return f"  <--[{data.get('relation_type', '?')}]--"
    return "  --[?]-->"


# ── Tool 12: Find Data Hubs ─────────────────────────────────────────


@mcp.tool()
def find_data_hubs(min_tables: int = 3, limit: int = 10) -> str:
    """Find clusters of interconnected functions (via CODE_TO_CODE call
    chains) that collectively touch many database tables (via CODE_TO_DATA).
    Useful for identifying the most complex and data-heavy parts of the
    codebase.

    Args:
        min_tables: Minimum number of tables a cluster must touch to be
                    included (default 3).
        limit: Maximum number of clusters to return (default 10).
    """
    # Build subgraph of CODE_TO_CODE edges only
    c2c_triples = store.edges_by_type.get("CODE_TO_CODE", [])
    c2c_nodes: set[str] = set()
    c2c_graph = nx.DiGraph()
    for from_urn, to_urn, _key in c2c_triples:
        c2c_graph.add_edge(from_urn, to_urn)
        c2c_nodes.add(from_urn)
        c2c_nodes.add(to_urn)

    # Find weakly connected components (each = a call chain cluster)
    components: list[set[str]] = list(nx.weakly_connected_components(c2c_graph))

    # Also add isolated code nodes that have CODE_TO_DATA but no CODE_TO_CODE
    c2d_triples = store.edges_by_type.get("CODE_TO_DATA", [])
    code_nodes_with_data = set()
    for from_urn, _to, _key in c2d_triples:
        code_nodes_with_data.add(from_urn)

    isolated_code = code_nodes_with_data - c2c_nodes
    for urn in isolated_code:
        components.append({urn})

    # For each component, collect unique tables reachable via CODE_TO_DATA
    clusters = []
    for component in components:
        tables_by_func: dict[str, list[str]] = {}  # table_name → [function_urns]
        for func_urn in component:
            for _, to_urn, data in store.G.out_edges(func_urn, data=True):
                if data.get("relation_type") != "CODE_TO_DATA":
                    continue
                target = store.node_dict(to_urn)
                if target and target.get("node_type") == "table":
                    t_name = target["metadata"].get("table_name", to_urn)
                    tables_by_func.setdefault(t_name, []).append(func_urn)

        if len(tables_by_func) >= min_tables:
            clusters.append((component, tables_by_func))

    if not clusters:
        return f"No data hub clusters found with >= {min_tables} tables."

    # Sort by table count descending
    clusters.sort(key=lambda x: len(x[1]), reverse=True)
    clusters = clusters[:limit]

    lines = [f"Data Hub Clusters ({len(clusters)} found):"]

    for idx, (component, tables_by_func) in enumerate(clusters, 1):
        lines.append(f"\nCluster {idx}: {len(component)} functions → {len(tables_by_func)} tables")

        # Functions
        lines.append("  Functions:")
        for func_urn in sorted(component):
            node = store.node_dict(func_urn)
            if node:
                fname = node["metadata"].get("function_name", "?")
                fpath = node["metadata"].get("file_path", "?")
                lines.append(f"    {fname}() in {fpath}")

        # Call chain edges within this component
        internal_edges = []
        for func_urn in component:
            if func_urn not in c2c_graph:
                continue
            for _, to_urn in c2c_graph.out_edges(func_urn):
                if to_urn in component:
                    from_node = store.node_dict(func_urn)
                    to_node = store.node_dict(to_urn)
                    from_name = from_node["metadata"].get("function_name", "?") if from_node else "?"
                    to_name = to_node["metadata"].get("function_name", "?") if to_node else "?"
                    internal_edges.append(f"    {from_name} --[CODE_TO_CODE]--> {to_name}")

        if internal_edges:
            lines.append("  Call chain:")
            lines.extend(internal_edges)

        # Tables accessed
        lines.append("  Tables accessed:")
        for t_name in sorted(tables_by_func.keys()):
            func_urns = tables_by_func[t_name]
            func_names = []
            for fu in func_urns:
                n = store.node_dict(fu)
                func_names.append(n["metadata"].get("function_name", "?") if n else "?")
            lines.append(f"    {t_name} (via: {', '.join(func_names)})")

    return "\n".join(lines)


# ── Tool 13: Add Soft Link ────────────────────────────────────────────


@mcp.tool()
def add_soft_link(
    from_urn: str,
    to_urn: str,
    confidence: float = 0.7,
    note: str = "",
) -> str:
    """Manually create a CODE_TO_DATA edge between two URNs. The link is
    persisted to soft_links.json and immediately added to the in-memory
    graph.

    Args:
        from_urn: Source node URN (must exist in graph).
        to_urn: Target node URN (must exist in graph).
        confidence: Confidence score between 0.0 and 1.0 (default 0.7).
        note: Human-readable reason for the link.
    """
    if from_urn not in store.G:
        return f"Error: from_urn not found in graph: {from_urn}"
    if to_urn not in store.G:
        return f"Error: to_urn not found in graph: {to_urn}"

    for existing in store.soft_links:
        if existing["from_urn"] == from_urn and existing["to_urn"] == to_urn:
            return f"Error: soft link already exists (id={existing['id']})"

    relation_type = "CODE_TO_DATA"
    link_id = str(uuid.uuid4())
    edge_key = str(uuid.uuid5(
        EDGE_NAMESPACE, f"{from_urn}:{to_urn}:{relation_type}"
    ))
    org_id = store.G.nodes[from_urn].get("organization_id")

    link = {
        "id": link_id,
        "from_urn": from_urn,
        "to_urn": to_urn,
        "relation_type": relation_type,
        "detection_method": "soft_link",
        "confidence": confidence,
        "note": note,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }

    store.soft_links.append(link)
    store.G.add_edge(
        from_urn, to_urn, key=edge_key,
        relation_type=relation_type,
        metadata={
            "detection_method": "soft_link",
            "confidence": confidence,
            "note": note,
        },
        organization_id=org_id,
    )
    store.edges_by_type.setdefault(relation_type, []).append(
        (from_urn, to_urn, edge_key)
    )
    store._save_soft_links()

    from_label = _node_label(store.node_dict(from_urn))
    to_label = _node_label(store.node_dict(to_urn))
    return (
        f"Soft link created (id={link_id}):\n"
        f"  {from_label} --[CODE_TO_DATA]--> {to_label}\n"
        f"  confidence={confidence}, note={note!r}"
    )


# ── Tool 14: List Soft Links ─────────────────────────────────────────


@mcp.tool()
def list_soft_links() -> str:
    """List all manually-created soft links with their IDs, node labels,
    confidence scores, and notes."""
    if not store.soft_links:
        return "No soft links defined."

    lines = [f"Soft links ({len(store.soft_links)}):"]
    for link in store.soft_links:
        from_node = store.node_dict(link["from_urn"])
        to_node = store.node_dict(link["to_urn"])
        from_label = _node_label(from_node) if from_node else link["from_urn"]
        to_label = _node_label(to_node) if to_node else link["to_urn"]
        lines.append(f"\n  id: {link['id']}")
        lines.append(f"  from: [{from_node['node_type'] if from_node else '?'}] {from_label}")
        lines.append(f"        {link['from_urn']}")
        lines.append(f"  to:   [{to_node['node_type'] if to_node else '?'}] {to_label}")
        lines.append(f"        {link['to_urn']}")
        lines.append(f"  confidence: {link.get('confidence', '?')}")
        lines.append(f"  note: {link.get('note', '')}")
        lines.append(f"  created_at: {link.get('created_at', '?')}")

    return "\n".join(lines)


# ── Tool 15: Remove Soft Link ────────────────────────────────────────


@mcp.tool()
def remove_soft_link(soft_link_id: str) -> str:
    """Remove a soft link by its UUID. The edge is removed from the
    in-memory graph and the change is persisted to soft_links.json.

    Args:
        soft_link_id: UUID of the soft link to remove.
    """
    target_link = None
    for link in store.soft_links:
        if link["id"] == soft_link_id:
            target_link = link
            break

    if target_link is None:
        return f"Error: no soft link found with id={soft_link_id}"

    from_urn = target_link["from_urn"]
    to_urn = target_link["to_urn"]
    relation_type = target_link.get("relation_type", "CODE_TO_DATA")
    edge_key = str(uuid.uuid5(
        EDGE_NAMESPACE, f"{from_urn}:{to_urn}:{relation_type}"
    ))

    if store.G.has_edge(from_urn, to_urn, key=edge_key):
        store.G.remove_edge(from_urn, to_urn, key=edge_key)

    edge_tuple = (from_urn, to_urn, edge_key)
    if relation_type in store.edges_by_type:
        try:
            store.edges_by_type[relation_type].remove(edge_tuple)
        except ValueError:
            pass

    store.soft_links.remove(target_link)
    store._save_soft_links()

    from_node = store.node_dict(from_urn)
    to_node = store.node_dict(to_urn)
    from_label = _node_label(from_node) if from_node else from_urn
    to_label = _node_label(to_node) if to_node else to_urn
    return f"Soft link removed (id={soft_link_id}): {from_label} → {to_label}"


# ── Entry point ──────────────────────────────────────────────────────

if __name__ == "__main__":
    mcp.run(transport="stdio")
