import json
import tempfile
import unittest
from pathlib import Path

from engine_compat import QWEN_CHAT_TEMPLATE, CompatibilityError, prepare_model


class EngineCompatTests(unittest.TestCase):
    def _model(self, root: Path, tokenizer: dict) -> Path:
        model = root / "model"
        model.mkdir()
        (model / "config.json").write_text(
            json.dumps({"model_type": "qwen2", "architectures": ["Qwen2ForCausalLM"]}),
            encoding="utf-8",
        )
        (model / "tokenizer_config.json").write_text(
            json.dumps(tokenizer), encoding="utf-8"
        )
        (model / "model.safetensors").write_bytes(b"sample-weights")
        return model

    def test_repairs_list_and_missing_qwen_chat_template(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = self._model(
                root,
                {
                    "eos_token": "<|im_end|>",
                    "extra_special_tokens": ["<|im_start|>", "<|im_end|>"],
                    "additional_special_tokens": ["<|im_start|>"],
                },
            )
            view = Path(prepare_model(source, work_root=root))
            fixed = json.loads((view / "tokenizer_config.json").read_text("utf-8"))
            original = json.loads((source / "tokenizer_config.json").read_text("utf-8"))
            self.assertIn("extra_special_tokens", original)
            self.assertNotIn("extra_special_tokens", fixed)
            self.assertEqual(
                fixed["additional_special_tokens"], ["<|im_start|>", "<|im_end|>"]
            )
            self.assertEqual(fixed["chat_template"], QWEN_CHAT_TEMPLATE)
            self.assertEqual((view / "model.safetensors").read_bytes(), b"sample-weights")
            self.assertTrue((view / ".modelhub-compat-report.json").is_file())

    def test_no_change_does_not_create_view(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = self._model(
                root,
                {"eos_token": "<|im_end|>", "chat_template": "custom template"},
            )
            self.assertIsNone(prepare_model(source, work_root=root))
            self.assertEqual([p.name for p in root.iterdir()], ["model"])

    def test_standalone_chat_template_is_preserved(self):
        for filename in ("chat_template.jinja", "chat_template.json"):
            with self.subTest(filename=filename), tempfile.TemporaryDirectory() as directory:
                root = Path(directory)
                source = self._model(root, {"eos_token": "<|im_end|>"})
                (source / filename).write_text("{{ messages }}", encoding="utf-8")
                self.assertIsNone(prepare_model(source, work_root=root))

    def test_known_tokenizer_class_is_replaced_only_with_tokenizer_json(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = self._model(
                root,
                {"tokenizer_class": "TokenizersBackend", "chat_template": "custom"},
            )
            with self.assertRaisesRegex(CompatibilityError, "tokenizer.json"):
                prepare_model(source, work_root=root)
            (source / "tokenizer.json").write_text("{}", encoding="utf-8")
            view = Path(prepare_model(source, work_root=root))
            fixed = json.loads((view / "tokenizer_config.json").read_text("utf-8"))
            self.assertEqual(fixed["tokenizer_class"], "PreTrainedTokenizerFast")

    def test_unknown_model_family_does_not_get_qwen_template(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = self._model(root, {"eos_token": "<|im_end|>"})
            (source / "config.json").write_text(
                json.dumps({"model_type": "llama"}), encoding="utf-8"
            )
            self.assertIsNone(prepare_model(source, work_root=root))


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