from __future__ import annotations

"""Graded relevance recipe for advanced evaluation.

This recipe generates (query, context, grade) triples with multi-level relevance
scores (0-3) instead of binary labels. Useful for training and evaluating
ranking models with graded relevance.
"""

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="graded_relevance",
        description=(
            "Graded relevance pairs: paragraph as query, other paragraphs as contexts "
            "with relevance grades 0 (irrelevant), 1 (marginally relevant), "
            "2 (relevant), 3 (highly relevant) based on section hierarchy."
        ),
        input_query="",
        prompt_template="",
        output_schema={
            "query_text": "str",
            "context_text": "str",
            "grade": "int",  # 0-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_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_graded_relevance(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate graded relevance pairs for advanced evaluation.

    Grading logic:
    - Grade 3 (highly relevant): Same section, adjacent paragraphs
    - Grade 2 (relevant): Same section, non-adjacent paragraphs
    - Grade 1 (marginally relevant): Different section, same parent document
    - Grade 0 (irrelevant): Random paragraphs from different sections

    For each query paragraph, we emit multiple contexts with different grades.
    """

    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_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)

    # Build adjacency map (next_sibling edges)
    next_sibling_map: dict[str, str] = {}
    for edge in graph.edges:
        if edge.edge_type == "next_sibling":
            next_sibling_map[edge.source] = edge.target

    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 i, query_para in enumerate(paragraphs):
                query_text = (query_para.label or "").strip()
                if not query_text:
                    continue

                # Grade 3: Adjacent paragraphs in same section
                if i > 0:
                    prev_para = paragraphs[i - 1]
                    ctx_text = (prev_para.label or "").strip()
                    if ctx_text:
                        record = {
                            "document_id": document_id,
                            "query_node_id": query_para.id,
                            "context_node_id": prev_para.id,
                            "query_text": query_text,
                            "context_text": ctx_text,
                            "grade": 3,
                            "metadata": {
                                **(graph.metadata or {}),
                                "section_id": section_id,
                                "relationship": "adjacent_prev",
                            },
                        }
                        f.write(json.dumps(record, ensure_ascii=False) + "\n")

                if i < len(paragraphs) - 1:
                    next_para = paragraphs[i + 1]
                    ctx_text = (next_para.label or "").strip()
                    if ctx_text:
                        record = {
                            "document_id": document_id,
                            "query_node_id": query_para.id,
                            "context_node_id": next_para.id,
                            "query_text": query_text,
                            "context_text": ctx_text,
                            "grade": 3,
                            "metadata": {
                                **(graph.metadata or {}),
                                "section_id": section_id,
                                "relationship": "adjacent_next",
                            },
                        }
                        f.write(json.dumps(record, ensure_ascii=False) + "\n")

                # Grade 2: Non-adjacent paragraphs in same section
                for j, ctx_para in enumerate(paragraphs):
                    if abs(i - j) <= 1:  # Skip self and adjacent
                        continue
                    ctx_text = (ctx_para.label or "").strip()
                    if ctx_text:
                        record = {
                            "document_id": document_id,
                            "query_node_id": query_para.id,
                            "context_node_id": ctx_para.id,
                            "query_text": query_text,
                            "context_text": ctx_text,
                            "grade": 2,
                            "metadata": {
                                **(graph.metadata or {}),
                                "section_id": section_id,
                                "relationship": "same_section_non_adjacent",
                            },
                        }
                        f.write(json.dumps(record, ensure_ascii=False) + "\n")

                # Grade 1: Paragraphs from different sections (marginally relevant)
                # Pick up to 2 paragraphs from other sections
                other_sections = [sid for sid in section_paragraphs.keys() if sid != section_id and sid is not None]
                count_grade1 = 0
                for other_sid in other_sections:
                    if count_grade1 >= 2:
                        break
                    other_paras = section_paragraphs[other_sid]
                    if other_paras:
                        ctx_para = other_paras[0]
                        ctx_text = (ctx_para.label or "").strip()
                        if ctx_text:
                            record = {
                                "document_id": document_id,
                                "query_node_id": query_para.id,
                                "context_node_id": ctx_para.id,
                                "query_text": query_text,
                                "context_text": ctx_text,
                                "grade": 1,
                                "metadata": {
                                    **(graph.metadata or {}),
                                    "query_section_id": section_id,
                                    "context_section_id": other_sid,
                                    "relationship": "different_section",
                                },
                            }
                            f.write(json.dumps(record, ensure_ascii=False) + "\n")
                            count_grade1 += 1

                # Grade 0: Irrelevant (from distant sections)
                # Pick 1 paragraph from a section far away
                if len(other_sections) > 1:
                    distant_sid = other_sections[-1]  # Pick last section as "distant"
                    distant_paras = section_paragraphs[distant_sid]
                    if distant_paras:
                        ctx_para = distant_paras[-1]  # Pick last paragraph
                        ctx_text = (ctx_para.label or "").strip()
                        if ctx_text:
                            record = {
                                "document_id": document_id,
                                "query_node_id": query_para.id,
                                "context_node_id": ctx_para.id,
                                "query_text": query_text,
                                "context_text": ctx_text,
                                "grade": 0,
                                "metadata": {
                                    **(graph.metadata or {}),
                                    "query_section_id": section_id,
                                    "context_section_id": distant_sid,
                                    "relationship": "irrelevant",
                                },
                            }
                            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path
