import json

from stego_olmoe.data import (
    DataCollatorForRouterSFT,
    JsonlSFTDataset,
    TokenizedSFTDataset,
    render_messages_prefix_answer,
    write_tokenized_cache,
)


class TinyTokenizer:
    eos_token = "<eos>"
    eos_token_id = 0
    pad_token = "<pad>"
    pad_token_id = 0

    def __len__(self):
        return 512

    def __call__(self, text, truncation=True, max_length=512, padding=False, add_special_tokens=True, return_tensors=None):
        ids = [ord(ch) % 255 + 1 for ch in text]
        if add_special_tokens:
            ids = [self.eos_token_id] + ids
        ids = ids[:max_length]
        return {"input_ids": ids, "attention_mask": [1] * len(ids)}


def test_prompt_answer_router_mask_answer_only(tmp_path):
    path = tmp_path / "data.jsonl"
    path.write_text(json.dumps({"prompt": "A", "answer": "B"}) + "\n", encoding="utf-8")
    dataset = JsonlSFTDataset(path, TinyTokenizer(), max_length=64)
    item = dataset[0]
    assert item["router_loss_mask"].sum().item() >= 1
    assert item["router_loss_mask"][0].item() == 0


def test_collator_pads_router_mask(tmp_path):
    path = tmp_path / "data.jsonl"
    path.write_text(
        json.dumps({"text": "short"}) + "\n" + json.dumps({"text": "a little longer"}) + "\n",
        encoding="utf-8",
    )
    tokenizer = TinyTokenizer()
    dataset = JsonlSFTDataset(path, tokenizer, max_length=64)
    batch = DataCollatorForRouterSFT(tokenizer)([dataset[0], dataset[1]])
    assert batch["input_ids"].shape[0] == 2
    assert batch["router_loss_mask"].shape == batch["attention_mask"].shape


def test_messages_render_last_assistant_target():
    prefix, answer = render_messages_prefix_answer(
        [
            {"role": "system", "content": "Be concise."},
            {"role": "user", "content": "Question"},
            {"role": "assistant", "content": "First answer"},
            {"role": "user", "content": "Follow-up"},
            {"role": "assistant", "content": "Final answer"},
        ]
    )
    assert prefix.endswith("Assistant:")
    assert "First answer" in prefix
    assert "Final answer" not in prefix
    assert answer == "Final answer"


def test_tokenized_cache_round_trips(tmp_path):
    jsonl_path = tmp_path / "data.jsonl"
    cache_path = tmp_path / "data.pt"
    jsonl_path.write_text(
        json.dumps({"prompt": "A", "answer": "B"}) + "\n" + json.dumps({"text": "free text"}) + "\n",
        encoding="utf-8",
    )

    metadata = write_tokenized_cache(jsonl_path, cache_path, TinyTokenizer(), max_length=64)
    jsonl_dataset = JsonlSFTDataset(jsonl_path, TinyTokenizer(), max_length=64)
    cache_dataset = TokenizedSFTDataset(cache_path)

    assert metadata["n_records"] == 2
    assert len(cache_dataset) == len(jsonl_dataset)
    for key, value in jsonl_dataset[0].items():
        assert value.equal(cache_dataset[0][key])