"""Traversal queries over the knowledge graph: entity/citation -> chunks, neighbors."""

from __future__ import annotations

from typing import Any

from .builder import _normalize_entity


def _entity_key(text: str) -> str:
    return f"entity:{_normalize_entity(text)}"


def _chunk_payload(graph: Any, chunk_node: str) -> dict[str, Any]:
    data = graph.nodes[chunk_node]
    doc_node = next(
        (src for src, _, e in graph.in_edges(chunk_node, data=True) if e.get("kind") == "CONTAINS"),
        None,
    )
    doc = graph.nodes[doc_node] if doc_node is not None else {}
    return {
        "chunk": chunk_node,
        "chunk_index": data.get("chunk_index", 0),
        "text": data.get("text", ""),
        "source_path": doc.get("source_path", ""),
        "ocr_engine": doc.get("ocr_engine", ""),
        "ocr_confidence_mean": doc.get("ocr_confidence_mean", 0.0),
        "timestamp": doc.get("timestamp", ""),
    }


def chunks_for_entity(graph: Any, entity_text: str) -> list[dict[str, Any]]:
    """All chunks mentioning the given entity (case-insensitive)."""
    key = _entity_key(entity_text)
    if key not in graph:
        return []
    chunks = [
        src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "MENTIONS"
    ]
    return [_chunk_payload(graph, c) for c in sorted(chunks)]


def chunks_for_citation(graph: Any, identifier: str) -> list[dict[str, Any]]:
    """All chunks citing the given identifier (e.g. '10.1038/nature12345')."""
    key = identifier if identifier.startswith("citation:") else f"citation:{identifier}"
    if key not in graph:
        return []
    chunks = [src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "CITES"]
    return [_chunk_payload(graph, c) for c in sorted(chunks)]


def related_entities(graph: Any, entity_text: str, limit: int = 10) -> list[tuple[str, int]]:
    """Entities co-occurring with the given one, as (text, shared-chunk weight)."""
    key = _entity_key(entity_text)
    if key not in graph:
        return []
    neighbors: dict[str, int] = {}
    for _, target, edge in graph.out_edges(key, data=True):
        if edge.get("kind") == "CO_OCCURS":
            neighbors[target] = edge.get("weight", 1)
    for source, _, edge in graph.in_edges(key, data=True):
        if edge.get("kind") == "CO_OCCURS":
            neighbors[source] = max(neighbors.get(source, 0), edge.get("weight", 1))
    ranked = sorted(
        ((graph.nodes[n].get("text", n), w) for n, w in neighbors.items()),
        key=lambda item: item[1],
        reverse=True,
    )
    return ranked[:limit]


def expand_context(graph: Any, entity_text: str, limit: int = 5) -> list[dict[str, Any]]:
    """Chunks mentioning the entity or its strongest co-occurring neighbors.

    This is the graph-RAG primitive: one hop of CO_OCCURS expansion, so a query
    for 'BRCA1' also surfaces chunks that only mention its interactors.
    """
    seen: set[str] = set()
    results: list[dict[str, Any]] = []

    seeds = [entity_text] + [text for text, _ in related_entities(graph, entity_text, limit=limit)]
    for seed in seeds:
        for payload in chunks_for_entity(graph, seed):
            if payload["chunk"] in seen:
                continue
            seen.add(payload["chunk"])
            payload = {**payload, "via_entity": seed}
            results.append(payload)
    return results