"""Shared UMLS graph backend (FalkorDB + subgraph builder). Used by app.py."""

import os
import threading
import time
from collections import deque
from itertools import combinations
from pathlib import Path

import pandas as pd
from falkordb import FalkorDB

from tui_definitions import format_tui_list, tui_codes, tui_display

# ─────────────────────────────────────────────────────────────────────────────
# FalkorDB connection
# ─────────────────────────────────────────────────────────────────────────────
FALKORDB_HOST = os.environ.get("FALKORDB_HOST", "localhost")
FALKORDB_PORT = int(os.environ.get("FALKORDB_PORT", "6379"))
FALKORDB_GRAPH = os.environ.get("FALKORDB_GRAPH", "clinical_graph")

DATA_DIR = Path(__file__).resolve().parent
NODES_PATH = DATA_DIR / "graph_nodes.parquet"

# Safety caps (keeps the browser responsive on dense concepts)
MAX_SAME_CUI_AUIS = 14
MAX_MRREL_EDGES = 800
MRREL_CAP_MIN_EDGES = 2_000  # only cap when the subgraph has a very large MRREL set
NODE_WARN_THRESHOLD = 2500
NODE_TRIM_TARGET = 2000
QUERY_TIMEOUT_MS = 120_000
MAX_BFS_AUIS = 15_000

# FalkorDB silently truncates any query result to RESULTSET_SIZE rows (default
# 10,000). Fetching metadata/edges for the whole BFS frontier (up to MAX_BFS_AUIS)
# therefore drops rows — including the seed's own neighbours — which disconnects
# the graph. We never display more than ~NODE_TRIM_TARGET nodes, so we pick the
# lowest-hop AUIs up to MAX_DISPLAY_AUIS *before* fetching, and chunk every
# id-keyed query below FETCH_CHUNK so no single query can be truncated.
MAX_DISPLAY_AUIS = NODE_TRIM_TARGET
FETCH_CHUNK = 5_000

# Path finder (AUI ↔ AUI via HIER_ISA + MRREL, undirected)
MAX_PATH_ENUM = 50
MAX_PATH_SEARCH_HOPS = 25
PATH_EXTRA_HOPS = 4
MAX_PATH_BFS_VISITED = 12_000
MAX_PATH_BFS_FRONTIER = 3_000
MAX_PATH_ADJ_NODES = 8_000

# FalkorDB schema: Concept nodes (aui, cui, str, sab, tui, tty); HIER_ISA + MRREL edges
REL_HIER = "HIER_ISA"
REL_MRREL = "MRREL"
NODE_LABEL = "Concept"


def _truncate(text: str, n: int) -> str:
    return (text[: n - 1] + "…") if len(text) > n else text


def _aui_node_label(aui: str, s: str) -> str:
    short = _truncate(s, 36)
    return f"{aui}\n{short}"


def _cui_node_label(cui: str, _s: str = "") -> str:
    return cui


def _prop_str(val) -> str:
    if val is None:
        return ""
    return str(val).strip()


def _edge_props(props) -> dict:
    if not props:
        return {}
    if isinstance(props, dict):
        return props
    return dict(props)


def _clamp_depth(depth: int) -> int:
    return max(1, min(15, int(depth or 1)))


class _LockedGraph:
    """Serialize FalkorDB access (Dash may invoke callbacks concurrently)."""

    def __init__(self, graph):
        self._graph = graph
        self._lock = threading.Lock()

    def query(
        self, cypher: str, params: dict | None = None, *, timeout: int | None = None
    ):
        with self._lock:
            return self._graph.query(cypher, params=params or {}, timeout=timeout)


print(
    f"▶ Connecting to FalkorDB ({FALKORDB_HOST}:{FALKORDB_PORT}, graph={FALKORDB_GRAPH}) …",
    flush=True,
)
_t = time.time()
_fdb = FalkorDB(host=FALKORDB_HOST, port=FALKORDB_PORT)
GRAPH = _LockedGraph(_fdb.select_graph(FALKORDB_GRAPH))

# Smoke test
_probe = GRAPH.query(
    f"MATCH (n:{NODE_LABEL}) RETURN count(n) AS c LIMIT 1", timeout=QUERY_TIMEOUT_MS
)
_node_count = _probe.result_set[0][0] if _probe.result_set else "?"
print(f"  {FALKORDB_GRAPH}: {_node_count:,} Concept nodes  ({time.time() - _t:.1f}s)")

# Warm index/cache so the first user explore is less likely to time out.
_warm = GRAPH.query(
    f"MATCH (n:{NODE_LABEL}) WHERE n.aui IS NOT NULL RETURN n.aui LIMIT 1",
    timeout=QUERY_TIMEOUT_MS,
)
if _warm.result_set:
    print("  index warm-up ok", flush=True)

# Category lives in parquet only (not stored in FalkorDB Concept nodes).
_CATEGORY_BY_AUI: dict[str, str] = {}


def _load_category_lookup() -> None:
    global _CATEGORY_BY_AUI
    if not NODES_PATH.exists():
        print("  ⚠ graph_nodes.parquet missing — category filter disabled.", flush=True)
        _CATEGORY_BY_AUI = {}
        return
    t0 = time.time()
    df = pd.read_parquet(NODES_PATH, columns=["node_id", "category"])
    _CATEGORY_BY_AUI = {
        str(row.node_id).strip(): str(row.category).strip()
        for row in df.itertuples(index=False)
        if row.node_id
    }
    print(
        f"  category lookup: {len(_CATEGORY_BY_AUI):,} AUIs from parquet ({time.time() - t0:.1f}s)",
        flush=True,
    )


def _category_for_aui(aui: str) -> str:
    return _CATEGORY_BY_AUI.get(aui, "")


_load_category_lookup()
print("▶ FalkorDB ready.", flush=True)

# Legacy alias (app previously used DB for DuckDB)
DB = GRAPH


# ─────────────────────────────────────────────────────────────────────────────
# Cypher helpers
# ─────────────────────────────────────────────────────────────────────────────
def _chunked(seq: list[str], size: int = FETCH_CHUNK):
    for i in range(0, len(seq), size):
        yield seq[i : i + size]


def _fetch_concepts_by_auis(auis: list[str]) -> list[dict]:
    if not auis:
        return []
    out: list[dict] = []
    # Chunk so a single query never returns more than RESULTSET_SIZE rows
    # (one row per AUI), which FalkorDB would otherwise silently truncate.
    for chunk in _chunked(list(auis)):
        rows = GRAPH.query(
            f"""
            UNWIND $auis AS a
            MATCH (n:{NODE_LABEL} {{aui: a}})
            RETURN n.aui AS aui, n.cui AS cui, n.str AS str,
                   n.sab AS sab, n.tui AS tui, n.tty AS tty
            ORDER BY aui
            """,
            params={"auis": chunk},
            timeout=QUERY_TIMEOUT_MS,
        ).result_set
        for row in rows:
            out.append(
                {
                    "AUI": _prop_str(row[0]),
                    "CUI": _prop_str(row[1]),
                    "STR": _prop_str(row[2]),
                    "SAB": _prop_str(row[3]),
                    "TUI": _prop_str(row[4]),
                    "TTY": _prop_str(row[5]),
                }
            )
    return out


def _fetch_concepts_by_cuis(cuis: list[str]) -> list[dict]:
    if not cuis:
        return []
    out: list[dict] = []
    # One CUI can map to many AUIs, so chunk the input AND guard against the
    # RESULTSET cap by shrinking the chunk if a query comes back maxed out.
    for chunk in _chunked(list(cuis), max(1, FETCH_CHUNK // 4)):
        rows = GRAPH.query(
            f"""
            UNWIND $cuis AS c
            MATCH (n:{NODE_LABEL} {{cui: c}})
            RETURN n.cui AS cui, n.aui AS aui, n.str AS str, n.tui AS tui
            ORDER BY cui, aui
            """,
            params={"cuis": chunk},
            timeout=QUERY_TIMEOUT_MS,
        ).result_set
        out.extend(
            {
                "CUI": _prop_str(row[0]),
                "AUI": _prop_str(row[1]),
                "STR": _prop_str(row[2]),
                "TUI": _prop_str(row[3]),
            }
            for row in rows
        )
    return out


def _bfs_auis(
    seed_auis: list[str], depth: int
) -> tuple[list[str], dict[str, int], bool]:
    """
    Layered BFS over HIER_ISA + MRREL (undirected).

    Returns (all AUIs sorted, hop distance per AUI from seed, capped_flag).
    Layered expansion avoids FalkorDB's ~10k row cap on variable-length patterns
    and yields stable hop ranks for large-graph trimming. Metadata is fetched
    separately (after hop-budget selection) so we never pull more rows than we
    display, which would otherwise hit the RESULTSET_SIZE truncation cap.
    """
    if not seed_auis:
        return [], {}, False
    depth = _clamp_depth(depth)
    hop: dict[str, int] = {}
    for a in seed_auis:
        hop[a] = 0
    frontier = set(seed_auis)
    all_auis: set[str] = set(seed_auis)
    capped = False

    for d in range(1, depth + 1):
        if not frontier:
            break
        if len(all_auis) >= MAX_BFS_AUIS:
            capped = True
            break
        nbrs = _neighbors_one_hop_auis(frontier) - all_auis
        if len(all_auis) + len(nbrs) > MAX_BFS_AUIS:
            room = MAX_BFS_AUIS - len(all_auis)
            nbrs = set(sorted(nbrs)[:room])
            capped = True
        for a in nbrs:
            hop[a] = d
        all_auis |= nbrs
        frontier = nbrs
        if capped:
            break

    return sorted(all_auis), hop, capped


def _select_auis_by_hop(
    all_auis: list[str],
    aui_hop: dict[str, int],
    seed_auis: list[str],
    max_auis: int,
) -> list[str]:
    """
    Pick the AUIs closest to the seed (lowest hop) up to *max_auis*.

    Keeping the nearest hops guarantees the seed and its immediate neighbours
    survive, so the displayed graph stays connected. Selecting before fetching
    also keeps every downstream query well under the RESULTSET_SIZE cap.
    """
    if len(all_auis) <= max_auis:
        return list(all_auis)
    seed_set = set(seed_auis)
    ordered = sorted(
        all_auis,
        key=lambda a: (0 if a in seed_set else 1, aui_hop.get(a, 10**9), a),
    )
    return sorted(ordered[:max_auis])


def fetch_aui_filter_meta(aui_ids: list[str]) -> dict[str, dict]:
    """Return {AUI_<id>: {sab, tty, tuis}} for sidebar filters."""
    if not aui_ids:
        return {}
    rows = _fetch_concepts_by_auis(aui_ids)
    meta: dict[str, dict] = {}
    for r in rows:
        meta[f"AUI_{r['AUI']}"] = {
            "sab": r["SAB"],
            "tty": r["TTY"],
            "tuis": tui_codes(r["TUI"]),
            "category": _category_for_aui(r["AUI"]),
        }
    return meta


# ─────────────────────────────────────────────────────────────────────────────
# Subgraph post-processing (connectivity + size cap)
# ─────────────────────────────────────────────────────────────────────────────
def _split_elements(elements: list[dict]) -> tuple[list[dict], list[dict]]:
    nodes = [e for e in elements if "source" not in e["data"]]
    edges = [e for e in elements if "source" in e["data"]]
    return nodes, edges


def _build_adjacency(edges: list[dict]) -> dict[str, set[str]]:
    adj: dict[str, set[str]] = {}
    for el in edges:
        s, t = el["data"]["source"], el["data"]["target"]
        adj.setdefault(s, set()).add(t)
        adj.setdefault(t, set()).add(s)
    return adj


def _bfs_reachable(
    starts: set[str], adj: dict[str, set[str]], universe: set[str]
) -> set[str]:
    reachable: set[str] = set()
    q: deque[str] = deque()
    for s in starts:
        if s in universe and s not in reachable:
            reachable.add(s)
            q.append(s)
    while q:
        u = q.popleft()
        for v in adj.get(u, ()):
            if v in universe and v not in reachable:
                reachable.add(v)
                q.append(v)
    return reachable


def _bfs_distances(
    starts: set[str], adj: dict[str, set[str]], universe: set[str]
) -> dict[str, int]:
    dist: dict[str, int] = {}
    q: deque[str] = deque()
    for s in starts:
        if s in universe and s not in dist:
            dist[s] = 0
            q.append(s)
    while q:
        u = q.popleft()
        du = dist[u]
        for v in adj.get(u, ()):
            if v in universe and v not in dist:
                dist[v] = du + 1
                q.append(v)
    return dist


def _seed_node_ids(
    seed: str, seed_auis: list[str], seed_type: str, show_cui_layer: bool
) -> set[str]:
    if seed_type == "CUI" and show_cui_layer:
        return {f"CUI_{seed}"}
    if seed_auis:
        return {f"AUI_{seed_auis[0]}"}
    return set()


def _subset_for_kept_auis(
    nodes: list[dict], edges: list[dict], keep_aui_ids: set[str]
) -> tuple[list[dict], list[dict]]:
    """Keep selected AUI nodes and CUI nodes tied to those AUIs (plus seed CUI)."""
    keep_cui: set[str] = set()
    for n in nodes:
        d = n["data"]
        if d.get("node_type") == "AUI" and d["id"] in keep_aui_ids:
            cui = d.get("cui")
            if cui:
                keep_cui.add(f"CUI_{cui}")
        elif d.get("node_type") == "CUI" and d.get("is_seed"):
            keep_cui.add(d["id"])
    keep_ids = keep_aui_ids | keep_cui
    nodes_out = [n for n in nodes if n["data"]["id"] in keep_ids]
    edges_out = [
        e
        for e in edges
        if e["data"]["source"] in keep_ids and e["data"]["target"] in keep_ids
    ]
    return nodes_out, edges_out


def _trim_auis_by_hop_budget(
    nodes: list[dict], edges: list[dict], aui_hop: dict[str, int], target: int
) -> tuple[list[dict], list[dict], int, set[str]]:
    """
    Add AUIs in ascending hop order until the AUI+CUI subset would exceed *target*.
    Ensures a larger BFS depth never keeps fewer AUIs than a smaller depth.
    """
    aui_nodes = [n for n in nodes if n["data"].get("node_type") == "AUI"]
    aui_nodes.sort(
        key=lambda n: (aui_hop.get(n["data"].get("aui"), 10**9), n["data"]["id"])
    )
    kept_aui: set[str] = set()
    best_nodes, best_edges = [], []
    max_hop_kept = 0
    for n in aui_nodes:
        trial_aui = kept_aui | {n["data"]["id"]}
        trial_nodes, trial_edges = _subset_for_kept_auis(nodes, edges, trial_aui)
        if len(trial_nodes) > target:
            break
        kept_aui = trial_aui
        best_nodes, best_edges = trial_nodes, trial_edges
        max_hop_kept = max(max_hop_kept, aui_hop.get(n["data"].get("aui"), 0))
    if not kept_aui and aui_nodes:
        seed_n = aui_nodes[0]
        best_nodes, best_edges = _subset_for_kept_auis(
            nodes, edges, {seed_n["data"]["id"]}
        )
        max_hop_kept = aui_hop.get(seed_n["data"].get("aui"), 0)
    return best_nodes, best_edges, max_hop_kept, kept_aui


def _apply_graph_limits(
    elements: list[dict],
    seed: str,
    seed_auis: list[str],
    seed_type: str,
    show_cui_layer: bool,
    aui_hop: dict[str, int] | None = None,
) -> tuple[list[dict], list[str]]:
    notes: list[str] = []
    pinned_auis: set[str] = set()
    nodes, edges = _split_elements(elements)
    all_ids = {n["data"]["id"] for n in nodes}
    if not all_ids:
        return elements, notes

    starts = _seed_node_ids(seed, seed_auis, seed_type, show_cui_layer) & all_ids
    if not starts and seed_auis:
        starts = {f"AUI_{seed_auis[0]}"}

    n_nodes = len(nodes)
    if n_nodes > NODE_WARN_THRESHOLD and aui_hop:
        best_nodes, best_edges, max_hop_kept, pinned_aui = _trim_auis_by_hop_budget(
            nodes, edges, aui_hop, NODE_TRIM_TARGET
        )
        pinned_auis = pinned_aui
        trimmed = n_nodes - len(best_nodes)
        n_aui_kept = sum(1 for n in best_nodes if n["data"].get("node_type") == "AUI")
        notes.append(
            f"⚠ {n_nodes:,} nodes exceeds {NODE_WARN_THRESHOLD:,} — trimmed to "
            f"{len(best_nodes):,} ({n_aui_kept:,} AUIs, farthest hop kept: {max_hop_kept}, "
            f"removed {trimmed:,})"
        )
        nodes, edges = best_nodes, best_edges
        all_ids = {n["data"]["id"] for n in nodes}
        n_nodes = len(nodes)
    elif n_nodes > NODE_WARN_THRESHOLD:
        adj = _build_adjacency(edges)
        dist = _bfs_distances(starts, adj, all_ids)
        order_remove = sorted(
            nodes,
            key=lambda n: (-dist.get(n["data"]["id"], 10**9), n["data"]["id"]),
        )
        keep_ids = {n["data"]["id"] for n in nodes}
        for n in order_remove:
            if len(keep_ids) <= NODE_TRIM_TARGET:
                break
            keep_ids.discard(n["data"]["id"])

        trimmed = n_nodes - len(keep_ids)
        notes.append(
            f"⚠ {n_nodes:,} nodes exceeds {NODE_WARN_THRESHOLD:,} — trimmed to "
            f"{len(keep_ids):,} (removed {trimmed:,} farthest from seed)"
        )
        nodes = [n for n in nodes if n["data"]["id"] in keep_ids]
        edges = [
            e
            for e in edges
            if e["data"]["source"] in keep_ids and e["data"]["target"] in keep_ids
        ]
        all_ids = keep_ids
        adj = _build_adjacency(edges)
        reachable2 = _bfs_reachable(starts, adj, all_ids)
        if len(reachable2) < len(all_ids):
            extra = len(all_ids) - len(reachable2)
            notes.append(
                f"removed {extra:,} node{'s' if extra != 1 else ''} after trim (no path to seed)"
            )
            nodes = [n for n in nodes if n["data"]["id"] in reachable2]
            edges = [
                e
                for e in edges
                if e["data"]["source"] in reachable2
                and e["data"]["target"] in reachable2
            ]

    all_ids = {n["data"]["id"] for n in nodes}
    adj = _build_adjacency(edges)
    reachable = _bfs_reachable(starts, adj, all_ids)
    if pinned_auis:
        reachable |= pinned_auis
    dropped = len(all_ids) - len(reachable)
    if dropped:
        notes.append(
            f"removed {dropped:,} disconnected node{'s' if dropped != 1 else ''}"
        )
        nodes = [n for n in nodes if n["data"]["id"] in reachable]
        edges = [
            e
            for e in edges
            if e["data"]["source"] in reachable and e["data"]["target"] in reachable
        ]

    return nodes + edges, notes


def _seed_starts_from_nodes(nodes: list[dict]) -> set[str]:
    cui_seeds = sorted(
        n["data"]["id"]
        for n in nodes
        if n["data"].get("is_seed") and n["data"].get("node_type") == "CUI"
    )
    if cui_seeds:
        return {cui_seeds[0]}
    aui_seeds = sorted(
        n["data"]["id"]
        for n in nodes
        if n["data"].get("is_seed") and n["data"].get("node_type") == "AUI"
    )
    if aui_seeds:
        return {aui_seeds[0]}
    legacy = sorted(
        n["data"]["id"] for n in nodes if "seed-node" in (n.get("classes") or "")
    )
    return {legacy[0]} if legacy else set()


def filter_to_seed_component(
    nodes: list[dict], edges: list[dict]
) -> tuple[list[dict], list[dict]]:
    node_ids = {n["data"]["id"] for n in nodes}
    if not node_ids:
        return nodes, edges

    starts = _seed_starts_from_nodes(nodes) & node_ids
    if not starts:
        aui_seeds = sorted(
            n["data"]["id"]
            for n in nodes
            if n["data"].get("is_seed") and n["data"].get("node_type") == "AUI"
        )
        if aui_seeds:
            starts = {aui_seeds[0]}
        else:
            return nodes, edges

    adj = _build_adjacency(edges)
    reachable = _bfs_reachable(starts, adj, node_ids)
    nodes_out = [n for n in nodes if n["data"]["id"] in reachable]
    edges_out = [
        e
        for e in edges
        if e["data"]["source"] in reachable and e["data"]["target"] in reachable
    ]
    return nodes_out, edges_out


def _neighbors_one_hop_auis(frontier: set[str], *, directed: bool = False) -> set[str]:
    if not frontier:
        return set()
    frontier_list = list(frontier)
    if directed:
        rows = GRAPH.query(
            f"""
            MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(m:{NODE_LABEL})
            WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
            RETURN DISTINCT m.aui AS aui
            """,
            params={"frontier": frontier_list},
            timeout=QUERY_TIMEOUT_MS,
        ).result_set
    else:
        rows = GRAPH.query(
            f"""
            MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]-(m:{NODE_LABEL})
            WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
            RETURN DISTINCT m.aui AS aui
            """,
            params={"frontier": frontier_list},
            timeout=QUERY_TIMEOUT_MS,
        ).result_set
    return {_prop_str(row[0]) for row in rows if row[0]}


def _neighbors_one_hop_predecessors(frontier: set[str]) -> set[str]:
    if not frontier:
        return set()
    rows = GRAPH.query(
        f"""
        MATCH (m:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(n:{NODE_LABEL})
        WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
        RETURN DISTINCT m.aui AS aui
        """,
        params={"frontier": list(frontier)},
    ).result_set
    return {_prop_str(row[0]) for row in rows if row[0]}


def _neighbors_undirected_edges(frontier: set[str]) -> list[tuple[str, str]]:
    if not frontier:
        return []
    rows = GRAPH.query(
        f"""
        MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]-(m:{NODE_LABEL})
        WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
        RETURN DISTINCT n.aui AS u, m.aui AS v
        """,
        params={"frontier": list(frontier)},
    ).result_set
    out: list[tuple[str, str]] = []
    for row in rows:
        u, v = _prop_str(row[0]), _prop_str(row[1])
        if u and v:
            out.append((u, v) if u <= v else (v, u))
    return out


def _directed_edges_forward(frontier: set[str]) -> list[tuple[str, str]]:
    if not frontier:
        return []
    rows = GRAPH.query(
        f"""
        MATCH (n:{NODE_LABEL})-[:{REL_HIER}|{REL_MRREL}]->(m:{NODE_LABEL})
        WHERE n.aui IN $frontier AND m.aui IS NOT NULL AND m.aui <> n.aui
        RETURN DISTINCT n.aui AS u, m.aui AS v
        """,
        params={"frontier": list(frontier)},
    ).result_set
    return [(_prop_str(row[0]), _prop_str(row[1])) for row in rows if row[0] and row[1]]


def _aui_exists(aui: str) -> bool:
    hit = GRAPH.query(
        f"MATCH (n:{NODE_LABEL} {{aui: $a}}) RETURN n.aui LIMIT 1",
        params={"a": aui},
    ).result_set
    return bool(hit)


def _resolve_seed(seed: str) -> tuple[list[str], str] | tuple[None, None]:
    """Return (seed_auis, seed_type) or (None, None) if not found."""
    hit = GRAPH.query(
        f"MATCH (n:{NODE_LABEL} {{aui: $s}}) RETURN n.aui LIMIT 1",
        params={"s": seed},
    ).result_set
    if hit:
        return [seed], "AUI"

    rows = GRAPH.query(
        f"MATCH (n:{NODE_LABEL} {{cui: $s}}) RETURN DISTINCT n.aui AS aui ORDER BY aui",
        params={"s": seed},
    ).result_set
    if rows:
        return [_prop_str(row[0]) for row in rows if row[0]], "CUI"
    return None, None


def _shortest_hop_distance(
    aui_start: str, aui_end: str, max_hops: int, *, directed: bool = False
) -> int | None:
    if aui_start == aui_end:
        return 0
    dist_fwd: dict[str, int] = {aui_start: 0}
    dist_bwd: dict[str, int] = {aui_end: 0}
    front_fwd: set[str] = {aui_start}
    front_bwd: set[str] = {aui_end}
    for step in range(1, max_hops + 1):
        if len(dist_fwd) + len(dist_bwd) > MAX_PATH_BFS_VISITED:
            return None
        meet = set(dist_fwd) & set(dist_bwd)
        if meet:
            return min(dist_fwd[m] + dist_bwd[m] for m in meet)
        if len(front_fwd) <= len(front_bwd):
            if not front_fwd:
                break
            nbrs = _neighbors_one_hop_auis(front_fwd, directed=directed)
            next_fwd: set[str] = set()
            for v in nbrs:
                if v not in dist_fwd:
                    dist_fwd[v] = step
                    next_fwd.add(v)
                    if len(dist_fwd) > MAX_PATH_BFS_VISITED:
                        return None
            if len(next_fwd) > MAX_PATH_BFS_FRONTIER:
                return None
            front_fwd = next_fwd
        else:
            if not front_bwd:
                break
            nbrs = (
                _neighbors_one_hop_predecessors(front_bwd)
                if directed
                else _neighbors_one_hop_auis(front_bwd)
            )
            next_bwd: set[str] = set()
            for v in nbrs:
                if v not in dist_bwd:
                    dist_bwd[v] = step
                    next_bwd.add(v)
                    if len(dist_bwd) > MAX_PATH_BFS_VISITED:
                        return None
            if len(next_bwd) > MAX_PATH_BFS_FRONTIER:
                return None
            front_bwd = next_bwd
    return None


def _build_local_adjacency(
    seed_auis: set[str], max_hops: int, *, directed: bool = False
) -> tuple[set[str], dict[str, set[str]]]:
    visited: set[str] = set(seed_auis)
    adj: dict[str, set[str]] = {a: set() for a in seed_auis}
    frontier = set(seed_auis)
    for _ in range(max_hops):
        if not frontier or len(visited) >= MAX_PATH_BFS_VISITED:
            break
        new_frontier: set[str] = set()
        if directed:
            for u, v in _directed_edges_forward(frontier):
                if len(visited) >= MAX_PATH_BFS_VISITED:
                    break
                adj.setdefault(u, set()).add(v)
                adj.setdefault(v, set())
                if v not in visited:
                    visited.add(v)
                    new_frontier.add(v)
        else:
            for u, v in _neighbors_undirected_edges(frontier):
                if len(visited) >= MAX_PATH_BFS_VISITED:
                    break
                adj.setdefault(u, set()).add(v)
                adj.setdefault(v, set()).add(u)
                for n in (u, v):
                    if n not in visited:
                        visited.add(n)
                        adj.setdefault(n, set())
                        new_frontier.add(n)
                        if len(visited) >= MAX_PATH_BFS_VISITED:
                            break
        if (
            len(visited) >= MAX_PATH_BFS_VISITED
            or len(new_frontier) > MAX_PATH_BFS_FRONTIER
        ):
            break
        frontier = new_frontier
    return visited, adj


def _enumerate_simple_paths(
    aui_start: str,
    aui_end: str,
    adj: dict[str, set[str]],
    max_depth: int,
    max_paths: int,
) -> list[list[str]]:
    paths: list[list[str]] = []

    def dfs(curr: str, visited: set[str], trail: list[str]) -> None:
        if len(paths) >= max_paths:
            return
        if curr == aui_end:
            paths.append(trail[:])
            return
        if len(trail) >= max_depth:
            return
        for nxt in sorted(adj.get(curr, ())):
            if nxt not in visited:
                visited.add(nxt)
                trail.append(nxt)
                dfs(nxt, visited, trail)
                trail.pop()
                visited.discard(nxt)
                if len(paths) >= max_paths:
                    return

    dfs(aui_start, {aui_start}, [aui_start])
    return paths


def find_paths_between_auis(
    aui_start: str,
    aui_end: str,
    *,
    max_paths: int = MAX_PATH_ENUM,
    max_hops: int = MAX_PATH_SEARCH_HOPS,
    directed: bool = False,
) -> tuple[list[list[str]], str]:
    aui_start = aui_start.strip().upper()
    aui_end = aui_end.strip().upper()
    mode = "directed" if directed else "undirected"
    if not _aui_exists(aui_start):
        return [], f"AUI '{aui_start}' not found."
    if not _aui_exists(aui_end):
        return [], f"AUI '{aui_end}' not found."
    if aui_start == aui_end:
        return [[aui_start]], f"Same start and end AUI (trivial path, {mode})."

    shortest = _shortest_hop_distance(aui_start, aui_end, max_hops, directed=directed)
    if shortest is None:
        return (
            [],
            f"No {mode} path from {aui_start} to {aui_end} within {max_hops} hops "
            f"(HIER_ISA + MRREL).",
        )

    search_depth = min(shortest + PATH_EXTRA_HOPS, max_hops)
    _, adj = _build_local_adjacency({aui_start}, search_depth, directed=directed)
    if aui_end not in adj:
        _, adj = _build_local_adjacency(
            {aui_start, aui_end}, search_depth + 2, directed=directed
        )

    if len(adj) > MAX_PATH_ADJ_NODES:
        return (
            [],
            f"Path search neighborhood too large ({len(adj):,} AUIs near "
            f"{aui_start}). Pick AUIs farther from high-degree hubs, use directed "
            f"mode, or reduce hop depth.",
        )

    paths = _enumerate_simple_paths(aui_start, aui_end, adj, search_depth, max_paths)
    if not paths:
        return (
            [],
            f"Reachable in {mode} graph (shortest ≈ {shortest} hops) but no simple path "
            f"found within {search_depth} hops (enumeration limit).",
        )

    capped = len(paths) >= max_paths
    hop_lens = [len(p) - 1 for p in paths]
    msg = (
        f"Found {len(paths)} {mode} path{'s' if len(paths) != 1 else ''} "
        f"({min(hop_lens)}–{max(hop_lens)} hops, shortest={shortest})"
    )
    if capped:
        msg += f" — showing first {max_paths} (cap)"
    if search_depth > shortest:
        msg += f" · searched up to {search_depth} hops"
    return paths, msg


def _path_nav_neighbor_map(paths: list[list[str]]) -> dict[str, set[str]]:
    nav: dict[str, set[str]] = {}
    for path in paths:
        for i, aui in enumerate(path):
            nid = f"AUI_{aui}"
            nav.setdefault(nid, set())
            if i > 0:
                nav[nid].add(f"AUI_{path[i - 1]}")
            if i < len(path) - 1:
                nav[nid].add(f"AUI_{path[i + 1]}")
    return nav


def _path_edge_pairs(
    paths: list[list[str]], *, directed: bool = False
) -> set[tuple[str, str]]:
    pairs: set[tuple[str, str]] = set()
    for path in paths:
        for i in range(len(path) - 1):
            a, b = path[i], path[i + 1]
            if directed:
                pairs.add((a, b))
            else:
                pairs.add((a, b) if a <= b else (b, a))
    return pairs


def _fetch_hier_edges(auis: list[str]) -> list[dict]:
    if not auis:
        return []
    rows = GRAPH.query(
        f"""
        MATCH (x:{NODE_LABEL})-[r:{REL_HIER}]->(y:{NODE_LABEL})
        WHERE x.aui IN $auis AND y.aui IN $auis
        RETURN x.aui AS u, y.aui AS v, r.sab AS sab
        """,
        params={"auis": auis},
    ).result_set
    return [
        {
            "u": _prop_str(row[0]),
            "v": _prop_str(row[1]),
            "sab": _prop_str(row[2]),
            "rela": "isa",
        }
        for row in rows
    ]


def _fetch_mrrel_edges(auis: list[str], limit: int | None = None) -> list[dict]:
    if not auis:
        return []
    lim = f"LIMIT {int(limit)}" if limit else ""
    rows = GRAPH.query(
        f"""
        MATCH (x:{NODE_LABEL})-[r:{REL_MRREL}]->(y:{NODE_LABEL})
        WHERE x.aui IN $auis AND y.aui IN $auis AND x.aui <> y.aui
        RETURN x.aui AS u, y.aui AS v, r.sab AS sab, r.rela AS rela, r.rel AS rel
        {lim}
        """,
        params={"auis": auis},
    ).result_set
    return [
        {
            "u": _prop_str(row[0]),
            "v": _prop_str(row[1]),
            "sab": _prop_str(row[2]),
            "rela": _prop_str(row[3]) or _prop_str(row[4]) or "rel",
            "rel": _prop_str(row[4]),
        }
        for row in rows
    ]


def _fetch_path_edges(
    pairs: set[tuple[str, str]], *, directed: bool = False
) -> list[dict]:
    if not pairs:
        return []
    pair_list = [{"u": u, "v": v} for u, v in sorted(pairs)]
    if directed:
        hier_q = f"""
            UNWIND $pairs AS p
            MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_HIER}]->(y:{NODE_LABEL} {{aui: p.v}})
            RETURN 'mrhier' AS et, x.aui AS u, y.aui AS v, r.sab AS sab, 'isa' AS rela, '' AS rel
        """
        mrrel_q = f"""
            UNWIND $pairs AS p
            MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_MRREL}]->(y:{NODE_LABEL} {{aui: p.v}})
            RETURN 'mrrel' AS et, x.aui AS u, y.aui AS v, r.sab AS sab,
                   coalesce(r.rela, r.rel, 'rel') AS rela, coalesce(r.rel, '') AS rel
        """
    else:
        hier_q = f"""
            UNWIND $pairs AS p
            MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_HIER}]-(y:{NODE_LABEL} {{aui: p.v}})
            RETURN 'mrhier' AS et, x.aui AS u, y.aui AS v, r.sab AS sab, 'isa' AS rela, '' AS rel
        """
        mrrel_q = f"""
            UNWIND $pairs AS p
            MATCH (x:{NODE_LABEL} {{aui: p.u}})-[r:{REL_MRREL}]-(y:{NODE_LABEL} {{aui: p.v}})
            RETURN 'mrrel' AS et, x.aui AS u, y.aui AS v, r.sab AS sab,
                   coalesce(r.rela, r.rel, 'rel') AS rela, coalesce(r.rel, '') AS rel
        """
    rows = GRAPH.query(hier_q, params={"pairs": pair_list}).result_set
    rows += GRAPH.query(mrrel_q, params={"pairs": pair_list}).result_set
    out: list[dict] = []
    for row in rows:
        out.append(
            {
                "et": _prop_str(row[0]),
                "u": _prop_str(row[1]),
                "v": _prop_str(row[2]),
                "sab": _prop_str(row[3]),
                "rela": _prop_str(row[4]),
                "rel": _prop_str(row[5]),
            }
        )
    return out


def build_path_subgraph(
    aui_start: str,
    aui_end: str,
    paths: list[list[str]],
    *,
    directed: bool = False,
) -> tuple[list[dict], str]:
    if not paths:
        return [], "No paths to display."

    aui_start = aui_start.strip().upper()
    aui_end = aui_end.strip().upper()
    all_auis: set[str] = set()
    for p in paths:
        all_auis.update(p)

    meta_rows = _fetch_concepts_by_auis(list(all_auis))
    meta_by_aui = {r["AUI"]: r for r in meta_rows}

    elements: list[dict] = []
    aui_node_ids: set[str] = set()
    path_pairs = _path_edge_pairs(paths, directed=directed)

    for aui in sorted(all_auis):
        r = meta_by_aui.get(aui, {})
        nid = f"AUI_{aui}"
        aui_node_ids.add(nid)
        is_start = aui == aui_start
        is_end = aui == aui_end
        classes = "aui-node"
        if is_start:
            classes += " path-endpoint-start"
        elif is_end:
            classes += " path-endpoint-end"
        else:
            classes += " path-node"
        if is_start or is_end:
            classes += " seed-node"

        elements.append(
            {
                "data": {
                    "id": nid,
                    "label": _aui_node_label(aui, r.get("STR", "") or ""),
                    "node_type": "AUI",
                    "aui": aui,
                    "cui": r.get("CUI", "") or "",
                    "str": r.get("STR", "") or "",
                    "sab": r.get("SAB", "") or "",
                    "tui": r.get("TUI", "") or "",
                    "tui_label": format_tui_list(r.get("TUI", "") or ""),
                    "tty": r.get("TTY", "") or "",
                    "category": _category_for_aui(aui),
                    "is_seed": is_start or is_end,
                    "path_role": "start" if is_start else ("end" if is_end else "via"),
                },
                "classes": classes,
            }
        )

    nav_map = _path_nav_neighbor_map(paths)

    if not path_pairs:
        for el in elements:
            if "source" in el["data"]:
                continue
            nid = el["data"].get("id")
            if not nid:
                continue
            step_nbrs = sorted(nav_map.get(nid, ()))
            if step_nbrs:
                el["data"]["nav_neighbor_ids"] = step_nbrs
        n_edges = 0
        msg = (
            f"Path view: {aui_start} → {aui_end}  |  "
            f"{len(paths)} path(s)  |  {len(all_auis)} AUIs  |  {n_edges} edges"
        )
        return elements, msg

    edge_rows = _fetch_path_edges(path_pairs, directed=directed)
    seen_edge: set[tuple[str, str, str]] = set()
    for r in edge_rows:
        u, v = r["u"], r["v"]
        et = r["et"]
        key = (et, u, v)
        if key in seen_edge:
            continue
        seen_edge.add(key)
        src, tgt = f"AUI_{u}", f"AUI_{v}"
        if src not in aui_node_ids or tgt not in aui_node_ids:
            continue
        if et == "mrhier":
            elements.append(
                {
                    "data": {
                        "id": f"mrhier_{u}_{v}",
                        "source": src,
                        "target": tgt,
                        "edge_type": "mrhier",
                        "label": r["rela"],
                        "sab": r["sab"],
                        "on_path": True,
                    },
                    "classes": "mrhier-edge path-edge",
                }
            )
        else:
            rela = str(r["rela"] or "rel")
            label = rela if len(rela) <= 22 else rela[:20] + "…"
            elements.append(
                {
                    "data": {
                        "id": f"mrrel_{u}_{v}_{r['rel']}_{rela}",
                        "source": src,
                        "target": tgt,
                        "edge_type": "mrrel",
                        "label": label,
                        "rela": rela,
                        "rel": r["rel"] or "",
                        "sab": r["sab"] or "",
                        "on_path": True,
                    },
                    "classes": "mrrel-edge path-edge",
                }
            )

    for el in elements:
        if "source" in el["data"]:
            continue
        nid = el["data"].get("id")
        if not nid:
            continue
        step_nbrs = sorted(nav_map.get(nid, ()))
        if step_nbrs:
            el["data"]["nav_neighbor_ids"] = step_nbrs

    n_edges = sum(1 for e in elements if "source" in e["data"])
    msg = (
        f"Path view: {aui_start} → {aui_end}  |  "
        f"{len(paths)} path(s)  |  {len(all_auis)} AUIs  |  {n_edges} edges"
    )
    return elements, msg


def build_subgraph(
    seed: str,
    depth: int,
    show_mrhier: bool,
    show_aui_cui: bool,
    show_same_cui: bool,
    show_cui_cui: bool,
    show_mrrel: bool,
    show_cui_layer: bool = True,
) -> tuple[list[dict], str]:
    del show_cui_cui  # CUI parent hierarchy not in this dataset

    seed = seed.strip().upper()
    depth = _clamp_depth(depth)

    resolved = _resolve_seed(seed)
    if resolved[0] is None:
        return [], f"'{seed}' not found as AUI or CUI."
    seed_auis, seed_type = resolved

    all_auis, aui_hop, bfs_capped = _bfs_auis(seed_auis, depth)
    if not all_auis:
        return [], "No metadata found for the resolved AUIs."

    # Select the display set (nearest hops) BEFORE fetching metadata/edges so we
    # never exceed FalkorDB's RESULTSET_SIZE cap and keep the seed connected.
    kept_auis = _select_auis_by_hop(all_auis, aui_hop, seed_auis, MAX_DISPLAY_AUIS)
    prefetch_trimmed = len(all_auis) - len(kept_auis)
    meta_rows = _fetch_concepts_by_auis(kept_auis)
    if not meta_rows:
        return [], "No metadata found for the resolved AUIs."

    seed_set: set[str] = set(seed_auis)

    elements: list[dict] = []
    aui_node_ids: set[str] = set()
    cui_node_ids: set[str] = set()
    aui_to_cui: dict[str, str] = {}
    visited_auis = [r["AUI"] for r in meta_rows]

    for r in meta_rows:
        aui = r["AUI"]
        cui = r["CUI"] or ""
        s = r["STR"] or ""
        sab = r["SAB"] or ""
        tui = r["TUI"] or ""
        tty = r["TTY"] or ""
        aui_to_cui[aui] = cui

        nid = f"AUI_{aui}"
        aui_node_ids.add(nid)
        is_seed = aui in seed_set
        tui_fmt = format_tui_list(tui)

        elements.append(
            {
                "data": {
                    "id": nid,
                    "label": _aui_node_label(aui, s),
                    "node_type": "AUI",
                    "aui": aui,
                    "cui": cui,
                    "str": s,
                    "sab": sab,
                    "tui": tui,
                    "tui_label": tui_fmt,
                    "tty": tty,
                    "category": _category_for_aui(aui),
                    "is_seed": is_seed,
                },
                "classes": "aui-node" + (" seed-node" if is_seed else ""),
            }
        )

    all_cuis: set[str] = {c for c in aui_to_cui.values() if c}
    cui_label: dict[str, tuple[str, str]] = {}
    if show_cui_layer and all_cuis:
        lbl_rows = _fetch_concepts_by_cuis(list(all_cuis))
        cui_buckets: dict[str, list[dict]] = {}
        for row in lbl_rows:
            cui_buckets.setdefault(row["CUI"], []).append(row)
        for cui, bucket in cui_buckets.items():
            pref = sorted(bucket, key=lambda x: x["AUI"])[0]
            cui_label[cui] = (pref["STR"] or "", pref["TUI"] or "")

        for cui in all_cuis:
            nid = f"CUI_{cui}"
            cui_node_ids.add(nid)
            raw, tui_raw = cui_label.get(cui, ("", ""))
            tui_fmt = format_tui_list(tui_raw)
            is_seed_cui = seed_type == "CUI" and cui == seed
            elements.append(
                {
                    "data": {
                        "id": nid,
                        "label": _cui_node_label(cui),
                        "node_type": "CUI",
                        "cui": cui,
                        "str": raw,
                        "tui": tui_raw,
                        "tui_label": tui_fmt,
                        "is_seed": is_seed_cui,
                    },
                    "classes": "cui-node" + (" seed-node" if is_seed_cui else ""),
                }
            )

    all_ids = aui_node_ids | cui_node_ids
    mrrel_truncated = False

    if show_mrhier:
        for r in _fetch_hier_edges(visited_auis):
            src = f"AUI_{r['u']}"
            tgt = f"AUI_{r['v']}"
            if src in all_ids and tgt in all_ids:
                elements.append(
                    {
                        "data": {
                            "id": f"mrhier_{r['u']}_{r['v']}",
                            "source": src,
                            "target": tgt,
                            "edge_type": "mrhier",
                            "label": r["rela"],
                            "sab": r["sab"],
                        },
                        "classes": "mrhier-edge",
                    }
                )

    if show_cui_layer and show_aui_cui:
        for aui, cui in aui_to_cui.items():
            if not cui:
                continue
            src = f"AUI_{aui}"
            tgt = f"CUI_{cui}"
            if src in all_ids and tgt in all_ids:
                elements.append(
                    {
                        "data": {
                            "id": f"aui_cui_{aui}",
                            "source": src,
                            "target": tgt,
                            "edge_type": "aui_cui",
                            "label": "belongs_to",
                        },
                        "classes": "aui-cui-edge",
                    }
                )

    if show_cui_layer and show_same_cui:
        cui_buckets: dict[str, list[str]] = {}
        for aui, cui in aui_to_cui.items():
            if cui:
                cui_buckets.setdefault(cui, []).append(aui)
        for _cui, auis in cui_buckets.items():
            if len(auis) > MAX_SAME_CUI_AUIS:
                continue
            if len(auis) > 1:
                for a, b in combinations(sorted(auis), 2):
                    elements.append(
                        {
                            "data": {
                                "id": f"same_cui_{a}_{b}",
                                "source": f"AUI_{a}",
                                "target": f"AUI_{b}",
                                "edge_type": "same_cui",
                                "label": "same_cui",
                            },
                            "classes": "same-cui-edge",
                        }
                    )

    if show_mrrel:
        re = _fetch_mrrel_edges(visited_auis, limit=None)
        if len(re) > MAX_MRREL_EDGES and len(re) >= MRREL_CAP_MIN_EDGES:
            mrrel_truncated = True
            re = sorted(
                re,
                key=lambda r: aui_hop.get(r["u"], 10**9) + aui_hop.get(r["v"], 10**9),
            )[:MAX_MRREL_EDGES]
        for r in re:
            src = f"AUI_{r['u']}"
            tgt = f"AUI_{r['v']}"
            if src in all_ids and tgt in all_ids:
                rela = str(r["rela"] or "rel")
                label = rela if len(rela) <= 22 else rela[:20] + "…"
                elements.append(
                    {
                        "data": {
                            "id": f"mrrel_{r['u']}_{r['v']}_{r['rel']}_{rela}",
                            "source": src,
                            "target": tgt,
                            "edge_type": "mrrel",
                            "label": label,
                            "rela": rela,
                            "rel": r["rel"] or "",
                            "sab": r["sab"] or "",
                        },
                        "classes": "mrrel-edge",
                    }
                )

    elements, limit_notes = _apply_graph_limits(
        elements,
        seed,
        seed_auis,
        seed_type,
        show_cui_layer,
        aui_hop=aui_hop,
    )
    # Connectivity is enforced inside _apply_graph_limits (hop-trimmed AUIs are pinned).

    n_aui = sum(1 for e in elements if e["data"].get("node_type") == "AUI")
    n_cui = sum(1 for e in elements if e["data"].get("node_type") == "CUI")
    n_edges = sum(1 for e in elements if "source" in e["data"])
    extra: list[str] = []
    if bfs_capped:
        extra.append(f"AUI neighborhood capped at {MAX_BFS_AUIS:,} atoms")
    if prefetch_trimmed > 0:
        extra.append(
            f"⚠ {len(all_auis):,} AUIs in range — kept {len(kept_auis):,} nearest "
            f"(hop-priority) before fetch; raise hop depth narrows scope"
        )
    if mrrel_truncated:
        extra.append(f"MRREL capped at {MAX_MRREL_EDGES} edges (seed-adjacent first)")
    extra.extend(limit_notes)
    msg = (
        f"Seed: {seed} ({seed_type})  |  Depth: {depth}  |  "
        f"AUI nodes: {n_aui}  |  CUI nodes: {n_cui}  |  Edges: {n_edges}"
    )
    if extra:
        msg += "  |  " + "  ·  ".join(extra)
    return elements, msg


def _extract_filter_options(elements: list[dict]) -> dict:
    auis: set[str] = set()
    cuis: set[str] = set()
    for e in elements:
        d = e["data"]
        if d.get("node_type") == "AUI" and d.get("aui"):
            auis.add(d["aui"])
        elif d.get("node_type") == "CUI" and d.get("cui"):
            cuis.add(d["cui"])

    sabs: set[str] = set()
    tuis: set[str] = set()
    ttys: set[str] = set()
    categories: set[str] = set()

    if auis:
        for r in _fetch_concepts_by_auis(list(auis)):
            if r["SAB"]:
                sabs.add(r["SAB"])
            if r["TTY"]:
                ttys.add(r["TTY"])
            tuis |= tui_codes(r["TUI"] or "")
            cat = _category_for_aui(r["AUI"])
            if cat:
                categories.add(cat)

    if cuis:
        for r in _fetch_concepts_by_cuis(list(cuis)):
            if r.get("SAB"):
                sabs.add(r["SAB"])
            tuis |= tui_codes(r.get("TUI") or "")

    return {
        "sab": [{"label": s, "value": s} for s in sorted(sabs)],
        "tui": [{"label": tui_display(t), "value": t} for t in sorted(tuis)],
        "tty": [{"label": t, "value": t} for t in sorted(ttys)],
        "category": [{"label": c, "value": c} for c in sorted(categories)],
    }


def _aui_passes_filters(
    node_id: str,
    meta: dict[str, dict],
    sab_keep: set[str],
    tui_keep: set[str],
    tty_keep: set[str],
    category_keep: set[str] | None = None,
) -> bool:
    m = meta.get(node_id)
    if not m:
        return True
    if sab_keep and m["sab"] not in sab_keep:
        return False
    if tty_keep and m["tty"] not in tty_keep:
        return False
    if tui_keep and not (m["tuis"] & tui_keep):
        return False
    if category_keep and m.get("category") not in category_keep:
        return False
    return True