from __future__ import annotations

"""Graph-aware neighborhood context recipe.

This recipe emits RAG-style records that include a paragraph and its
structural neighborhood (same section siblings and parent section summary).
"""

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

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


def _ensure_recipe(recipe: DatasetRecipe | None) -> DatasetRecipe:
    if recipe is not None:
        return recipe
    return DatasetRecipe(
        name="graph_neighborhood",
        description="Graph-aware neighborhood context around paragraph nodes.",
        input_query="",
        prompt_template="",
        output_schema={
            "text": "str",
            "context": "list[str]",
        },
    )


def _build_index(graph: DocumentGraph) -> tuple[Dict[str, GraphNode], Dict[str, List[GraphEdge]]]:
    nodes_by_id: Dict[str, GraphNode] = {n.id: n for n in graph.nodes}
    edges_by_source: Dict[str, List[GraphEdge]] = {}
    for e in graph.edges:
        edges_by_source.setdefault(e.source, []).append(e)
    return nodes_by_id, edges_by_source


def _find_parent_section(node_id: str, graph: DocumentGraph, nodes_by_id: Dict[str, GraphNode]) -> GraphNode | None:
    # Walk incoming parent_of edges backwards.
    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 and parent.node_type is NodeType.SECTION:
            return parent
        if parent is not None:
            return _find_parent_section(parent.id, graph, nodes_by_id)
    return None


def _collect_sibling_paragraphs(
    paragraph_id: str,
    parent_section: GraphNode | None,
    graph: DocumentGraph,
    nodes_by_id: Dict[str, GraphNode],
) -> List[str]:
    if parent_section is None:
        return []

    # Collect children of the section in order via parent_of + next_sibling edges.
    child_ids: List[str] = []
    for e in graph.edges:
        if e.edge_type == "parent_of" and e.source == parent_section.id:
            child_ids.append(e.target)

    # Order children using next_sibling where available.
    ordered: List[str] = []
    remaining = set(child_ids)
    # naive: just keep child_ids order; neighborhood does not depend on exact
    # graph traversal.
    ordered.extend(child_ids)

    siblings: List[str] = []
    for cid in ordered:
        if cid == paragraph_id:
            continue
        node = nodes_by_id.get(cid)
        if node and node.node_type is NodeType.PARAGRAPH and node.label:
            siblings.append(node.label.strip())
    return siblings


def run_graph_neighborhood(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Generate graph-aware paragraph + neighborhood context records.

    Each record corresponds to a paragraph node and contains the paragraph
    text plus additional context texts from its structural neighborhood:

    - sibling paragraphs in the same section
    - the section-level summary if 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_index(graph)

    # Map from section id to its SUMMARY node text, if any.
    section_summary: Dict[str, str] = {}
    for node in graph.nodes:
        if node.node_type is NodeType.SUMMARY and node.metadata:
            sid = node.metadata.get("source_section")
            if isinstance(sid, str) and node.label:
                section_summary[sid] = node.label.strip()

    with out_path.open("w", encoding="utf-8") as f:
        for node in graph.nodes:
            if node.node_type is not NodeType.PARAGRAPH or not node.label:
                continue

            text = node.label.strip()
            if not text:
                continue

            parent_section = _find_parent_section(node.id, graph, nodes_by_id)
            sibling_texts = _collect_sibling_paragraphs(
                node.id, parent_section, graph, nodes_by_id
            )

            context_parts: List[str] = []
            if parent_section is not None:
                summary_text = section_summary.get(parent_section.id)
                if summary_text:
                    context_parts.append(summary_text)

            context_parts.extend(sibling_texts)

            record = {
                "document_id": document_id,
                "node_id": node.id,
                "text": text,
                "context": context_parts,
                "metadata": {
                    **(graph.metadata or {}),
                    **(node.metadata or {}),
                },
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    return out_path

