from __future__ import annotations

"""Entity-context ranking recipe for knowledge-enhanced evaluation.

This recipe generates ranking datasets that mix entity knowledge and section
summaries as contexts. Useful for evaluating retrieval systems that need to
handle both factual entity information and broader contextual summaries.
"""

from collections import defaultdict
from pathlib import Path
import json
from typing import Dict

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


def _ensure_recipe(recipe: DatasetRecipe | None) -> DatasetRecipe:
    if recipe is not None:
        return recipe
    return DatasetRecipe(
        name="entity_context_ranking",
        description=(
            "Entity-context ranking dataset: for each query paragraph, collect "
            "contexts from entity nodes (positive if entity mentioned in query) "
            "and summary nodes (positive if from same section), creating a mixed "
            "knowledge + context ranking scenario."
        ),
        input_query="",
        prompt_template="",
        output_schema={
            "query_text": "str",
            "contexts": "list[dict]",  # Each dict has context_text, label, context_type
        },
    )


def _build_nodes_index(graph: DocumentGraph) -> Dict[str, GraphNode]:
    return {n.id: n for n in graph.nodes}


def _find_parent_section(node_id: str, graph: DocumentGraph, nodes_index: Dict[str, GraphNode]) -> str | None:
    """Find the parent SECTION node for a given node."""
    for edge in graph.edges:
        if edge.edge_type == "parent_of" and edge.target == node_id:
            parent = nodes_index.get(edge.source)
            if parent and parent.node_type is NodeType.SECTION:
                return parent.id
    return None


def _get_entity_mentions(paragraph_text: str, entity_name: str) -> bool:
    """Simple check if entity is mentioned in paragraph (case-insensitive)."""
    return entity_name.lower() in paragraph_text.lower()


def run_entity_context_ranking(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate entity-context ranking dataset.

    For each query paragraph:
    - Positive entity contexts: entities mentioned in the query paragraph (label=1)
    - Negative entity contexts: entities NOT mentioned in the query (label=0)
    - Positive summary contexts: summaries from the same section (label=1)
    - Negative summary contexts: summaries from other sections (label=0)

    Output format: JSONL with fields:
    - query_text
    - query_node_id
    - contexts: list of {context_node_id, context_text, label, context_type}
    - metadata
    """

    recipe = _ensure_recipe(recipe)
    max_negative_entities = int((recipe.params or {}).get("max_negative_entities", 2))
    max_negative_summaries = int((recipe.params or {}).get("max_negative_summaries", 2))

    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"

    nodes_index = _build_nodes_index(graph)

    # Collect entities and summaries
    entities = [n for n in graph.nodes if n.node_type is NodeType.ENTITY]
    summaries = [n for n in graph.nodes if n.node_type is NodeType.SUMMARY]
    paragraphs = [n for n in graph.nodes if n.node_type is NodeType.PARAGRAPH]

    # Build section -> summaries map
    section_summaries: dict[str | None, list[GraphNode]] = defaultdict(list)
    for summary in summaries:
        section_id = _find_parent_section(summary.id, graph, nodes_index)
        section_summaries[section_id].append(summary)

    with out_path.open("w", encoding="utf-8") as f:
        for query_para in paragraphs:
            query_text = (query_para.label or "").strip()
            if not query_text:
                continue

            query_section_id = _find_parent_section(query_para.id, graph, nodes_index)
            contexts = []

            # Entity contexts
            for entity in entities:
                entity_name = entity.label or ""
                entity_desc = (entity.metadata or {}).get("description", "")
                entity_text = f"{entity_name}: {entity_desc}".strip()

                if not entity_text or entity_text == ":":
                    continue

                # Positive if entity mentioned in query
                is_mentioned = _get_entity_mentions(query_text, entity_name)
                label = 1 if is_mentioned else 0

                # Limit negatives
                if label == 0:
                    if sum(1 for c in contexts if c.get("context_type") == "entity" and c["label"] == 0) >= max_negative_entities:
                        continue

                contexts.append({
                    "context_node_id": entity.id,
                    "context_text": entity_text,
                    "label": label,
                    "context_type": "entity",
                    "relationship": "mentioned" if is_mentioned else "not_mentioned",
                })

            # Summary contexts
            for section_id, section_sums in section_summaries.items():
                for summary in section_sums:
                    summary_text = (summary.label or "").strip()
                    if not summary_text:
                        continue

                    # Positive if from same section
                    is_same_section = (section_id == query_section_id)
                    label = 1 if is_same_section else 0

                    # Limit negatives
                    if label == 0:
                        if sum(1 for c in contexts if c.get("context_type") == "summary" and c["label"] == 0) >= max_negative_summaries:
                            continue

                    contexts.append({
                        "context_node_id": summary.id,
                        "context_text": summary_text,
                        "label": label,
                        "context_type": "summary",
                        "relationship": "same_section" if is_same_section else "different_section",
                    })

            if not contexts:
                continue

            record = {
                "document_id": document_id,
                "query_node_id": query_para.id,
                "query_text": query_text,
                "contexts": contexts,
                "metadata": {
                    **(graph.metadata or {}),
                    "section_id": query_section_id,
                    "num_entity_positives": sum(1 for c in contexts if c.get("context_type") == "entity" and c["label"] == 1),
                    "num_entity_negatives": sum(1 for c in contexts if c.get("context_type") == "entity" and c["label"] == 0),
                    "num_summary_positives": sum(1 for c in contexts if c.get("context_type") == "summary" and c["label"] == 1),
                    "num_summary_negatives": sum(1 for c in contexts if c.get("context_type") == "summary" and c["label"] == 0),
                },
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

