src/services/api/controllers/retrieve.py

"""
File: /retrieve.py
Created Date: Sunday October 26th 2025
Author: Christian Nonis <alch.infoemail@gmail.com>
-----
Last Modified: Thursday January 29th 2026 8:43:59 pm
Modified By: Christian Nonis <alch.infoemail@gmail.com>
-----
"""

import asyncio
import json
import re
from collections import defaultdict
from datetime import datetime
from typing import Any, Optional

from fastapi import HTTPException
from starlette.responses import JSONResponse

from src.constants.embeddings import Vector
from src.constants.kg import IdentificationParams, Node, Predicate
from src.core.search.entities import search_entities
from src.core.search.fact_filter import (
    filter_relevant_facts,
    personalized_pagerank,
    reciprocal_rank_fusion,
)
from src.core.search.hybrid import (
    collect_bm25_passages,
    collect_dense_passages,
    collect_ilike_passages,
    fuse_passage_lists,
)
from src.core.search.relationships import search_relationships
from src.lib.tracing.profiler import profile_request, profile_stage
from src.utils.vector_search import VectorSearchFacade
from src.services.api.constants.requests import (
    GetContextRequestBody,
    GetContextResponse,
    GetContextTriple,
    RetrieveRequestResponse,
    RetrieveNeighborsRequestResponse,
    RetrievedNeighborNode,
)
from src.services.kg_agent.main import graph_adapter, kg_agent
from src.services.data.main import data_adapter
from src.services.kg_agent.main import embeddings_adapter, vector_store_adapter
from src.utils.dates import parse_date_string, to_naive_utc
from src.utils.similarity.vectors import cosine_similarity
from src.utils.nlp.ner import _entity_extractor
from src.config import config

vector_search = VectorSearchFacade(vector_store_adapter)

_MAX_NOUN_CHUNKS = 5
_SEED_K = 25
_PASSAGE_K = 24
_PPR_DAMPING = 0.85
_PPR_ITERS = 20
_SESSION_RE = re.compile(r"session_(\d+)", re.IGNORECASE)
_BATCH_TURN_RE = re.compile(r"\bb(\d+)_t(\d+)\b", re.IGNORECASE)
_HISTORY_MODE_RE = re.compile(
    r"\b("
    r"order(?:ing|ed)?|chronolog|sequence|walk me through|"
    r"contradict|conflict(?:ing)?|inconsisten|"
    r"previously|originally|before (?:the |that )?(?:update|change|extension)|"
    r"first (?:said|stated|mentioned|sprint|deadline)|"
    r"how many (?:weeks|days)|"
    r"between .+ and |"
    r"have i|did i"
    r")\b",
    re.I,
)
_CURRENT_TRUTH_RE = re.compile(
    r"\b("
    r"now|currently|latest|updated (?:to|value)|after the (?:update|change)|"
    r"how many commits|average response time|response time of"
    r")\b",
    re.I,
)
_SPINE_ACTOR = ("MADE", "INITIATED", "PERFORMED", "EXPERIENCED", "COVERED")
_SPINE_TARGET = ("TARGETED", "AFFECTED", "RESULTED")
_CONTEXT_REL = ("OCCURRED", "WITHIN")
_DIVERSIFY_LAMBDA = 0.65
_BRIDGE_SCORE_PENALTY = 0.15
_BRIDGE_SEED_HUB_CAP = 12
_BRIDGE_RESERVE_HUB_CAP = 6
_BRIDGE_RESERVE_MIN = 3
_BRIDGE_RESERVE_FRAC_DENOM = 4
_DISTANCE_QUANT_DECIMALS = 3


def _quantize_distance(distance: float) -> float:
    try:
        value = float(distance)
    except (TypeError, ValueError):
        return float("inf")
    if value != value or value == float("inf") or value == float("-inf"):
        return value
    return round(value, _DISTANCE_QUANT_DECIMALS)


def _candidate_event_uuid(candidate: dict[str, Any]) -> str:
    return _event_hub_id(candidate)


def _best_hub_fact_texts(candidates: list[dict[str, Any]]) -> dict[str, str]:
    best: dict[str, tuple[float, str]] = {}
    for candidate in candidates:
        hub = _candidate_event_uuid(candidate)
        if not hub:
            continue
        text = str(candidate.get("text") or "").strip()
        if not text:
            continue
        score = _candidate_distance(candidate)
        prev = best.get(hub)
        if prev is None or score < prev[0] or (score == prev[0] and text < prev[1]):
            best[hub] = (score, text)
    return {hub: text for hub, (_score, text) in best.items()}


def _paths_for_curated(
    paths: list[dict[str, Any]],
    curated: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Keep bridge paths whose hubs both survived curation; attach fact legs."""
    hub_texts = _best_hub_fact_texts(curated)
    if not hub_texts or not paths:
        return []
    curated_hubs = set(hub_texts)
    out: list[dict[str, Any]] = []
    for path in paths:
        hubs = [str(h) for h in (path.get("hubs") or []) if h]
        if len(hubs) < 2 or not all(hub in curated_hubs for hub in hubs):
            continue
        out.append(
            {
                "hubs": hubs,
                "shared_entity": path.get("shared_entity"),
                "shared_entity_name": path.get("shared_entity_name"),
                "weight": path.get("weight"),
                "legs": [hub_texts[hub] for hub in hubs],
            }
        )
    out.sort(
        key=lambda p: (
            tuple(p["hubs"]),
            str(p.get("shared_entity") or ""),
        )
    )
    return out


def _candidate_distance(candidate: dict[str, Any]) -> float:
    raw = candidate.get("score")
    if raw is None:
        return float("inf")
    try:
        return _quantize_distance(float(raw))
    except (TypeError, ValueError):
        return float("inf")


def _rank_bridge_seed_hubs(
    candidates: list[dict[str, Any]],
    *,
    cap: int,
) -> list[str]:
    hub_best_score: dict[str, float] = {}
    for candidate in candidates:
        hub = _candidate_event_uuid(candidate)
        if not hub:
            continue
        score = _candidate_distance(candidate)
        prev = hub_best_score.get(hub)
        if prev is None or score < prev:
            hub_best_score[hub] = score
    return sorted(
        hub_best_score.keys(),
        key=lambda hub: (hub_best_score[hub], hub),
    )[: max(0, cap)]


def _resolve_chunk_sessions(
    chunk_ids: list[str], brain_id: str
) -> dict[str, list[str]]:
    unique: list[str] = []
    seen: set[str] = set()
    for chunk_id in chunk_ids:
        cid = str(chunk_id or "").strip()
        if cid and cid not in seen:
            seen.add(cid)
            unique.append(cid)
    if not unique:
        return {}
    out: dict[str, list[str]] = {}
    try:
        chunks, _ = data_adapter.get_text_chunks_by_ids(
            unique[:120], False, brain_id
        )
    except Exception:
        return {}
    for chunk in chunks:
        chunk_id = str(getattr(chunk, "id", "") or "")
        body = (getattr(chunk, "text", None) or "").strip()
        if not chunk_id or not body:
            continue
        sessions = _session_ids_from_text(body)
        if sessions:
            out[chunk_id] = sessions
    return out


def _sessions_from_chunk_map(
    chunk_ids: list[str], chunk_sessions: dict[str, list[str]]
) -> list[str]:
    found: list[str] = []
    seen: set[str] = set()
    for chunk_id in chunk_ids:
        for sid in chunk_sessions.get(str(chunk_id), []) or []:
            if sid not in seen:
                seen.add(sid)
                found.append(sid)
    return found


def _expand_cross_event_bridges(
    candidates: list[dict[str, Any]],
    brain_id: str,
    *,
    max_per_hub: int,
    include_history: bool = False,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    from src.core.saving.hub_bridges import (
        HubBridge,
        select_bridge_neighbors,
    )

    if max_per_hub <= 0 or not candidates:
        return candidates, []

    hub_best_score: dict[str, float] = {}
    seed_sessions: set[str] = set()
    seed_chunk_ids: list[str] = []
    for candidate in candidates:
        hub = _candidate_event_uuid(candidate)
        if hub:
            score = _candidate_distance(candidate)
            prev = hub_best_score.get(hub)
            if prev is None or score < prev:
                hub_best_score[hub] = score
        for sid in _candidate_session_ids(candidate):
            seed_sessions.add(sid)
        seed_chunk_ids.extend(str(c) for c in (candidate.get("chunk_ids") or []) if c)

    ranked_seeds = _rank_bridge_seed_hubs(
        candidates, cap=_BRIDGE_SEED_HUB_CAP
    )
    if not ranked_seeds:
        return candidates, []
    reserve_hubs = set(
        ranked_seeds[: min(_BRIDGE_RESERVE_HUB_CAP, len(ranked_seeds))]
    )

    try:
        bridges = graph_adapter.get_hub_bridges(ranked_seeds, brain_id=brain_id)
    except Exception:
        return candidates, []
    if not bridges:
        return candidates, []

    typed_bridges: list[HubBridge] = []
    for item in bridges:
        if isinstance(item, HubBridge):
            typed_bridges.append(item)
        else:
            typed_bridges.append(
                HubBridge(
                    event_a=str(item.get("event_a") or ""),
                    event_b=str(item.get("event_b") or ""),
                    shared_entity=str(item.get("shared_entity") or ""),
                    shared_entity_name=str(item.get("shared_entity_name") or ""),
                    weight=float(item.get("weight") or 1.0),
                )
            )

    expansions = select_bridge_neighbors(
        ranked_seeds,
        typed_bridges,
        max_per_hub=max_per_hub,
    )
    if not expansions:
        return candidates, []

    selected_neighbors = {neighbor for neighbor, _ in expansions}
    neighbor_uuids = sorted(selected_neighbors)
    with profile_stage(
        "graph.bridge_neighbors", neighbors=len(neighbor_uuids)
    ) as detail:
        neighbors = graph_adapter.get_event_hub_facts(
            neighbor_uuids, brain_id=brain_id
        )
        detail["rows"] = len(neighbors)

    neighbor_chunk_ids: list[str] = []
    facts_by_hub: dict[str, list[tuple[Any, Any, Any, Any, Any]]] = {}
    for n, r, m, r2, b in neighbors:
        event_uuid = str(getattr(m, "uuid", "") or "")
        if not event_uuid:
            continue
        facts_by_hub.setdefault(event_uuid, []).append((n, r, m, r2, b))
        neighbor_chunk_ids.extend(_extract_source_chunk_ids(n, r, m, r2, b))

    chunk_sessions = _resolve_chunk_sessions(
        seed_chunk_ids + neighbor_chunk_ids, brain_id
    )
    if not seed_sessions:
        seed_sessions.update(_sessions_from_chunk_map(seed_chunk_ids, chunk_sessions))

    hub_sessions: dict[str, set[str]] = {}
    for event_uuid, rows in facts_by_hub.items():
        sessions: set[str] = set()
        for n, r, m, r2, b in rows:
            chunk_ids = _extract_source_chunk_ids(n, r, m, r2, b)
            sessions.update(_sessions_from_chunk_map(chunk_ids, chunk_sessions))
        hub_sessions[event_uuid] = sessions

    bridge_by_neighbor = {neighbor: bridge for neighbor, bridge in expansions}
    seed_by_neighbor: dict[str, str] = {}
    for neighbor, bridge in expansions:
        if bridge.event_a == neighbor:
            seed_by_neighbor[neighbor] = bridge.event_b
        else:
            seed_by_neighbor[neighbor] = bridge.event_a

    path_meta: list[dict[str, Any]] = []
    for neighbor, bridge in expansions:
        seed = seed_by_neighbor.get(neighbor, "")
        path_meta.append(
            {
                "hubs": sorted([seed, neighbor]),
                "shared_entity": bridge.shared_entity,
                "shared_entity_name": bridge.shared_entity_name,
                "weight": bridge.weight,
            }
        )
    path_meta.sort(
        key=lambda p: (
            tuple(p["hubs"]),
            p["shared_entity"],
        )
    )

    existing_keys = {
        candidate.get("key")
        for candidate in candidates
        if candidate.get("key") is not None
    }
    expanded = list(candidates)
    for event_uuid in sorted(selected_neighbors):
        bridge = bridge_by_neighbor.get(event_uuid)
        if bridge is None:
            continue
        seed = seed_by_neighbor.get(event_uuid, "")
        base = hub_best_score.get(seed, float("inf"))
        novel = hub_sessions.get(event_uuid, set()) - seed_sessions
        score = (
            base + _BRIDGE_SCORE_PENALTY
            if base != float("inf")
            else float("inf")
        )
        if novel and score != float("inf"):
            score = max(0.0, score - min(0.05, 0.02 * len(novel)))
        reserve_ok = seed in reserve_hubs
        for n, r, m, r2, b in facts_by_hub.get(event_uuid, []):
            if not _fact_predicates_allowed(r, r2, include_history=include_history):
                continue
            key = _fact_key(r, r2)
            if key in existing_keys:
                continue
            existing_keys.add(key)
            chunk_ids = _extract_source_chunk_ids(n, r, m, r2, b)
            session_ids = _sessions_from_chunk_map(chunk_ids, chunk_sessions)
            fact_text = _format_event_fact(n, r, m, r2, b)
            if include_history and not (
                _is_currently_valid(r) and _is_currently_valid(r2)
            ):
                fact_text = f"[superseded] {fact_text}"
            expanded.append(
                {
                    "identified_entity": f"bridge:{bridge.shared_entity}",
                    "triple": (n, r, m, r2, b),
                    "score": score,
                    "key": key,
                    "text": fact_text,
                    "chunk_ids": chunk_ids,
                    "session_ids": session_ids,
                    "bridge": {
                        "from_hub": seed,
                        "to_hub": event_uuid,
                        "shared_entity": bridge.shared_entity,
                        "shared_entity_name": bridge.shared_entity_name,
                        "reserve_ok": reserve_ok,
                    },
                }
            )
    return expanded, path_meta


async def retrieve_data(
    text: str, limit: int, preferred_entities: str, brain_id: str = "default"
) -> RetrieveRequestResponse:
    """
    Retrieve data from the knowledge graph and data store.
    """
    if preferred_entities:
        preferred_entities = [
            e.strip() for e in preferred_entities.split(",") if e.strip()
        ]
    else:
        preferred_entities = []

    def _get_data():
        text_embeddings = embeddings_adapter.embed_query(text)

        data_vectors = vector_search.search_data(
            text_embeddings.embeddings,
            brain_id=brain_id,
            k=limit,
        )
        triple_vectors = vector_search.search_triplets(
            text_embeddings.embeddings,
            brain_id=brain_id,
            k=limit,
        )

        search_result = data_adapter.search(text, brain_id)

        ts_text_chunks = search_result.text_chunks
        ts_observations = search_result.observations

        v_text_chunks, v_observations = data_adapter.get_text_chunks_by_ids(
            [dv.metadata.get("resource_id") for dv in data_vectors], True, brain_id
        )

        node_ids = [
            node_id
            for tv in triple_vectors
            for node_id in tv.metadata.get("node_ids", [])
        ]
        nodes = graph_adapter.get_nodes_by_uuid(
            uuids=node_ids,
            brain_id=brain_id,
            with_relationships=True,
            relationships_depth=1,
            relationships_type=[
                tv.metadata.get("predicate")
                for tv in triple_vectors
                if tv.metadata.get("predicate", None)
            ],
            preferred_labels=preferred_entities or [],
        )
        return ts_text_chunks, ts_observations, v_text_chunks, v_observations, nodes

    ts_text_chunks, ts_observations, v_text_chunks, v_observations, nodes = (
        await asyncio.to_thread(_get_data)
    )

    return RetrieveRequestResponse(
        data=[*ts_text_chunks, *v_text_chunks],
        observations=[*ts_observations, *v_observations],
        relationships=nodes,
    )


async def retrieve_neighbors(
    uuid: Optional[str] = None,
    look_for: Optional[str] = None,
    identification_params: Optional[IdentificationParams] = None,
    limit: int = 10,
    brain_id: str = "default",
) -> RetrieveNeighborsRequestResponse:
    """
    Retrieve neighboring nodes related to a specified main node.

    If `uuid` is provided it is used to locate the main node; otherwise `identification_params` is used. Optionally filters first-degree neighbors by semantic similarity to `look_for`, expands those matches to similar nodes, and returns a deduplicated list of neighbor nodes (up to `limit`) with their relationship and the most common matching similar node.

    Parameters:
        uuid (Optional[str]): UUID of the main node to retrieve neighbors for. If omitted, `identification_params` must be provided.
        look_for (Optional[str]): Text used to filter first-degree neighbors by embedding similarity before expanding to similar nodes.
        identification_params (Optional[IdentificationParams]): Identification parameters used to find the main node when `uuid` is not provided.
        limit (int): Maximum number of neighbor results to include in the response.
        brain_id (str): Identifier of the brain / dataset to query.

    Returns:
        RetrieveNeighborsRequestResponse: Object containing:
          - count: total number of unique neighbors found,
          - main_node: the resolved main Node,
          - neighbors: list of RetrievedNeighborNode objects (neighbor, relationship, most_common) limited to `limit`.

    Raises:
        HTTPException: 404 if the main node cannot be found.
    """

    async def _get_neighbors():

        # ---------------------------------------------------------
        # ================= GETTING THE MAIN NODE =================
        # ---------------------------------------------------------
        def _get_node() -> Node:
            node = None
            if uuid:
                node = graph_adapter.get_by_uuid(uuid, brain_id)
            elif identification_params:
                node = graph_adapter.get_by_identification_params(
                    identification_params,
                    brain_id=brain_id,
                    entity_types=identification_params.entity_types,
                )
            if not node:
                raise HTTPException(status_code=404, detail="Entity not found")
            return node

        node = await asyncio.to_thread(_get_node)
        target_node_types = node.labels

        looking_for_v = embeddings_adapter.embed_query(look_for) if look_for else None

        # ---------------------------------------------------------
        # ===== Getting 1st degree neighbors of the main node =====
        # ---------------------------------------------------------
        def _get_fd_neighbors() -> (
            tuple[dict[str, list[tuple[Predicate, Node]]], list[str]]
        ):
            fd_neighbors = graph_adapter.get_neighbors(
                [node.uuid], limit=limit, brain_id=brain_id
            )
            fd_v_neighbors_ids = [
                fd[1].properties.get("v_id")
                for fd in fd_neighbors[node.uuid]
                if fd[1].properties.get("v_id") is not None
            ]
            if look_for:
                fd_v_neighbors_embeddings = vector_store_adapter.get_by_ids(
                    fd_v_neighbors_ids, brain_id=brain_id, store="nodes"
                )
                fd_v_neighbors_embeddings_map = {
                    v.id: v.embeddings
                    for v in fd_v_neighbors_embeddings
                    if (
                        cosine_similarity(looking_for_v.embeddings, v.embeddings) > 0.5
                        and v.id
                        and not v.id.replace(
                            "-", ""
                        ).isalpha()  # likely not a UUID if all numeric (may have hyphens for uuid standard)
                    )
                }
                fd_v_neighbors_ids = list(fd_v_neighbors_embeddings_map.keys())

            return fd_neighbors, fd_v_neighbors_ids

        fd_neighbors, fd_v_neighbors_ids = await asyncio.to_thread(_get_fd_neighbors)

        # ---------------------------------------------------------
        # === Getting nodes similar to the 1st degree neighbors ===
        # ---------------------------------------------------------
        fd_v_similar_node_futures = []
        for fd_v_neighbor_id in fd_v_neighbors_ids:
            fd_v_similar_node_futures.append(
                asyncio.to_thread(
                    vector_store_adapter.search_similar_by_ids,
                    [fd_v_neighbor_id],
                    brain_id,
                    "nodes",
                    0.5,
                    limit,
                )
            )
        fd_v_similar_nodes_results: list[dict[str, list[Vector]]] = (
            await asyncio.gather(*fd_v_similar_node_futures)
        )
        fd_similar_node_ids = [
            v.metadata.get("uuid")
            for result_dict in fd_v_similar_nodes_results
            for vectors in result_dict.values()
            for v in vectors
            if v.metadata.get("uuid") is not None
        ]
        fd_similar_nodes = await asyncio.to_thread(
            graph_adapter.get_by_uuids, fd_similar_node_ids, brain_id
        )
        fd_similar_nodes_by_uuid = {n.uuid: n for n in fd_similar_nodes}

        # ---------------------------------------------------------
        # === Getting neighbors of the 1st degree similar nodes ===
        # ---------------------------------------------------------
        def _get_fd_similar_node_neighbors() -> dict[str, list[tuple[Predicate, Node]]]:
            fd_similar_node_neighbors = graph_adapter.get_neighbors(
                fd_similar_node_ids,
                limit=limit,
                brain_id=brain_id,
                of_types=list(set(target_node_types)),
            )
            return fd_similar_node_neighbors

        fd_similar_node_neighbors = await asyncio.to_thread(
            _get_fd_similar_node_neighbors
        )

        seen_neighbor_uuids = set()
        unique_neighbors = []
        for source_uuid, neighbors_list in fd_similar_node_neighbors.items():
            for neighbor_tuple in neighbors_list:
                neighbor_uuid = neighbor_tuple[1].uuid
                if neighbor_uuid not in seen_neighbor_uuids:
                    seen_neighbor_uuids.add(neighbor_uuid)
                    unique_neighbors.append(
                        RetrievedNeighborNode(
                            neighbor=neighbor_tuple[1],
                            relationship=neighbor_tuple[0],
                            most_common=fd_similar_nodes_by_uuid.get(source_uuid),
                        )
                    )

        return RetrieveNeighborsRequestResponse(
            count=len(unique_neighbors),
            main_node=node,
            neighbors=unique_neighbors[:limit],
        )

    return await _get_neighbors()


async def retrieve_neighbors_ai_mode(
    identification_params: IdentificationParams,
    looking_for: Optional[list[str]],
    limit: int,
    brain_id: str = "default",
) -> RetrieveNeighborsRequestResponse:
    """
    Retrieve neighbors of an entity from the knowledge graph.
    """

    def _get_neighbors():
        node = graph_adapter.get_by_identification_params(
            identification_params,
            brain_id=brain_id,
            entity_types=identification_params.entity_types,
        )
        if not node:
            raise HTTPException(status_code=404, detail="Entity not found")

        result = kg_agent.retrieve_neighbors(node, looking_for, limit, brain_id)

        ids = [neighbor.uuid for neighbor in result.neighbors]
        descriptions = [neighbor.description for neighbor in result.neighbors]

        nodes = graph_adapter.get_nodes_by_uuid(uuids=ids, brain_id=brain_id)
        paired = list(zip(nodes, descriptions))

        return RetrieveNeighborsRequestResponse(neighbors=paired)

    result = await asyncio.to_thread(_get_neighbors)

    return result


async def get_relationships(
    limit: int = 10,
    skip: int = 0,
    relationship_types: Optional[list[str]] = None,
    from_node_labels: Optional[list[str]] = None,
    to_node_labels: Optional[list[str]] = None,
    query_text: Optional[str] = None,
    query_search_target: Optional[str] = "all",
    brain_id: str = "default",
):
    """
    Retrieve relationships from the knowledge graph with optional filtering and pagination.

    Parameters:
        relationship_types (list[str], optional): Filter results to specific relationship types.
        from_node_labels (list[str], optional): Filter relationships originating from nodes with these labels.
        to_node_labels (list[str], optional): Filter relationships targeting nodes with these labels.
        query_text (str, optional): Text to search within relationship or node content.
        query_search_target (str, optional): Field to target for text search; commonly "all", "from", or "to".
        limit (int, optional): Maximum number of relationships to return.
        skip (int, optional): Number of relationships to skip (offset).
        brain_id (str, optional): Identifier of the brain/graph to query.

    Returns:
        JSONResponse: A response whose JSON content contains:
            - message: Confirmation string.
            - relationships: List of serialized relationship objects.
            - total: Total number of matching relationships.
    """
    relationships = await asyncio.to_thread(
        search_relationships,
        limit,
        skip,
        relationship_types,
        from_node_labels,
        to_node_labels,
        query_text,
        query_search_target,
        brain_id,
    )

    return JSONResponse(
        content={
            "message": "Relationships retrieved successfully",
            "relationships": [r.model_dump(mode="json") for r in relationships.results],
            "total": relationships.total,
        }
    )


async def get_entities(
    limit: int = 10,
    skip: int = 0,
    node_labels: Optional[list[str]] = None,
    query_text: Optional[str] = None,
    brain_id: str = "default",
):
    """
    Retrieve entities from the knowledge graph with optional label and text filters.

    Parameters:
        limit (int): Maximum number of entities to return (pagination).
        skip (int): Number of entities to skip (pagination offset).
        node_labels (Optional[list[str]]): If provided, only return entities whose labels match any value in this list.
        query_text (Optional[str]): If provided, filter entities by matching text content.
        brain_id (str): Identifier of the knowledge graph/brain to query.

    Returns:
        JSONResponse: Object containing:
            - message (str): Informational message.
            - entities (list): Serialized entity objects.
            - total (int): Total number of matching entities.
    """
    entities = await asyncio.to_thread(
        search_entities, limit, skip, node_labels, query_text, brain_id
    )

    return JSONResponse(
        content={
            "message": "Entities retrieved successfully",
            "entities": [e.model_dump(mode="json") for e in entities.results],
            "total": entities.total,
        }
    )


def _format_kg_item(item: Any) -> str:
    name = getattr(item, "name", None) or ""
    description = getattr(item, "description", None) or ""
    happened_at = getattr(item, "happened_at", None) or ""
    base = f"{name}: {description}" if description else name
    if happened_at:
        return f"{base} @{happened_at}"
    return base


def _format_event_fact(n: Node, r: Predicate, m: Node, r2: Predicate, b: Node) -> str:
    r2_name = (getattr(r2, "name", None) or "").upper()
    end_label = (
        "Context"
        if any(token in r2_name for token in _CONTEXT_REL)
        else "Target"
    )
    return (
        f"Actor: {_format_kg_item(n)} | {_format_kg_item(r)} | "
        f"Event: {_format_kg_item(m)} | {_format_kg_item(r2)} | "
        f"{end_label}: {_format_kg_item(b)}"
    )


def _fact_key(r: Predicate, r2: Predicate) -> tuple[str, str]:
    return (getattr(r, "uuid", None) or "", getattr(r2, "uuid", None) or "")


def _rel_name(predicate: Predicate | None) -> str:
    return (getattr(predicate, "name", None) or "").upper()


def _is_spine_actor_rel(name: str) -> bool:
    return any(token in name for token in _SPINE_ACTOR)


def _is_spine_target_rel(name: str) -> bool:
    return any(token in name for token in _SPINE_TARGET)


def _is_context_rel(name: str) -> bool:
    return any(token in name for token in _CONTEXT_REL)


def _event_hub_id(candidate: dict[str, Any]) -> str:
    triple = candidate.get("triple") or ()
    r = triple[1] if len(triple) > 1 else None
    m = triple[2] if len(triple) > 2 else None
    r2 = triple[3] if len(triple) > 3 else None
    event_uuid = getattr(m, "uuid", None) if m is not None else None
    if event_uuid:
        return str(event_uuid)
    flow_key = (
        getattr(r, "flow_key", None)
        or getattr(r2, "flow_key", None)
        or ""
    )
    if flow_key:
        return str(flow_key)
    key = candidate.get("key")
    if isinstance(key, tuple):
        return "|".join(str(part or "") for part in key)
    return str(candidate.get("text") or "")


def _candidate_session_ids(candidate: dict[str, Any]) -> list[str]:
    sessions = candidate.get("session_ids")
    if isinstance(sessions, list) and sessions:
        return [str(s) for s in sessions if s]
    return []


def _hub_leg_kind(candidate: dict[str, Any]) -> str:
    triple = candidate.get("triple") or ()
    r2 = triple[3] if len(triple) > 3 else None
    name = _rel_name(r2)
    if _is_context_rel(name):
        return "context"
    if _is_spine_target_rel(name):
        return "spine"
    return "other"


def _hub_completeness_score(candidate: dict[str, Any]) -> float:
    triple = candidate.get("triple") or ()
    if len(triple) < 5:
        return 0.0
    n, r, m, r2, b = triple
    score = 0.0
    m_labels = {str(label).upper() for label in (getattr(m, "labels", None) or [])}
    if "EVENT" in m_labels:
        score += 1.0
    r_name = _rel_name(r)
    r2_name = _rel_name(r2)
    if _is_spine_actor_rel(r_name):
        score += 1.0
    elif r_name:
        score += 0.2
    if _is_spine_target_rel(r2_name):
        score += 1.0
    elif _is_context_rel(r2_name):
        score += 0.35
    elif r2_name:
        score += 0.15
    if getattr(n, "uuid", None) and getattr(b, "uuid", None):
        score += 0.5
    if getattr(m, "happened_at", None):
        score += 0.15
    return score


def _fact_recency_score(
    candidate: dict[str, Any], *, now: datetime | None = None
) -> float:
    triple = candidate.get("triple") or ()
    best = 0.0
    now = now or datetime.now(tz=None)
    for node in triple[0::2] if triple else ():
        raw = getattr(node, "happened_at", None)
        if not raw:
            continue
        parsed = parse_date_string(raw) if isinstance(raw, str) else raw
        if not isinstance(parsed, datetime):
            continue
        days_ago = max(0, (now - to_naive_utc(parsed)).days)
        score = 1.0 if days_ago <= 0 else 1.0 / (1.0 + days_ago)
        if score > best:
            best = score
    return best


def _relevance_score(
    candidate: dict[str, Any], *, now: datetime | None = None
) -> float:
    distance = _candidate_distance(candidate)
    distance_term = 1.0 / (1.0 + max(0.0, distance))
    ppr_term = float(candidate.get("ppr_mass") or 0.0)
    completeness = _hub_completeness_score(candidate) / 3.8
    recency = _fact_recency_score(candidate, now=now)
    topic_boost = 0.08 if candidate.get("topic_preferred") else 0.0
    return (
        0.45 * distance_term
        + 0.25 * ppr_term
        + 0.20 * completeness
        + 0.10 * recency
        + topic_boost
    )


def _completeness_rank_key(candidate: dict[str, Any]) -> tuple:
    return (
        -_relevance_score(candidate),
        _candidate_distance(candidate),
    ) + _candidate_tiebreak_key(candidate)


def _rank_facts_with_completeness(
    candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    return sorted(candidates, key=_completeness_rank_key)


def _is_complementary_hub_leg(
    selected_meta: list[dict[str, Any]], item: dict[str, Any]
) -> bool:
    if item["kind"] not in {"spine", "context"}:
        return False
    for prev in selected_meta:
        if prev["hub"] != item["hub"]:
            continue
        if {item["kind"], prev["kind"]} == {"spine", "context"}:
            return True
    return False


def _bridge_reserved_slots(max_facts: int) -> int:
    if max_facts <= 0:
        return 0
    return min(
        max_facts,
        max(_BRIDGE_RESERVE_MIN, max_facts // _BRIDGE_RESERVE_FRAC_DENOM),
    )


def _is_bridge_candidate(candidate: dict[str, Any]) -> bool:
    return bool(candidate.get("bridge"))


def _bridge_reserve_ok(candidate: dict[str, Any]) -> bool:
    bridge = candidate.get("bridge")
    if not bridge:
        return False
    if isinstance(bridge, dict) and "reserve_ok" in bridge:
        return bool(bridge.get("reserve_ok"))
    return True


def _prepare_diversify_item(
    candidate: dict[str, Any], *, now: datetime | None = None
) -> dict[str, Any]:
    key = candidate.get("key") or ("", "")
    is_bridge = _is_bridge_candidate(candidate)
    return {
        "candidate": candidate,
        "relevance": _relevance_score(candidate, now=now),
        "hub": _event_hub_id(candidate),
        "sessions": frozenset(_candidate_session_ids(candidate)),
        "kind": _hub_leg_kind(candidate),
        "tie": (str(key[0]) if key else "", str(key[1]) if len(key) > 1 else ""),
        "is_bridge": is_bridge,
        "reserve_ok": _bridge_reserve_ok(candidate) if is_bridge else False,
    }


def _mmr_pick_index(
    remaining: list[dict[str, Any]],
    selected_meta: list[dict[str, Any]],
    selected_hubs: set[str],
    selected_sessions: set[str],
) -> int:
    best_idx = 0
    best_mmr = None
    best_tie = None
    selected_kinds_by_hub: dict[str, set[str]] = defaultdict(set)
    for selected in selected_meta:
        selected_kinds_by_hub[selected["hub"]].add(selected["kind"])
    for idx, item in enumerate(remaining):
        relevance = item["relevance"]
        hub = item["hub"]
        sessions = item["sessions"]
        complementary = (
            item["kind"] in {"spine", "context"}
            and bool(
                selected_kinds_by_hub.get(item["hub"], set())
                & ({"spine", "context"} - {item["kind"]})
            )
        )
        if complementary:
            hub_pen = 0.0
            sess_pen = 0.0
            relevance = relevance + 0.5
        else:
            hub_pen = 1.0 if hub in selected_hubs else 0.0
            if sessions and selected_sessions:
                sess_pen = len(sessions & selected_sessions) / max(
                    1, len(sessions)
                )
            else:
                sess_pen = 0.0
        redundancy = 0.55 * hub_pen + 0.45 * sess_pen
        mmr = _DIVERSIFY_LAMBDA * relevance - (1.0 - _DIVERSIFY_LAMBDA) * redundancy
        tie = item["tie"]
        if (
            best_mmr is None
            or mmr > best_mmr
            or (mmr == best_mmr and (best_tie is None or tie < best_tie))
        ):
            best_mmr = mmr
            best_tie = tie
            best_idx = idx
    return best_idx


def _mmr_take(
    remaining: list[dict[str, Any]],
    *,
    budget: int,
    selected: list[dict[str, Any]],
    selected_meta: list[dict[str, Any]],
    selected_hubs: set[str],
    selected_sessions: set[str],
) -> None:
    while budget > 0 and remaining:
        best_idx = _mmr_pick_index(
            remaining, selected_meta, selected_hubs, selected_sessions
        )
        chosen = remaining.pop(best_idx)
        selected.append(chosen["candidate"])
        selected_meta.append(chosen)
        selected_hubs.add(chosen["hub"])
        selected_sessions.update(chosen["sessions"])
        budget -= 1


def _take_novel_bridge_slots(
    remaining: list[dict[str, Any]],
    *,
    budget: int,
    covered_sessions: set[str],
    selected: list[dict[str, Any]],
    selected_meta: list[dict[str, Any]],
    selected_hubs: set[str],
    selected_sessions: set[str],
) -> None:
    while budget > 0 and remaining:
        novel_sessions: set[str] = set()
        for item in remaining:
            if not item["is_bridge"] or not item["reserve_ok"]:
                continue
            novel_sessions.update(item["sessions"] - covered_sessions)
        if not novel_sessions:
            break
        target = min(novel_sessions)
        best_idx = None
        best_key = None
        for idx, item in enumerate(remaining):
            if not item["is_bridge"] or not item["reserve_ok"]:
                continue
            if target not in item["sessions"]:
                continue
            key = (item["hub"], item["tie"])
            if best_key is None or key < best_key:
                best_key = key
                best_idx = idx
        if best_idx is None:
            break
        chosen = remaining.pop(best_idx)
        selected.append(chosen["candidate"])
        selected_meta.append(chosen)
        selected_hubs.add(chosen["hub"])
        selected_sessions.update(chosen["sessions"])
        covered_sessions.update(chosen["sessions"])
        budget -= 1


def _diversify_facts(
    ranked: list[dict[str, Any]],
    *,
    max_facts: int,
) -> list[dict[str, Any]]:
    if max_facts <= 0 or not ranked:
        return []
    if len(ranked) <= max_facts:
        return list(ranked)

    ranking_now = datetime.now(tz=None)
    prepared = [
        _prepare_diversify_item(candidate, now=ranking_now) for candidate in ranked
    ]
    prepared.sort(key=lambda item: (-item["relevance"], item["tie"]))

    selected: list[dict[str, Any]] = []
    selected_meta: list[dict[str, Any]] = []
    selected_hubs: set[str] = set()
    selected_sessions: set[str] = set()

    reserve = _bridge_reserved_slots(max_facts)
    has_bridge = any(item["is_bridge"] and item["reserve_ok"] for item in prepared)
    if not has_bridge:
        reserve = 0
    main_budget = max_facts - reserve

    non_bridge = [item for item in prepared if not item["is_bridge"]]
    bridge_items = [item for item in prepared if item["is_bridge"]]
    stable_bridges = [item for item in bridge_items if item["reserve_ok"]]

    _mmr_take(
        non_bridge,
        budget=main_budget,
        selected=selected,
        selected_meta=selected_meta,
        selected_hubs=selected_hubs,
        selected_sessions=selected_sessions,
    )

    covered_by_non_bridge = set(selected_sessions)
    _take_novel_bridge_slots(
        stable_bridges,
        budget=reserve,
        covered_sessions=covered_by_non_bridge,
        selected=selected,
        selected_meta=selected_meta,
        selected_hubs=selected_hubs,
        selected_sessions=selected_sessions,
    )

    remaining = non_bridge + stable_bridges
    remaining.sort(key=lambda item: (-item["relevance"], item["tie"]))
    leftover = max_facts - len(selected)
    _mmr_take(
        remaining,
        budget=leftover,
        selected=selected,
        selected_meta=selected_meta,
        selected_hubs=selected_hubs,
        selected_sessions=selected_sessions,
    )
    return selected


def _temporal_conflict_meta(
    candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    by_entity: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for candidate in candidates:
        triple = candidate.get("triple") or ()
        for node in (triple[0], triple[4]) if len(triple) >= 5 else ():
            uuid = getattr(node, "uuid", None)
            if not uuid:
                continue
            happened = None
            event = triple[2] if len(triple) > 2 else None
            raw = getattr(event, "happened_at", None) if event is not None else None
            if isinstance(raw, str) and raw.strip():
                happened = raw.strip()
            by_entity[str(uuid)].append(
                {
                    "fact_key": list(candidate.get("key") or ()),
                    "hub_id": _event_hub_id(candidate),
                    "happened_at": happened,
                }
            )
    conflicts: list[dict[str, Any]] = []
    for entity_uuid, entries in sorted(by_entity.items()):
        times = {
            entry["happened_at"]
            for entry in entries
            if entry.get("happened_at")
        }
        hubs = {entry["hub_id"] for entry in entries}
        if len(times) >= 2 and len(hubs) >= 2:
            conflicts.append(
                {
                    "entity_uuid": entity_uuid,
                    "hubs": sorted(hubs),
                    "happened_ats": sorted(times),
                    "facts": [entry["fact_key"] for entry in entries],
                }
            )
    return conflicts


def _edge_type_weight(predicate: Predicate | None, query: str | None = None) -> float:
    name = _rel_name(predicate)
    if _is_spine_actor_rel(name) or _is_spine_target_rel(name):
        weight = 1.0
    elif _is_context_rel(name):
        weight = 0.35
    else:
        weight = 0.55
    if query:
        q_tokens = set(re.findall(r"[a-z0-9]+", query.lower()))
        p_tokens = set(re.findall(r"[a-z0-9]+", name.lower()))
        if q_tokens and p_tokens:
            overlap = len(q_tokens & p_tokens) / max(1, len(p_tokens))
            weight *= 1.0 + 0.5 * overlap
    return weight


def _add_weighted_edge(
    adjacency: dict[str, dict[str, float]],
    src: str | None,
    dst: str | None,
    weight: float,
) -> None:
    if not src or not dst or src == dst or weight <= 0:
        return
    bucket = adjacency.setdefault(str(src), {})
    prev = bucket.get(str(dst), 0.0)
    if weight > prev:
        bucket[str(dst)] = weight


def _candidate_tiebreak_key(candidate: dict[str, Any]) -> tuple[str, ...]:
    triple = candidate.get("triple") or ()
    n = triple[0] if len(triple) > 0 else None
    r = triple[1] if len(triple) > 1 else None
    m = triple[2] if len(triple) > 2 else None
    r2 = triple[3] if len(triple) > 3 else None
    b = triple[4] if len(triple) > 4 else None
    flow_key = (
        getattr(r, "flow_key", None)
        or getattr(r2, "flow_key", None)
        or ""
    )
    key = candidate.get("key")
    if not (isinstance(key, tuple) and len(key) == 2):
        key = _fact_key(r, r2) if r is not None and r2 is not None else ("", "")
    r_uuid, r2_uuid = key
    return (
        str(flow_key or ""),
        str(r_uuid or ""),
        str(r2_uuid or ""),
        str(getattr(n, "uuid", None) or ""),
        str(getattr(m, "uuid", None) or ""),
        str(getattr(b, "uuid", None) or ""),
        str(candidate.get("identified_entity") or ""),
        str(candidate.get("text") or ""),
    )


def _distance_rank_key(candidate: dict[str, Any]) -> tuple:
    return (_candidate_distance(candidate),) + _candidate_tiebreak_key(
        candidate
    )


def _ppr_rank_key(
    candidate: dict[str, Any], ppr_scores: dict[str, float]
) -> tuple:
    return (_score_candidate_with_ppr(candidate, ppr_scores),) + _candidate_tiebreak_key(
        candidate
    )


def _dedupe_candidates(candidates: list[dict[str, Any]]) -> list[dict[str, Any]]:
    deduped: dict[tuple[str, str], dict[str, Any]] = {}
    for candidate in candidates:
        key = candidate["key"]
        existing = deduped.get(key)
        if existing is None:
            deduped[key] = candidate
            continue
        if candidate["score"] < existing["score"]:
            deduped[key] = candidate
            continue
        if candidate["score"] == existing["score"] and _candidate_tiebreak_key(
            candidate
        ) < _candidate_tiebreak_key(existing):
            deduped[key] = candidate
    return sorted(deduped.values(), key=_distance_rank_key)


def _merge_variant_candidate_lists(
    variant_lists: list[list[dict[str, Any]]],
) -> list[dict[str, Any]]:
    merged: list[dict[str, Any]] = []
    for local in variant_lists:
        ordered = sorted(local, key=_distance_rank_key)
        merged.extend(ordered)
    return merged


def _is_currently_valid(predicate: Predicate) -> bool:
    props = getattr(predicate, "properties", None) or {}
    if props.get("invalid_at"):
        return False
    if getattr(predicate, "deprecated", False):
        return False
    return True


def _wants_historical_facts(query: str) -> bool:
    text = (query or "").strip()
    if not text:
        return False
    if _CURRENT_TRUTH_RE.search(text) and not _HISTORY_MODE_RE.search(text):
        return False
    return bool(_HISTORY_MODE_RE.search(text))


def _fact_predicates_allowed(
    r: Predicate,
    r2: Predicate,
    *,
    include_history: bool,
) -> bool:
    if include_history:
        return True
    return _is_currently_valid(r) and _is_currently_valid(r2)


def _collect_query_variants(text: str, elements) -> list[str]:
    seen: set[str] = set()
    variants: list[str] = []

    def _add(value: str | None) -> None:
        if not value:
            return
        cleaned = value.strip()
        if not cleaned:
            return
        key = cleaned.lower()
        if key in seen:
            return
        seen.add(key)
        variants.append(cleaned)

    _add(text)
    for token in getattr(elements, "tokens", None) or []:
        if isinstance(token, dict):
            _add(token.get("text"))
    noun_chunks = getattr(elements, "noun_chunks", None) or []
    added_chunks = 0
    for chunk in noun_chunks:
        if added_chunks >= _MAX_NOUN_CHUNKS:
            break
        before = len(variants)
        _add(chunk if isinstance(chunk, str) else None)
        if len(variants) > before:
            added_chunks += 1
    for extra in _decompose_enumeration_queries(text, elements):
        _add(extra)
    for extra in _ordering_aspect_queries(text):
        _add(extra)
    return variants


_ENUMERATION_RE = re.compile(
    r"\b(what|which|list|ways?|types?|kinds?|events?|books?|items?|"
    r"changes?|traits?|hobbies|activities)\b",
    re.IGNORECASE,
)
_ORDERING_ASPECT_QUERIES_BUDGET = (
    "core functionality authentication expense tracking visualization",
    "transaction CRUD error handling response management",
    "security deployment integration tests gunicorn workers",
)
_ORDERING_ASPECT_QUERIES_TRANSLATION = (
    "translation API integration error handling DeepL",
    "rate limiting request queue caching Redis",
    "language detection libraries franc evaluation",
    "database schema optimization contextual memory store",
)


def _is_enumeration_question(text: str) -> bool:
    return bool(_ENUMERATION_RE.search(text or ""))


def _is_ordering_question(text: str) -> bool:
    return bool(
        re.search(
            r"\b(order|ordered|ordering|chronolog|sequence|walk me through)\b",
            text or "",
            re.I,
        )
    )


def _ordering_aspect_queries(text: str) -> list[str]:
    if not _is_ordering_question(text):
        return []
    q = text or ""
    if re.search(
        r"\b(translation|language detection|deepl|franc|microservice|"
        r"multi-language|rate limit)\b",
        q,
        re.I,
    ):
        return list(_ORDERING_ASPECT_QUERIES_TRANSLATION)[:2]
    return list(_ORDERING_ASPECT_QUERIES_BUDGET)[:2]


def _decompose_enumeration_queries(text: str, elements) -> list[str]:
    if not _is_enumeration_question(text):
        return []
    queries: list[str] = []
    entities: list[str] = []
    for token in getattr(elements, "tokens", None) or []:
        if not isinstance(token, dict):
            continue
        label = str(token.get("label") or token.get("ent_type") or "").upper()
        value = str(token.get("text") or "").strip()
        if not value:
            continue
        if label in {"PERSON", "PER", "ORG", "GPE"} or (
            value[:1].isupper() and " " not in value
        ):
            entities.append(value)
    relation_terms = [
        m.group(0).lower()
        for m in re.finditer(
            r"\b(books?|read|ways?|types?|kinds?|events?|hobbies|"
            r"activities|changes?|traits?|friends?|places?)\b",
            text or "",
            re.IGNORECASE,
        )
    ]
    seen: set[str] = set()
    for entity in entities[:4]:
        for relation in relation_terms[:4]:
            q = f"{entity} {relation}"
            key = q.lower()
            if key not in seen and key != (text or "").strip().lower():
                seen.add(key)
                queries.append(q)
        key = entity.lower()
        if key not in seen:
            seen.add(key)
            queries.append(entity)
    for relation in relation_terms[:3]:
        if relation not in seen:
            seen.add(relation)
            queries.append(relation)
    return queries[:8]


def _seed_personalization(seeds: list[tuple[str, float, str]]) -> dict[str, float]:
    weights: dict[str, float] = {}
    for uuid, distance, _ in seeds:
        weight = 1.0 / (1.0 + max(0.0, _quantize_distance(distance)))
        prev = weights.get(uuid)
        if prev is None or weight > prev:
            weights[uuid] = weight
    return weights


def _build_adjacency_from_seeds(
    seed_uuids: list[str],
    brain_id: str,
    *,
    query: str | None = None,
    query_embedding: list[float] | None = None,
    include_history: bool = False,
) -> dict[str, list[tuple[str, float]]]:
    weighted: dict[str, dict[str, float]] = {}
    if not seed_uuids:
        return {}
    try:
        neighbors = graph_adapter.get_event_centric_neighbors(
            seed_uuids, brain_id=brain_id
        )
    except Exception:
        return {}
    for n, r, m, r2, b in neighbors:
        if not _fact_predicates_allowed(r, r2, include_history=include_history):
            continue
        n_uuid = getattr(n, "uuid", None)
        m_uuid = getattr(m, "uuid", None)
        b_uuid = getattr(b, "uuid", None)
        w1 = _edge_type_weight(r, query=query)
        w2 = _edge_type_weight(r2, query=query)
        if query_embedding is not None:
            for pred, base in ((r, w1), (r2, w2)):
                emb = getattr(pred, "embeddings", None) or getattr(
                    pred, "embedding", None
                )
                if emb is None:
                    continue
                try:
                    sim = float(cosine_similarity(query_embedding, emb))
                    boost = 0.5 + 0.5 * max(0.0, sim)
                    if pred is r:
                        w1 *= boost
                    else:
                        w2 *= boost
                except Exception:
                    pass
        _add_weighted_edge(weighted, n_uuid, m_uuid, w1)
        _add_weighted_edge(weighted, m_uuid, n_uuid, w1)
        _add_weighted_edge(weighted, m_uuid, b_uuid, w2)
        _add_weighted_edge(weighted, b_uuid, m_uuid, w2)
        if n_uuid and b_uuid:
            bridge = min(w1, w2) * 0.5
            _add_weighted_edge(weighted, n_uuid, b_uuid, bridge)
            _add_weighted_edge(weighted, b_uuid, n_uuid, bridge)
    return {
        src: sorted(((dst, weight) for dst, weight in edges.items()), key=lambda x: x[0])
        for src, edges in sorted(weighted.items(), key=lambda item: item[0])
    }


def _score_candidate_with_ppr(
    candidate: dict[str, Any], ppr_scores: dict[str, float]
) -> float:
    n, _r, m, _r2, b = candidate["triple"]
    scores = [
        ppr_scores.get(str(getattr(node, "uuid", "")), 0.0)
        for node in (n, m, b)
        if getattr(node, "uuid", None)
    ]
    if not scores:
        return _candidate_distance(candidate)
    return -max(scores)


def _context_looks_insufficient(
    question: str, text_lines: list[str], passages: list[str]
) -> bool:
    blob = "\n".join(text_lines + passages).lower()
    if len(blob.strip()) < 80:
        return True
    tokens = [
        t
        for t in re.findall(r"[a-z0-9']+", (question or "").lower())
        if len(t) > 3
        and t
        not in {
            "what",
            "which",
            "when",
            "where",
            "who",
            "does",
            "did",
            "have",
            "has",
            "been",
            "with",
            "from",
            "about",
            "would",
            "could",
            "their",
            "there",
            "this",
            "that",
            "mentioned",
        }
    ]
    if not tokens:
        return False
    hits = sum(1 for token in tokens if token in blob)
    return hits / len(tokens) < 0.35


def _stabilize_seed_hits(
    seeds: list[tuple[str, float, str]],
) -> list[tuple[str, float, str]]:
    best: dict[str, tuple[float, str]] = {}
    for uuid, distance, name in seeds:
        uid = str(uuid or "").strip()
        if not uid:
            continue
        dist = _quantize_distance(distance)
        prev = best.get(uid)
        if prev is None or dist < prev[0]:
            best[uid] = (dist, str(name))
    return sorted(
        ((uid, dist, name) for uid, (dist, name) in best.items()),
        key=lambda item: (item[1], item[0]),
    )


def _seed_nodes_for_text(text: str, brain_id: str) -> list[tuple[str, float, str]]:
    """Return (node_uuid, distance, entity_name) seeds from nodes + relationships."""
    with profile_stage("embed.query"):
        text_embeddings = embeddings_adapter.embed_query(text)
    embeddings = text_embeddings.embeddings
    seeds: list[tuple[str, float, str]] = []

    with profile_stage("vector.search_nodes", k=_SEED_K):
        node_vectors = vector_search.search_nodes(
            embeddings,
            brain_id=brain_id,
            k=_SEED_K,
        )
    for vector in node_vectors:
        meta = vector.metadata or {}
        uuid = meta.get("uuid") or getattr(vector, "id", None)
        if not uuid:
            continue
        distance = (
            vector.distance if vector.distance is not None else float("inf")
        )
        name = meta.get("name") or text
        seeds.append((str(uuid), float(distance), str(name)))

    with profile_stage("vector.search_relationships", k=_SEED_K):
        rel_vectors = vector_search.search_relationships(
            embeddings,
            brain_id=brain_id,
            k=_SEED_K,
        )
    for vector in rel_vectors:
        meta = vector.metadata or {}
        distance = (
            vector.distance if vector.distance is not None else float("inf")
        )
        node_ids = sorted(str(n) for n in (meta.get("node_ids") or []) if n)
        for node_id in node_ids:
            seeds.append((node_id, float(distance), text))

    return _stabilize_seed_hits(seeds)


def _extract_source_chunk_ids(*items: Any) -> list[str]:
    ids: list[str] = []
    seen: set[str] = set()
    for item in items:
        props = getattr(item, "properties", None) or {}
        for value in props.get("source_chunk_ids") or []:
            chunk_id = str(value).strip()
            if chunk_id and chunk_id not in seen:
                seen.add(chunk_id)
                ids.append(chunk_id)
        single = props.get("source_chunk_id")
        if single:
            chunk_id = str(single).strip()
            if chunk_id and chunk_id not in seen:
                seen.add(chunk_id)
                ids.append(chunk_id)
    return ids


def _session_ids_from_text(text: str) -> list[str]:
    found: list[str] = []
    seen: set[str] = set()
    for match in _SESSION_RE.finditer(text or ""):
        sid = f"session_{match.group(1)}"
        if sid not in seen:
            seen.add(sid)
            found.append(sid)
    for match in _BATCH_TURN_RE.finditer(text or ""):
        sid = f"session_b{match.group(1)}_t{match.group(2)}"
        if sid not in seen:
            seen.add(sid)
            found.append(sid)
    return found


def _session_ids_for_chunks(
    chunk_ids: list[str], chunk_sessions: dict[str, list[str]]
) -> list[str]:
    found: list[str] = []
    seen: set[str] = set()
    for chunk_id in chunk_ids:
        for sid in chunk_sessions.get(chunk_id) or []:
            if sid not in seen:
                seen.add(sid)
                found.append(sid)
    return found


def _build_fact_channel(
    curated: list[dict[str, Any]],
    chunk_sessions: dict[str, list[str]],
) -> tuple[list[str], list[GetContextTriple], list[str]]:
    text_lines: list[str] = []
    triples: list[GetContextTriple] = []
    graph_session_ids: list[str] = []
    seen_sessions: set[str] = set()
    for candidate in curated:
        n, r, m, r2, b = candidate["triple"]
        chunk_ids = list(candidate.get("chunk_ids") or [])
        session_ids = _candidate_session_ids(candidate)
        if not session_ids:
            session_ids = _session_ids_for_chunks(chunk_ids, chunk_sessions)
        line = candidate["text"]
        if session_ids:
            line = f"{line} ({', '.join(session_ids)})"
        text_lines.append(line)
        triples.append(
            GetContextTriple(
                identified_entity=candidate["identified_entity"],
                triple=(n, r, m, r2, b),
                source_chunk_ids=chunk_ids or None,
                source_session_ids=session_ids or None,
            )
        )
        for sid in session_ids:
            if sid not in seen_sessions:
                seen_sessions.add(sid)
                graph_session_ids.append(sid)
    graph_session_ids.sort()
    return text_lines, triples, graph_session_ids


def _retrieve_passages(
    text: str, brain_id: str, *, limit: int
) -> list[tuple[str, float, str]]:
    """Return (chunk_id, score, text) ranked passages via vector + keyword fusion."""
    mode = (
        config.context_passage_mode if config.search_enabled else "ilike"
    )
    use_dense = mode in ("hybrid", "dense", "ilike")
    use_bm25 = mode in ("hybrid", "bm25")
    use_ilike = mode == "ilike"
    with profile_stage("passages.retrieve", blocking=True, queries=1):
        vector_ids: list[str] = []
        id_to_distance: dict[str, float] = {}
        if use_dense:
            with profile_stage("embed.query"):
                text_embeddings = embeddings_adapter.embed_query(text)
            with profile_stage("vector.search_data", k=max(limit, _PASSAGE_K)):
                vector_ids, id_to_distance = collect_dense_passages(
                    vector_search,
                    text_embeddings.embeddings,
                    brain_id,
                    max(limit, _PASSAGE_K),
                )

        keyword_ids: list[str] = []
        id_to_text: dict[str, str] = {}
        if use_bm25:
            with profile_stage("data.bm25_search"):
                try:
                    keyword_ids, _, id_to_text = collect_bm25_passages(
                        data_adapter,
                        text,
                        brain_id,
                        max(limit, _PASSAGE_K),
                    )
                except Exception:
                    pass
        elif use_ilike:
            with profile_stage("data.keyword_search"):
                try:
                    keyword_ids, id_to_text = collect_ilike_passages(
                        data_adapter, text, brain_id
                    )
                except Exception:
                    pass

        fused = fuse_passage_lists(
            vector_ids if use_dense else [],
            keyword_ids,
        )
        ranked_ids = [item for item, _ in fused[:limit]]
        missing = [cid for cid in ranked_ids if cid not in id_to_text]
        if missing:
            with profile_stage("data.fetch_chunks", chunks=len(missing)):
                try:
                    chunks, _ = data_adapter.get_text_chunks_by_ids(
                        missing, False, brain_id
                    )
                    for chunk in chunks:
                        id_to_text[str(chunk.id)] = chunk.text or ""
                except Exception:
                    pass

        passages: list[tuple[str, float, str]] = []
        for chunk_id in ranked_ids:
            body = (id_to_text.get(chunk_id) or "").strip()
            if not body:
                continue
            score = id_to_distance.get(chunk_id, 1.0 / (1.0 + len(passages)))
            passages.append((chunk_id, float(score), body))
        return passages


async def get_context(
    request: GetContextRequestBody,
    *,
    fact_filter_adapter: Any = None,
) -> GetContextResponse:
    """
    Retrieve contextual information for a text.

    `fact_filter_adapter` is only supplied by deep-navigation callers: the
    one-shot context path leaves it unset because an LLM call does not fit the
    sub-second budget, and without it no fact filtering happens.

    `profile_stages` attaches a per-stage wall-clock breakdown to the response
    and publishes it as a latency trace event; it is off unless the caller asks
    for it or `TRACE_STAGE_PROFILER_ENABLED` is set.
    """
    with profile_request(
        "retrieve.context", enabled=request.profile_stages
    ) as profiler:
        response = await _build_context(request, fact_filter_adapter)
    if profiler is not None:
        response.stage_timings = profiler.last_report
    return response


async def _build_context(
    request: GetContextRequestBody,
    fact_filter_adapter: Any,
) -> GetContextResponse:
    with profile_stage("nlp.extract_elements"):
        elements = _entity_extractor.extract_elements(request.text)
    with profile_stage("nlp.query_variants") as detail:
        variants = _collect_query_variants(request.text, elements)
        detail["variants"] = len(variants)
    max_passages = max(1, request.max_passages)
    historical_context: list[str] = []
    source_passages: list[str] = []
    candidates: list[dict[str, Any]] = []
    seed_hits: list[tuple[str, float, str]] = []
    passage_hits: list[tuple[str, float, str]] = []
    include_history = _wants_historical_facts(request.text)

    def _collect_facts_for_variant(
        text: str,
    ) -> tuple[list[tuple[str, float, str]], list[dict[str, Any]]]:
        with profile_stage("facts.variant") as detail:
            with profile_stage("facts.seed_search"):
                seeds = _seed_nodes_for_text(text, request.brain_id)
            detail["seeds"] = len(seeds)
            if not seeds:
                return [], []
            uuid_to_distance: dict[str, float] = {}
            for uuid, distance, _ in seeds:
                dist = _quantize_distance(distance)
                prev = uuid_to_distance.get(uuid)
                if prev is None or dist < prev:
                    uuid_to_distance[uuid] = dist
            seed_uuids = sorted(
                uuid_to_distance.keys(),
                key=lambda uid: (uuid_to_distance[uid], uid),
            )
            with profile_stage(
                "graph.event_neighbors", seed_uuids=len(seed_uuids)
            ) as neighbors_detail:
                neighbors = graph_adapter.get_event_centric_neighbors(
                    seed_uuids, brain_id=request.brain_id
                )
                neighbors_detail["rows"] = len(neighbors)
            local: list[dict[str, Any]] = []
            with profile_stage("facts.assemble"):
                for n, r, m, r2, b in neighbors:
                    if not _fact_predicates_allowed(
                        r, r2, include_history=include_history
                    ):
                        continue
                    distances = [
                        uuid_to_distance[u]
                        for u in (
                            getattr(n, "uuid", None),
                            getattr(m, "uuid", None),
                            getattr(b, "uuid", None),
                        )
                        if u in uuid_to_distance
                    ]
                    score = min(distances) if distances else float("inf")
                    fact_text = _format_event_fact(n, r, m, r2, b)
                    if include_history and not (
                        _is_currently_valid(r) and _is_currently_valid(r2)
                    ):
                        fact_text = f"[superseded] {fact_text}"
                    local.append(
                        {
                            "identified_entity": text,
                            "triple": (n, r, m, r2, b),
                            "score": score,
                            "key": _fact_key(r, r2),
                            "text": fact_text,
                            "chunk_ids": _extract_source_chunk_ids(n, r, m, r2, b),
                        }
                    )
            return list(seeds), local

    def _collect_passages() -> None:
        nonlocal passage_hits
        with profile_stage("passages.collect") as detail:
            queries = [request.text] + [
                v
                for v in variants
                if v.strip().lower() != request.text.strip().lower()
            ]
            if _is_enumeration_question(request.text):
                queries = queries[:6]
            else:
                queries = queries[:3]
            detail["queries"] = len(queries)
            fused_lists: list[list[str]] = []
            id_to_text: dict[str, str] = {}
            id_to_score: dict[str, float] = {}
            for query in queries:
                hits = _retrieve_passages(
                    query,
                    request.brain_id,
                    limit=max(max_passages, _PASSAGE_K),
                )
                ordered_ids: list[str] = []
                for chunk_id, score, body in hits:
                    ordered_ids.append(chunk_id)
                    id_to_text[chunk_id] = body
                    prev = id_to_score.get(chunk_id)
                    if prev is None or score < prev:
                        id_to_score[chunk_id] = score
                if ordered_ids:
                    fused_lists.append(ordered_ids)
            if not fused_lists:
                passage_hits = []
                return
            fused = reciprocal_rank_fusion(fused_lists)
            ranked_ids = [item for item, _ in fused[: max(max_passages, _PASSAGE_K)]]
            passage_hits = [
                (cid, id_to_score.get(cid, 1.0), id_to_text.get(cid, ""))
                for cid in ranked_ids
                if id_to_text.get(cid)
            ]

    async def _get_historical_context():
        nonlocal historical_context
        with profile_stage("historical.context", blocking=False):
            passages = await asyncio.to_thread(
                _retrieve_passages,
                request.text,
                request.brain_id,
                limit=max(request.historical_limit, _PASSAGE_K),
            )
            if passages:
                historical_context = [
                    text for _, _, text in passages[: request.historical_limit]
                ]
                return historical_context
            with profile_stage("historical.fallback", blocking=False):
                text_chunks, structured_data = await asyncio.gather(
                    asyncio.to_thread(
                        data_adapter.get_last_text_chunks,
                        brain_id=request.brain_id,
                        limit=request.historical_limit,
                    ),
                    asyncio.to_thread(
                        data_adapter.get_last_structured_data,
                        brain_id=request.brain_id,
                        limit=request.historical_limit,
                    ),
                )
            historical_context = [text_chunk.text for text_chunk in text_chunks] + [
                json.dumps(structured_data.data)
                for structured_data in structured_data
                if len(structured_data.data.items()) > 0
            ]
            return historical_context

    with profile_stage(
        "retrieval.fanout", blocking=False, variants=len(variants)
    ):
        fact_futures = [
            asyncio.to_thread(_collect_facts_for_variant, variant)
            for variant in variants
        ]
        other_futures = [
            asyncio.to_thread(_collect_passages),
            _get_historical_context(),
        ]
        all_results = await asyncio.gather(*fact_futures, *other_futures)
        variant_outputs = all_results[: len(variants)]

    seed_hits = []
    variant_lists: list[list[dict[str, Any]]] = []
    for seeds, local in variant_outputs:
        seed_hits.extend(seeds)
        variant_lists.append(local)
    candidates = _merge_variant_candidate_lists(variant_lists)

    bridge_paths: list[dict[str, Any]] = []
    with profile_stage(
        "facts.bridge_expand",
        bridges=request.cross_event_bridges,
    ) as detail:
        candidates, bridge_paths = _expand_cross_event_bridges(
            candidates,
            request.brain_id,
            max_per_hub=request.cross_event_bridges,
            include_history=include_history,
        )
        detail["candidates"] = len(candidates)
        detail["paths"] = len(bridge_paths)

    with profile_stage("facts.dedup_rank") as detail:
        ranked = _dedupe_candidates(candidates)
        detail["candidates"] = len(candidates)
        detail["deduped"] = len(ranked)

    if request.use_ppr and seed_hits and ranked:
        with profile_stage("ppr"):
            personalization = _seed_personalization(seed_hits)
            with profile_stage(
                "ppr.adjacency", seeds=len(personalization)
            ) as adjacency_detail:
                adjacency = _build_adjacency_from_seeds(
                    sorted(personalization.keys()),
                    request.brain_id,
                    query=request.text,
                    include_history=include_history,
                )
                adjacency_detail["nodes"] = len(adjacency)
            with profile_stage("ppr.iterations", iterations=_PPR_ITERS):
                ppr_scores = personalized_pagerank(
                    adjacency,
                    personalization,
                    damping=_PPR_DAMPING,
                    iterations=_PPR_ITERS,
                )
            if ppr_scores:
                with profile_stage("ppr.reorder", scored=len(ppr_scores)):
                    for candidate in ranked:
                        n, _r, m, _r2, b = candidate["triple"]
                        masses = [
                            ppr_scores.get(str(getattr(node, "uuid", "")), 0.0)
                            for node in (n, m, b)
                            if getattr(node, "uuid", None)
                        ]
                        candidate["ppr_mass"] = max(masses) if masses else 0.0
                    ranked = sorted(
                        ranked, key=lambda c: _ppr_rank_key(c, ppr_scores)
                    )

    max_facts = max(0, request.max_facts)

    provenance_ids: list[str] = []
    for candidate in ranked:
        provenance_ids.extend(candidate.get("chunk_ids") or [])
    for chunk_id, _, body in passage_hits:
        if body and body not in source_passages:
            source_passages.append(body)
        if chunk_id not in provenance_ids:
            provenance_ids.append(chunk_id)

    chunk_sessions: dict[str, list[str]] = {}
    if provenance_ids:
        with profile_stage(
            "provenance.chunks",
            blocking=False,
            chunk_ids=len(provenance_ids[:80]),
        ):
            try:
                chunks, _ = await asyncio.to_thread(
                    data_adapter.get_text_chunks_by_ids,
                    provenance_ids[:80],
                    False,
                    request.brain_id,
                )
                for chunk in chunks:
                    body = (chunk.text or "").strip()
                    chunk_id = str(getattr(chunk, "id", "") or "")
                    if chunk_id and body:
                        sessions = _session_ids_from_text(body)
                        if sessions:
                            chunk_sessions[chunk_id] = sessions
                    if body and body not in source_passages:
                        source_passages.append(body)
            except Exception:
                pass

    for candidate in ranked:
        if not candidate.get("session_ids"):
            candidate["session_ids"] = _session_ids_for_chunks(
                list(candidate.get("chunk_ids") or []), chunk_sessions
            )

    response_topics: list[dict[str, Any]] = []
    with profile_stage("topics.coarse_to_fine") as topic_detail:
        try:
            from src.core.saving.topic_hyperedges import select_topics_and_sessions

            memberships = graph_adapter.list_topic_memberships(request.brain_id)
            seed_sessions: list[str] = []
            seen_seed: set[str] = set()
            for candidate in ranked:
                for sid in _candidate_session_ids(candidate):
                    if sid not in seen_seed:
                        seen_seed.add(sid)
                        seed_sessions.append(sid)
            response_topics, preferred_sessions = select_topics_and_sessions(
                request.text,
                memberships,
                seed_sessions,
                k_topics=10,
                k_sessions=20 if include_history else 10,
            )
            preferred = set(preferred_sessions)
            topic_detail["topics"] = len(response_topics)
            topic_detail["preferred_sessions"] = len(preferred)
            if preferred:
                for candidate in ranked:
                    sessions = set(_candidate_session_ids(candidate))
                    if sessions & preferred:
                        candidate["topic_preferred"] = True
        except Exception as exc:
            topic_detail["error"] = str(exc)[:120]

    with profile_stage("fact_filter") as detail:
        if request.apply_fact_filter and fact_filter_adapter is not None and ranked:
            detail["filter_applied"] = True
            keep = filter_relevant_facts(
                request.text,
                [c["text"] for c in ranked],
                llm_adapter=fact_filter_adapter,
                max_keep=max(max_facts, len(ranked)),
            )
            filtered = [ranked[i] for i in keep if 0 <= i < len(ranked)]
            if not filtered:
                filtered = ranked
        else:
            detail["filter_applied"] = False
            filtered = ranked

    with profile_stage("facts.diversify", max_facts=max_facts) as detail:
        filtered = _rank_facts_with_completeness(filtered)
        curated = _diversify_facts(filtered, max_facts=max_facts)
        detail["candidates"] = len(filtered)
        detail["curated"] = len(curated)

    with profile_stage("context.assemble") as detail:
        temporal_conflicts = _temporal_conflict_meta(curated)
        text_lines, triples, graph_session_ids = _build_fact_channel(
            curated, chunk_sessions
        )
        source_passages = source_passages[:max_passages]
        detail["facts"] = len(text_lines)
        detail["passages"] = len(source_passages)

    insufficient = False
    if request.sufficiency_retry:
        with profile_stage("sufficiency.check") as detail:
            insufficient = _context_looks_insufficient(
                request.text, text_lines, source_passages + historical_context
            )
            detail["fired"] = insufficient

    if insufficient:
        with profile_stage("sufficiency.retry", blocking=False) as detail:
            followups = _decompose_enumeration_queries(request.text, elements)
            if not followups:
                tokens = [
                    t
                    for t in re.findall(r"[A-Za-z][A-Za-z']+", request.text)
                    if len(t) > 3
                ]
                followups = [" ".join(tokens[:4])] if tokens else []
            detail["followups"] = len(followups[:3])
            extra_passages: list[tuple[str, float, str]] = []
            extra_variant_lists: list[list[dict[str, Any]]] = []
            for query in followups[:3]:
                extra_passages.extend(
                    await asyncio.to_thread(
                        _retrieve_passages,
                        query,
                        request.brain_id,
                        limit=max_passages,
                    )
                )
                seeds, local = await asyncio.to_thread(
                    _collect_facts_for_variant, query
                )
                seed_hits.extend(seeds)
                extra_variant_lists.append(local)
            for chunk_id, _, body in extra_passages:
                if body and body not in source_passages:
                    source_passages.append(body)
                if chunk_id not in provenance_ids:
                    provenance_ids.append(chunk_id)
            source_passages = source_passages[:max_passages]
            candidates.extend(_merge_variant_candidate_lists(extra_variant_lists))
            ranked = _dedupe_candidates(candidates)
            for candidate in ranked:
                if not candidate.get("session_ids"):
                    candidate["session_ids"] = _session_ids_for_chunks(
                        list(candidate.get("chunk_ids") or []), chunk_sessions
                    )
            curated = _diversify_facts(
                _rank_facts_with_completeness(ranked), max_facts=max_facts
            )
            temporal_conflicts = _temporal_conflict_meta(curated)
            for candidate in curated:
                for chunk_id in candidate.get("chunk_ids") or []:
                    if chunk_id not in provenance_ids:
                        provenance_ids.append(chunk_id)
            missing_ids = [
                cid for cid in provenance_ids[:80] if cid not in chunk_sessions
            ]
            if missing_ids:
                try:
                    chunks, _ = await asyncio.to_thread(
                        data_adapter.get_text_chunks_by_ids,
                        missing_ids,
                        False,
                        request.brain_id,
                    )
                    for chunk in chunks:
                        body = (chunk.text or "").strip()
                        chunk_id = str(getattr(chunk, "id", "") or "")
                        if chunk_id and body:
                            sessions = _session_ids_from_text(body)
                            if sessions:
                                chunk_sessions[chunk_id] = sessions
                except Exception:
                    pass
            text_lines, triples, graph_session_ids = _build_fact_channel(
                curated, chunk_sessions
            )

    response_paths = _paths_for_curated(bridge_paths, curated)

    with profile_stage("context.render", passages=len(source_passages)):
        topic_block = [
            f"[topic] {t.get('label')}"
            for t in response_topics
            if isinstance(t, dict) and t.get("label")
        ]
        passage_block = [f"[passage] {p}" for p in source_passages]
        text_context = "\n".join(topic_block + passage_block + text_lines)

    return GetContextResponse(
        text_context=text_context,
        triples=triples,
        historical_context=historical_context,
        source_passages=source_passages,
        graph_session_ids=graph_session_ids or None,
        temporal_conflicts=temporal_conflicts or None,
        paths=response_paths or None,
        topics=response_topics or None,
    )