import json
import sys
from types import SimpleNamespace

import pytest

from hessian_recipes.german_pilot import execute, prepare, validate_splits


def test_split_leakage_is_rejected_by_source_identity():
    with pytest.raises(ValueError, match="disjoint"):
        validate_splits([{"source_id": "same"}], [{"source_id": "same"}], [{"source_id": "external"}])
    with pytest.raises(ValueError, match="duplicate"):
        validate_splits([{"source_id": "x"}, {"source_id": "x"}], [{"source_id": "dev"}], [{"source_id": "external"}])


def test_invalid_compute_budget_fails_before_loading_a_model(tmp_path):
    with pytest.raises(ValueError, match="steps"):
        prepare(tmp_path, tmp_path, tmp_path / "out", steps=0)


def test_training_driver_refuses_a_cpu_node(tmp_path, monkeypatch):
    plan = tmp_path / "plan.json"
    plan.write_text(json.dumps({"model": "must-not-load"}))
    fake_torch = SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))
    monkeypatch.setitem(sys.modules, "torch", fake_torch)
    with pytest.raises(RuntimeError, match="GPU allocation"):
        execute(plan)
    assert not (tmp_path / "run-manifest.json").exists()


def test_preparation_rejects_whole_examples_that_would_truncate_targets(tmp_path, monkeypatch):
    class Tokenizer:
        chat_template = "{{ messages }}"

        def apply_chat_template(self, messages, *, return_dict, **kwargs):
            # Transformers 5 otherwise returns BatchEncoding, whose length is
            # the number of fields rather than the sequence's token count.
            assert return_dict is False
            return [1] * (20 if messages[-1]["content"] == "long" else 2)

        def save_pretrained(self, destination):
            destination.mkdir()
            (destination / "tokenizer.json").write_text("{}")

    model, data = tmp_path / "model", tmp_path / "data"
    model.mkdir()
    (model / "config.json").write_text("{}")
    data.mkdir()
    for name, rows in {
        "train": [{"source_id": "train-a", "prompt": "a", "completion": "long"},
                  {"source_id": "train-b", "prompt": "b", "completion": "short"}],
        "dev": [{"source_id": "dev-a"}],
        "mgsm-de": [{"source_id": "external-a"}],
    }.items():
        (data / f"{name}.jsonl").write_text("".join(json.dumps(row) + "\n" for row in rows))
    (data / "manifest.json").write_text("{}")
    fake = SimpleNamespace(AutoTokenizer=SimpleNamespace(from_pretrained=lambda *a, **k: Tokenizer()))
    monkeypatch.setitem(sys.modules, "transformers", fake)
    result = prepare(model, data, tmp_path / "output", eval_size=1, max_length=5)
    assert result["train_examples"] == 1
    assert result["overlength_rejected"] == [{"source_id": "train-a", "tokens": 20}]
    assert json.loads((tmp_path / "output/train.jsonl").read_text())["source_id"] == "train-b"