"""Offline one-step integration with random tiny weights; no quality claims.

Run after installing .[training]: python examples/training/cpu_smoke.py
Without --output, the temporary model, data and checkpoints are removed on exit.
"""

from __future__ import annotations

import json
import argparse
from contextlib import nullcontext
from pathlib import Path
import tempfile


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, help="Keep the model, checkpoints and manifests in a fresh directory")
    args = parser.parse_args()
    if args.output is not None:
        destination = args.output.resolve()
        if destination.exists() and any(destination.iterdir()):
            parser.error(f"Output directory must be empty: {destination}")
        destination.mkdir(parents=True, exist_ok=True)
    else:
        destination = None
    import torch
    from tokenizers import Tokenizer
    from tokenizers.models import WordLevel
    from tokenizers.pre_tokenizers import Whitespace
    from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast
    from hessian_recipes.training import run
    from hessian_recipes.distillation import generate

    torch.set_num_threads(1)
    torch.manual_seed(42)
    context = nullcontext(str(destination)) if destination is not None else tempfile.TemporaryDirectory(prefix="hessian-recipes-cpu-")
    with context as temporary:
        root = Path(temporary)
        checkpoint = root / "tiny-model"
        vocabulary = {token: i for i, token in enumerate([
            "<pad>", "<eos>", "<unk>", "<user>", "<assistant>",
            "What", "is", "plus", "?", "1", "2", "3", "4", "5", "6",
        ])}
        backend = Tokenizer(WordLevel(vocabulary, unk_token="<unk>"))
        backend.pre_tokenizer = Whitespace()
        tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, pad_token="<pad>",
                                              eos_token="<eos>", unk_token="<unk>")
        tokenizer.chat_template = (
            "{% for message in messages %}{{ '<' + message['role'] + '>' }} "
            "{{ message['content'] }} {{ eos_token }} {% endfor %}"
            "{% if add_generation_prompt %}{{ '<assistant> ' }}{% endif %}"
        )
        tokenizer.save_pretrained(checkpoint)
        model = GPT2LMHeadModel(GPT2Config(vocab_size=len(vocabulary), n_positions=128,
                                         n_embd=16, n_layer=1, n_head=2,
                                         bos_token_id=1, eos_token_id=1, pad_token_id=0))
        model.save_pretrained(checkpoint, safe_serialization=True)
        dataset = root / "train.jsonl"
        dataset.write_text("\n".join(json.dumps({"prompt": prompt, "completion": target, "solution": target})
                                     for prompt, target in [("What is 1 plus 2 ?", "3"),
                                                            ("What is 2 plus 2 ?", "4")]) + "\n")
        teacher_config = root / "teacher.json"
        teacher_config.write_text(json.dumps({"schema_version": 1, "model": {"name": str(checkpoint)},
                                              "dataset": str(dataset), "output": str(root / "teacher.jsonl"),
                                              "max_new_tokens": 4, "device": "cpu"}))
        teacher_report = generate(teacher_config)
        assert teacher_report["status"] == "completed"
        assert {"model.safetensors", "tokenizer.json", "config.json"} <= teacher_report["local_model_sha256"].keys()
        assert len((root / "teacher.jsonl").read_text().splitlines()) == 2
        print(json.dumps({"method": "local_teacher_generation", "status": teacher_report["status"]}))
        for method in ("sft", "grpo"):
            settings = {"max_steps": 1, "use_cpu": True, "bf16": False,
                        "per_device_train_batch_size": 2}
            settings.update({"max_length": 64} if method == "sft" else
                            {"max_completion_length": 1, "num_generations": 32,
                             "per_device_train_batch_size": 32})
            config = {"schema_version": 1, "method": method, "model": {"name": str(checkpoint)},
                      "dataset": str(dataset), "output_dir": str(root / method), "training": settings}
            config_path = root / f"{method}.json"
            config_path.write_text(json.dumps(config))
            report = run(config_path)
            assert report["status"] == "completed"
            assert (Path(report["checkpoint"]) / "model.safetensors").exists()
            assert (root / method / "recipe-manifest.json").exists()
            assert {"model.safetensors", "tokenizer.json", "config.json"} <= report["local_model_sha256"].keys()
            changed = GPT2LMHeadModel.from_pretrained(report["checkpoint"])
            assert any(not torch.equal(before, after) for before, after in
                       zip(model.parameters(), changed.parameters())), f"{method} made no parameter update"
            if method == "grpo":
                reward_logs = [log for log in report["train_log_history"] if "reward_std" in log]
                assert reward_logs and reward_logs[0]["reward_std"] > 0, "GRPO needs varied group rewards"
                assert reward_logs[0]["grad_norm"] > 0, "GRPO gradient must be nonzero"
            print(json.dumps({"method": method, "status": report["status"], "metrics": report["metrics"]}))


if __name__ == "__main__":
    main()