from __future__ import annotations

"""Entity-centric knowledge recipe built on semantic graph entities.

This recipe is deterministic and builds a small knowledge dataset from
ENTITY nodes produced by the semantic graph builder.
"""

from collections import defaultdict
from pathlib import Path
import json

from ..core.models import DatasetRecipe, DocumentGraph, NodeType


def _ensure_recipe(recipe: DatasetRecipe | None) -> DatasetRecipe:
    if recipe is not None:
        return recipe
    return DatasetRecipe(
        name="entity_knowledge",
        description="Default entity-centric knowledge recipe from ENTITY nodes.",
        input_query="",
        prompt_template="",
        output_schema={
            "entity": "str",
            "keyword": "str",
            "mentions": "list[str]",
        },
    )


def run_entity_knowledge(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate an entity-centric knowledge dataset from ENTITY nodes.

    Each record aggregates all paragraphs that mention a given entity keyword.
    """

    recipe = _ensure_recipe(recipe)

    base = Path(base_dir)
    ds_dir = base / "datasets" / recipe.name
    ds_dir.mkdir(parents=True, exist_ok=True)
    out_path = ds_dir / f"{document_id}.jsonl"

    # Index paragraphs by id for quick lookup
    paragraph_by_id = {
        n.id: n for n in graph.nodes if n.node_type is NodeType.PARAGRAPH
    }

    # Map keyword -> list of (entity_id, paragraph_text)
    agg: dict[str, list[tuple[str, str]]] = defaultdict(list)

    for node in graph.nodes:
        if node.node_type is not NodeType.ENTITY:
            continue
        kw = (node.metadata or {}).get("keyword") or (node.label or "")
        kw = str(kw).lower().strip()
        if not kw:
            continue

        # Find the paragraph that mentions this entity via 'mentions' edges
        for edge in graph.edges:
            if edge.edge_type != "mentions":
                continue
            if edge.target != node.id:
                continue
            para = paragraph_by_id.get(edge.source)
            if not para or not para.label:
                continue
            agg[kw].append((node.id, para.label.strip()))

    with out_path.open("w", encoding="utf-8") as f:
        for kw, mentions in agg.items():
            record = {
                "document_id": document_id,
                "entity": kw,
                "keyword": kw,
                "mentions": [text for _eid, text in mentions],
                "metadata": graph.metadata or {},
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

