"""Prepared data stays small without losing training or evaluation inputs."""
import json
import re

import pytest

from lab.common import read_jsonl, verify_bundle
from lab.prepare import prepare
from lab.data import read_dataset


class Tokenizer:
    eos_token = "<|im_end|>"
    eos_token_id = 1
    pad_token_id = 0
    chat_template = "test chat template"
    special_tokens_map = {"eos_token": eos_token}

    def get_vocab(self):
        return {self.eos_token: 1}

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
        text = "".join(m["role"] + ":" + m["content"] + self.eos_token + "\n" for m in messages)
        return text + ("assistant:" if add_generation_prompt else "")

    def __call__(self, text, add_special_tokens=False):
        tokens = re.findall(r"<\|im_end\|>|[\s\S]", text)
        return {"input_ids": [1 if t == self.eos_token else ord(t) + 2 for t in tokens]}


def test_seed_bundle_has_six_files_and_preserves_the_training_targets(tmp_path):
    episodes, facts = read_dataset("data/dataset.json")
    episodes = [e for e in episodes if e["split"] == "train"][:3]
    manifest = prepare(episodes, facts, Tokenizer(), tmp_path, allow_draft=True)
    assert manifest["schema"] == 2
    assert {p.relative_to(tmp_path).as_posix() for p in tmp_path.rglob("*") if p.is_file()} == {
        "facts.jsonl", "train.episodes.jsonl", "fewshots.json", "manifest.json",
        "_derived/train.sft.jsonl", "_derived/token-audit.json",
    }
    assert read_jsonl(tmp_path / "train.episodes.jsonl") == episodes
    rows = read_jsonl(tmp_path / "_derived/train.sft.jsonl")
    assert len(rows) == sum(len(e["turns"]) for e in episodes)
    assert rows[0]["completion"] == episodes[0]["turns"][0]["assistant"] + Tokenizer.eos_token
    assert episodes[0]["turns"][0]["assistant"] in rows[1]["prompt"]
    assert episodes[0]["turns"][1]["user"] not in rows[0]["prompt"]
    assert verify_bundle(tmp_path) == manifest
    with pytest.raises(ValueError, match="not empty"):
        prepare(episodes, facts, Tokenizer(), tmp_path, allow_draft=True)


def test_nonempty_validation_and_test_splits_remain_available(tmp_path):
    episodes, facts = read_dataset("data/dataset.json")
    episodes = episodes[:3] + [next(e for e in episodes if e["split"] == split) for split in ("val", "test")]
    manifest = prepare(episodes, facts, Tokenizer(), tmp_path, allow_draft=True)
    assert manifest["counts"] == {"train": 3, "val": 1, "test": 1}
    assert len(read_jsonl(tmp_path / "_derived/val.sft.jsonl")) == 2
    assert read_jsonl(tmp_path / "test.episodes.jsonl") == [episodes[-1]]
    assert not (tmp_path / "_derived/test.sft.jsonl").exists()
    fewshots = json.loads((tmp_path / "fewshots.json").read_text())
    train_ids = {e["id"] for e in episodes if e["split"] == "train"}
    assert all(e["episode_id"] in train_ids for e in fewshots)
    payload = tmp_path / "_derived/train.sft.jsonl"
    payload.write_text(payload.read_text() + "\n")
    with pytest.raises(ValueError, match="data changed"):
        verify_bundle(tmp_path)


@pytest.mark.parametrize("schema,relative", [(1, "train.sft.jsonl"), (2, "_derived/train.sft.jsonl")])
def test_training_reader_supports_old_and_new_bundles(tmp_path, schema, relative):
    from lab.common import write_jsonl
    from lab.data import read_sft
    rows = [{"input_ids": [10, 11, 1], "completion_mask": [0, 1, 1]}]
    write_jsonl(tmp_path / relative, rows)
    manifest = {"schema": schema, "counts": {"train": 1}}
    assert read_sft(tmp_path, "train", manifest) == rows
    assert read_sft(tmp_path, "val", manifest) == []
    manifest["counts"]["val"] = 1
    with pytest.raises(FileNotFoundError):
        read_sft(tmp_path, "val", manifest)


def test_inference_rejects_missing_split_before_loading_a_model(tmp_path, monkeypatch):
    from lab.infer import main
    episodes, facts = read_dataset("data/dataset.json")
    prepare(episodes[:3], facts, Tokenizer(), tmp_path, allow_draft=True)
    monkeypatch.setattr("sys.argv", ["infer", "--data", str(tmp_path), "--output", str(tmp_path / "replies.jsonl")])
    with pytest.raises(ValueError, match="selected split is empty"):
        main()


def test_fewshots_cover_all_grade_bands_using_train_only(tmp_path):
    episodes, facts = read_dataset("data/dataset.json")
    prepare(episodes, facts, Tokenizer(), tmp_path, allow_draft=True)
    fewshots = json.loads((tmp_path / "fewshots.json").read_text())
    assert {row["grade_band"] for row in fewshots} == {"1-2", "3-4", "5-6"}
    assert all(row["episode_id"] in {e["id"] for e in episodes if e["split"] == "train"} for row in fewshots)