from __future__ import annotations

"""Recipe implementation for simple RAG chunks.

For v0.1 this recipe treats each paragraph node in the ``DocumentGraph`` as
an individual chunk. More advanced splitting (by token length, sliding
windows, etc.) can be added later using the same interface.
"""

from pathlib import Path
import json

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


def _ensure_recipe(recipe: DatasetRecipe | None) -> DatasetRecipe:
    """Allow tests to pass ``None`` and fall back to a default recipe.

    This keeps :func:`run_rag_chunks` convenient in unit tests while real
    CLI usage always provides a concrete recipe configuration.
    """

    if recipe is not None:
        return recipe

    return DatasetRecipe(
        name="rag_chunks",
        description="Default RAG chunks recipe (one paragraph per chunk).",
        input_query="",
        prompt_template="",
        output_schema={},
    )


def run_rag_chunks(
    recipe: DatasetRecipe | None,
    graph: DocumentGraph,
    base_dir: str | Path,
    document_id: str,
) -> Path:
    """Materialise a RAG chunk dataset as JSONL.

    The output path is ``<base_dir>/datasets/<recipe.name>/<document_id>.jsonl``.
    Each line contains a single JSON object with at least the fields:

    - ``document_id``
    - ``node_id``
    - ``text``
    - ``metadata`` (merged graph + node metadata)
    """

    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"

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

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

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

    return out_path

