from __future__ import annotations

"""Multi-context ranking recipe.

This recipe groups multiple positive and negative contexts per query paragraph,
producing a structure suitable for learning-to-rank and evaluation tasks.
"""

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="multi_context_ranking",
        description=(
            "Multi-context ranking: multiple positive and negative contexts "
            "per query paragraph."
        ),
        input_query="",
        prompt_template="",
        output_schema={
            "query_text": "str",
            "contexts": "list[dict]",
        },
        params={"max_negatives": 3},
    )


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_multi_context_ranking(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate multi-context ranking records for evaluation.

    For each paragraph that belongs to a section we emit a single record with:

    - ``query_text``: the paragraph text.
    - ``contexts``: a list of context paragraphs, each with a binary label where
      1 denotes a positive (same section) and 0 denotes a negative (different
      section).
    """

    recipe = _ensure_recipe(recipe)

    max_negatives = 3
    params = getattr(recipe, "params", None) or {}
    try:
        maybe = int(params.get("max_negatives", max_negatives))
        if maybe > 0:
            max_negatives = maybe
    except Exception:
        # Ignore invalid values and keep the default.
        pass

    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:
            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
            ]

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

                contexts: list[dict] = []

                # Positive contexts: 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
                    contexts.append(
                        {
                            "context_node_id": ctx.id,
                            "context_text": ctx_text,
                            "label": 1,
                        }
                    )

                # Negative contexts: paragraphs from other sections, up to max_negatives.
                if negative_candidates:
                    neg_count = 0
                    idx = 0
                    while neg_count < max_negatives and idx < len(negative_candidates):
                        neg_ctx = negative_candidates[idx]
                        idx += 1
                        neg_text = (neg_ctx.label or "").strip()
                        if not neg_text:
                            continue
                        contexts.append(
                            {
                                "context_node_id": neg_ctx.id,
                                "context_text": neg_text,
                                "label": 0,
                            }
                        )
                        neg_count += 1

                if not contexts:
                    continue

                record = {
                    "document_id": document_id,
                    "query_node_id": query.id,
                    "query_text": query_text,
                    "contexts": contexts,
                    "metadata": {
                        **(graph.metadata or {}),
                        "query_section_id": section_id,
                    },
                }
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

