"""Small offline checks; no pretrained weights, dataset downloads, or GPU jobs."""
import hashlib
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

from .data import encode_example, normalize_messages, prompt_key
from .runtime import CONFIG, conversation, load_engine
from .train import collate, read_data


class PipelineTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        from tokenizers import Tokenizer
        from tokenizers.models import WordLevel
        from tokenizers.pre_tokenizers import Whitespace
        from transformers import PreTrainedTokenizerFast
        words = ["[UNK]", "[PAD]", "<|im_start|>", "<|im_end|>",
                 "system", "user", "assistant", "first", "second", "answer", "other"]
        tokenizer = Tokenizer(WordLevel({word: i for i, word in enumerate(words)}, unk_token="[UNK]"))
        tokenizer.pre_tokenizer = Whitespace()
        cls.tokenizer = PreTrainedTokenizerFast(
            tokenizer_object=tokenizer, unk_token="[UNK]", pad_token="[PAD]", eos_token="<|im_end|>",
            additional_special_tokens=["<|im_start|>", "<|im_end|>"])
        cls.tokenizer.chat_template = (
            "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n' + "
            "message['content'] + '<|im_end|>\n' }}{% endfor %}"
            "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}")

    def test_multiturn_assistant_targets_and_padding(self):
        messages = normalize_messages({"messages": [
            {"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"},
            {"role": "user", "content": "second"}, {"role": "assistant", "content": "other"}]})
        encoded = encode_example(self.tokenizer, messages, 2048)
        targets = [x for x in encoded["labels"] if x != -100]
        self.assertEqual(self.tokenizer.decode(targets), "answer <|im_end|> other <|im_end|>")
        short = encode_example(self.tokenizer, messages[:3], 2048)
        batch = collate([encoded, short], self.tokenizer.pad_token_id)
        self.assertTrue((batch["labels"][1, len(short["input_ids"]):] == -100).all())
        self.assertTrue((batch["attention_mask"][1, len(short["input_ids"]):] == 0).all())
        with self.assertRaises(ValueError):
            encode_example(self.tokenizer, messages, 4)

    def test_no_image_records_and_no_partial_turns(self):
        for messages in [
            [{"role": "user", "content": [{"type": "image"}]}],
            [{"role": "user", "content": "first"}],
            [{"role": "assistant", "content": "answer"}, {"role": "user", "content": "first"}],
        ]:
            with self.assertRaises(ValueError):
                normalize_messages({"messages": messages})
        self.assertEqual(conversation("first")[-1], {"role": "user", "content": "first"})
        with self.assertRaises(ValueError):
            conversation("first", [{"role": "user", "content": "second"}])

    def test_cpu_training_and_inference_fail_before_download(self):
        from argparse import Namespace
        from .train import train
        with patch("torch.cuda.is_available", return_value=False):
            with self.assertRaisesRegex(RuntimeError, "cloud model"):
                load_engine()
            with self.assertRaisesRegex(RuntimeError, "No training or model download"):
                train(Namespace())

    def test_manifest_integrity_and_prompt_separation(self):
        with tempfile.TemporaryDirectory() as temp:
            folder = Path(temp)
            manifest = {"config": CONFIG, "sha256": {}, "counts": {}}
            for split in ["train", "validation", "test"]:
                messages = normalize_messages({"messages": [
                    {"role": "user", "content": f"question {split}"},
                    {"role": "assistant", "content": "answer"}]})
                row = {"messages": messages, "prompt_hash": prompt_key(messages)}
                raw = (json.dumps(row) + "\n").encode()
                (folder / f"{split}.jsonl").write_bytes(raw)
                manifest["sha256"][split] = hashlib.sha256(raw).hexdigest()
                manifest["counts"][split] = 1
            (folder / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
            self.assertEqual(set(read_data(folder)), {"train", "validation", "test"})
            (folder / "train.jsonl").write_text("changed", encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "Data changed"):
                read_data(folder)

    def test_small_qwen_lora_forward_and_gradients(self):
        import torch
        from peft import LoraConfig, TaskType, get_peft_model
        from transformers import Qwen3Config, Qwen3ForCausalLM
        torch.set_num_threads(2)
        model = get_peft_model(Qwen3ForCausalLM(Qwen3Config(
            vocab_size=len(self.tokenizer), hidden_size=32, intermediate_size=64,
            num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1,
            head_dim=16, max_position_embeddings=512)),
            LoraConfig(r=2, lora_alpha=4, target_modules="all-linear", task_type=TaskType.CAUSAL_LM))
        messages = normalize_messages({"messages": [
            {"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"}]})
        batch = collate([encode_example(self.tokenizer, messages, 512)], self.tokenizer.pad_token_id)
        loss = model(**batch).loss
        self.assertTrue(torch.isfinite(loss))
        loss.backward()
        self.assertTrue(any(p.grad is not None and p.grad.abs().sum() > 0
                            for name, p in model.named_parameters() if "lora_B" in name))
        self.assertFalse(any(p.requires_grad for name, p in model.named_parameters() if "lora_" not in name))


if __name__ == "__main__":
    unittest.main()