#!/usr/bin/env python3
"""
Filtered graph query for graphify-out/graph.json.

Reduces noise by excluding test nodes and helper-only symbols.
Supports BFS query and shortest-path modes.

Usage:
  python scripts/graph_query.py "job lifecycle"
  python scripts/graph_query.py --path "job()" "_destroy_clone()"
  python scripts/graph_query.py "transport" --context calls --context implements
  python scripts/graph_query.py "lifecycle" --include-tests
  python scripts/graph_query.py "build" --depth 3 --budget 4000
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from collections import deque
from pathlib import Path

DEFAULT_EXCLUDE = [
    r"^tests/",
    r"/test_",
    r"_test\.py$",
    r"\.Tests\.",
    r"Test-",
    r"Measure-",
]

SIGNAL_RELATIONS = {
    "calls",
    "implements",
    "uses",
    "defines",
    "inherits",
    "shares_data_with",
    "references",
}


def load_graph(graph_path: Path) -> tuple[dict, dict]:
    data = json.loads(graph_path.read_text(encoding="utf-8"))
    nodes: dict[str, dict] = {n["id"]: n for n in data["nodes"]}
    adj: dict[str, list[tuple[str, str, float]]] = {nid: [] for nid in nodes}
    for link in data["links"]:
        src, tgt = link["source"], link["target"]
        rel = link.get("relation", "")
        cs = link.get("confidence_score", 1.0)
        if src in adj:
            adj[src].append((tgt, rel, cs))
        if tgt in adj:
            adj[tgt].append((src, rel, cs))
    return nodes, adj


def allowed_set(nodes: dict, exclude_patterns: list[str]) -> set[str]:
    allowed: set[str] = set()
    for nid, node in nodes.items():
        sf = node.get("source_file") or ""
        if not any(re.search(pat, sf) for pat in exclude_patterns):
            allowed.add(nid)
    return allowed


def find_starts(nodes: dict, query: str, allowed: set[str]) -> list[str]:
    q = query.lower()
    exact, partial = [], []
    for nid in allowed:
        label = (nodes[nid].get("label") or "").lower()
        sf = (nodes[nid].get("source_file") or "").lower()
        if q == label or q == nid:
            exact.append(nid)
        elif q in label or q in sf:
            partial.append(nid)
    return exact or partial


def bfs(
    nodes: dict,
    adj: dict,
    starts: list[str],
    allowed: set[str],
    max_depth: int,
    relations: set[str] | None,
) -> list[tuple[dict, int, str | None]]:
    visited: dict[str, str | None] = {}
    queue: deque[tuple[str, int, str | None]] = deque()
    for nid in starts:
        if nid in allowed:
            visited[nid] = None
            queue.append((nid, 0, None))

    results: list[tuple[dict, int, str | None]] = []
    while queue:
        nid, depth, via_rel = queue.popleft()
        results.append((nodes[nid], depth, via_rel))
        if depth < max_depth:
            for neighbor, rel, _cs in adj.get(nid, []):
                if neighbor not in visited and neighbor in allowed:
                    if relations is None or rel in relations:
                        visited[neighbor] = nid
                        queue.append((neighbor, depth + 1, rel))
    return results


def shortest_path(
    nodes: dict,
    adj: dict,
    a_ids: list[str],
    b_ids: list[str],
    allowed: set[str],
) -> list[tuple[str, str]]:
    b_set = set(b_ids)
    for start in a_ids:
        if start not in allowed:
            continue
        visited: dict[str, tuple[str, str] | None] = {start: None}
        queue: deque[str] = deque([start])
        while queue:
            cur = queue.popleft()
            if cur in b_set:
                path: list[tuple[str, str]] = []
                node: str | None = cur
                rel_used = ""
                while node is not None:
                    prev = visited[node]
                    if prev is None:
                        path.append((node, ""))
                    else:
                        prev_node, rel_used = prev
                        path.append((node, rel_used))
                        node = prev_node
                        continue
                    break
                return list(reversed(path))
            for neighbor, rel, _cs in adj.get(cur, []):
                if neighbor not in visited and neighbor in allowed:
                    visited[neighbor] = (cur, rel)
                    queue.append(neighbor)
    return []


def fmt_node(node: dict, depth: int, via_rel: str | None) -> str:
    indent = "  " * depth
    label = node.get("label", node["id"])
    sf = node.get("source_file", "")
    loc = node.get("source_location", "")
    loc_str = f" loc={loc}" if loc else ""
    rel_str = f" ←{via_rel}" if via_rel else ""
    return f"{indent}NODE {label} [src={sf}{loc_str}]{rel_str}"


def main() -> None:
    ap = argparse.ArgumentParser(description="Filtered graphify query")
    ap.add_argument("query", nargs="?", help="Search term for BFS query")
    ap.add_argument("--path", nargs=2, metavar=("A", "B"), help="Shortest path between two nodes")
    ap.add_argument("--context", action="append", metavar="REL", help="Limit traversal to these edge relations (repeatable)")
    ap.add_argument("--depth", type=int, default=2, help="BFS max depth (default 2)")
    ap.add_argument("--budget", type=int, default=3000, help="Approximate output char budget (default 3000)")
    ap.add_argument("--include-tests", action="store_true", help="Do not exclude test nodes")
    ap.add_argument("--graph", type=Path, default=Path("graphify-out/graph.json"), help="Path to graph.json")
    args = ap.parse_args()

    if not args.query and not args.path:
        ap.print_help()
        sys.exit(1)

    if not args.graph.exists():
        print(f"error: {args.graph} not found. Run /graphify first.", file=sys.stderr)
        sys.exit(1)

    nodes, adj = load_graph(args.graph)
    exclude = [] if args.include_tests else DEFAULT_EXCLUDE
    allowed = allowed_set(nodes, exclude)
    relations = set(args.context) if args.context else None

    total = len(nodes)
    filtered_out = total - len(allowed)
    print(f"Graph: {total} nodes, {filtered_out} excluded ({filtered_out*100//total}% noise filtered)\n")

    if args.path:
        a_ids = find_starts(nodes, args.path[0], allowed)
        b_ids = find_starts(nodes, args.path[1], allowed)
        if not a_ids:
            print(f"No match for '{args.path[0]}'", file=sys.stderr)
            sys.exit(1)
        if not b_ids:
            print(f"No match for '{args.path[1]}'", file=sys.stderr)
            sys.exit(1)
        path = shortest_path(nodes, adj, a_ids, b_ids, allowed)
        if not path:
            print(f"No path found between '{args.path[0]}' and '{args.path[1]}'")
        else:
            print(f"Path ({len(path)} hops):")
            for i, (nid, rel) in enumerate(path):
                node = nodes[nid]
                label = node.get("label", nid)
                sf = node.get("source_file", "")
                loc = node.get("source_location", "")
                loc_str = f":{loc}" if loc else ""
                arrow = f"  --{rel}--> " if rel else ""
                print(f"  {'  ' * i}{arrow}{label} [{sf}{loc_str}]")
        return

    starts = find_starts(nodes, args.query, allowed)
    if not starts:
        print(f"No matching nodes for '{args.query}' (excluding tests)")
        sys.exit(0)

    print(f"BFS from {len(starts)} start node(s) · depth={args.depth}" +
          (f" · context={','.join(sorted(relations))}" if relations else "") + "\n")

    results = bfs(nodes, adj, starts, allowed, args.depth, relations)

    chars = 0
    shown = 0
    for node, depth, via_rel in results:
        line = fmt_node(node, depth, via_rel)
        if chars + len(line) > args.budget:
            print(f"\n... ({len(results) - shown} more nodes, increase --budget to see)")
            break
        print(line)
        chars += len(line) + 1
        shown += 1


if __name__ == "__main__":
    main()