from __future__ import annotations

"""Section-based relevance pairs recipe.

This recipe builds a simple evaluation dataset of (query, context, label)
pairs, where queries and contexts are paragraphs and the label indicates
whether they come from the same section (relevant) or a different section
(irrelevant).
"""

from pathlib import Path
import json
from typing import Dict, List

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="section_relevance",
        description=(
            "Section-based relevance pairs: paragraph as query, other paragraphs "
            "as positive/negative contexts based on section membership."
        ),
        input_query="",
        prompt_template="",
        output_schema={
            "query_text": "str",
            "context_text": "str",
            "label": "int",
        },
    )


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_by_id: Dict[str, GraphNode],
) -> GraphNode | None:
    """Return the nearest SECTION ancestor of ``node_id``, if any."""

    for e in graph.edges:
        if e.edge_type != "parent_of" or e.target != node_id:
            continue
        parent = nodes_by_id.get(e.source)
        if parent is None:
            continue
        if parent.node_type is NodeType.SECTION:
            return parent
        return _find_parent_section(parent.id, graph, nodes_by_id)
    return None


def run_section_relevance(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate section-based relevance pairs for evaluation.

    For each paragraph that belongs to a section we emit:

    - positive pairs (label=1) with other paragraphs in the same section
    - at least one negative pair (label=0) with a paragraph from a different
      section, when available.
    """

    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"

    nodes_by_id = _build_nodes_index(graph)

    # Map paragraph -> section id and section -> list of paragraphs.
    section_id_by_paragraph: Dict[str, str] = {}
    paragraphs_by_section: Dict[str, List[GraphNode]] = {}
    all_paragraphs: List[GraphNode] = []

    for node in graph.nodes:
        if node.node_type is not NodeType.PARAGRAPH or not (node.label and node.label.strip()):
            continue
        all_paragraphs.append(node)
        parent_section = _find_parent_section(node.id, graph, nodes_by_id)
        if parent_section is None:
            # Skip paragraphs that are not under a section to keep semantics
            # simple and deterministic.
            continue
        sid = parent_section.id
        section_id_by_paragraph[node.id] = sid
        paragraphs_by_section.setdefault(sid, []).append(node)

    paragraphs_in_sections: List[GraphNode] = [
        p for p in all_paragraphs if p.id in section_id_by_paragraph
    ]

    with out_path.open("w", encoding="utf-8") as f:
        for section_id, paras in paragraphs_by_section.items():
            # Candidates for negative contexts: paragraphs from other sections.
            negative_candidates = [
                p
                for p in paragraphs_in_sections
                if section_id_by_paragraph.get(p.id) != section_id
            ]

            neg_idx = 0

            for query in paras:
                query_text = (query.label or "").strip()
                if not query_text:
                    continue

                # Positive pairs: same section, excluding self.
                for ctx in paras:
                    if ctx.id == query.id:
                        continue
                    ctx_text = (ctx.label or "").strip()
                    if not ctx_text:
                        continue

                    record = {
                        "document_id": document_id,
                        "query_node_id": query.id,
                        "context_node_id": ctx.id,
                        "query_text": query_text,
                        "context_text": ctx_text,
                        "label": 1,
                        "metadata": {
                            **(graph.metadata or {}),
                            "query_section_id": section_id,
                            "context_section_id": section_id,
                        },
                    }
                    f.write(json.dumps(record, ensure_ascii=False) + "\n")

                # One negative pair per query when possible.
                if not negative_candidates:
                    continue

                neg_ctx = negative_candidates[neg_idx % len(negative_candidates)]
                neg_idx += 1

                neg_text = (neg_ctx.label or "").strip()
                if not neg_text:
                    continue

                neg_sid = section_id_by_paragraph.get(neg_ctx.id)

                record = {
                    "document_id": document_id,
                    "query_node_id": query.id,
                    "context_node_id": neg_ctx.id,
                    "query_text": query_text,
                    "context_text": neg_text,
                    "label": 0,
                    "metadata": {
                        **(graph.metadata or {}),
                        "query_section_id": section_id,
                        "context_section_id": neg_sid,
                    },
                }
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

