from __future__ import annotations

"""Cross-document ranking recipe for multi-document evaluation.

This recipe is designed for scenarios where multiple documents are processed
and we want to generate ranking datasets that include cross-document negatives.
For single-document graphs, it behaves similarly to multi_context_ranking but
with explicit cross-document metadata.
"""

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="cross_document_ranking",
        description=(
            "Cross-document ranking dataset: for each query paragraph, collect "
            "positive contexts (same section) and negative contexts (other sections, "
            "potentially from other documents if metadata indicates multi-doc scenario)."
        ),
        input_query="",
        prompt_template="",
        output_schema={
            "query_text": "str",
            "contexts": "list[dict]",  # Each dict has context_text, label, source_doc
        },
    )


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 run_cross_document_ranking(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate cross-document ranking dataset.

    For each query paragraph:
    - Positive contexts: paragraphs from the same section (label=1)
    - Negative contexts: paragraphs from other sections (label=0)
    - Metadata includes source_document_id for cross-doc scenarios

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

    recipe = _ensure_recipe(recipe)
    max_negatives = int((recipe.params or {}).get("max_negatives", 3))

    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)

    # Group paragraphs by section
    section_paragraphs: dict[str | None, list[GraphNode]] = defaultdict(list)
    for node in graph.nodes:
        if node.node_type is not NodeType.PARAGRAPH:
            continue
        section_id = _find_parent_section(node.id, graph, nodes_index)
        section_paragraphs[section_id].append(node)

    # Collect all paragraphs for negative sampling
    all_paragraphs = [p for paras in section_paragraphs.values() for p in paras]

    with out_path.open("w", encoding="utf-8") as f:
        for section_id, paragraphs in section_paragraphs.items():
            if section_id is None or len(paragraphs) < 2:
                continue

            for query_para in paragraphs:
                query_text = (query_para.label or "").strip()
                if not query_text:
                    continue

                contexts = []

                # Positive contexts: same section
                for ctx_para in paragraphs:
                    if ctx_para.id == query_para.id:
                        continue
                    ctx_text = (ctx_para.label or "").strip()
                    if not ctx_text:
                        continue

                    contexts.append({
                        "context_node_id": ctx_para.id,
                        "context_text": ctx_text,
                        "label": 1,
                        "source_document_id": document_id,
                        "relationship": "same_section",
                    })

                # Negative contexts: other sections
                negatives_added = 0
                for neg_para in all_paragraphs:
                    if negatives_added >= max_negatives:
                        break

                    neg_section_id = _find_parent_section(neg_para.id, graph, nodes_index)
                    if neg_section_id == section_id:
                        continue

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

                    contexts.append({
                        "context_node_id": neg_para.id,
                        "context_text": neg_text,
                        "label": 0,
                        "source_document_id": document_id,
                        "relationship": "different_section",
                    })
                    negatives_added += 1

                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": section_id,
                        "num_positives": sum(1 for c in contexts if c["label"] == 1),
                        "num_negatives": sum(1 for c in contexts if c["label"] == 0),
                    },
                }
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

