# training/document_reasoning_loader.py

from __future__ import annotations

from typing import Any

from datasets import load_dataset, Dataset


def load_document_reasoning_sft(
    tokenizer,
    train_path: str,
    validation_path: str | None = None,
    stress_validation_path: str | None = None,
    add_eos_to_completion: bool = True,
):
    """
    Loads JSONL made by DocumentsReasoningDataset.py and converts:

        messages = [system, user, assistant]

    into:

        prompt     = rendered system+user chat
        completion = assistant content

    This is safer than relying on assistant_only_loss=True because
    prompt-completion SFT naturally trains only the completion side.
    """

    data_files = {"train": train_path}

    if validation_path:
        data_files["validation"] = validation_path

    if stress_validation_path:
        data_files["validation_stress"] = stress_validation_path

    raw = load_dataset("json", data_files=data_files)

    def convert(example: dict[str, Any]) -> dict[str, str]:
        messages = example["messages"]

        if len(messages) < 2:
            raise ValueError("Expected at least user + assistant messages.")

        if messages[-1]["role"] != "assistant":
            raise ValueError("Last message must be assistant.")

        prompt_messages = messages[:-1]
        assistant_content = messages[-1]["content"]

        prompt = tokenizer.apply_chat_template(
            prompt_messages,
            tokenize=False,
            add_generation_prompt=True,
        )

        completion = assistant_content

        if add_eos_to_completion and tokenizer.eos_token:
            completion = completion + tokenizer.eos_token

        return {
            "prompt": prompt,
            "completion": completion,
        }

    keep_cols = ["prompt", "completion"]

    converted = raw.map(
        convert,
        remove_columns=raw["train"].column_names,
        desc="Converting messages to prompt/completion",
    )

    train_dataset = converted["train"]
    eval_dataset = converted.get("validation", None)
    stress_eval_dataset = converted.get("validation_stress", None)

    return train_dataset, eval_dataset, stress_eval_dataset