import json
import shutil
import tempfile
import unittest
from pathlib import Path

import apply_patches

ROOT = Path(__file__).resolve().parents[1]
FIXTURE = ROOT / "tests" / "fixtures" / "fake-model"
TEMPLATE = (ROOT / "patches" / "chat_template" / "default_chat_template.jinja").read_text(encoding="utf-8")
WHITELIST = apply_patches.load_whitelist(ROOT / "patches" / "architectures" / "whitelist.txt")
ALIASES = apply_patches.load_aliases(ROOT / "patches" / "architectures" / "aliases.json")


class ChatTemplateTest(unittest.TestCase):
    def setUp(self) -> None:
        self.tmp = tempfile.TemporaryDirectory()
        self.model = Path(self.tmp.name) / "model"
        shutil.copytree(FIXTURE, self.model)

    def tearDown(self) -> None:
        self.tmp.cleanup()

    def test_injects_missing_template(self) -> None:
        result = apply_patches.apply_chat_template(self.model, TEMPLATE)
        self.assertTrue(result["injected"])
        data = json.loads((self.model / "tokenizer_config.json").read_text(encoding="utf-8"))
        self.assertEqual(data["chat_template"], TEMPLATE)

    def test_keeps_existing_template(self) -> None:
        path = self.model / "tokenizer_config.json"
        data = json.loads(path.read_text(encoding="utf-8"))
        data["chat_template"] = "already-set"
        path.write_text(json.dumps(data), encoding="utf-8")
        result = apply_patches.apply_chat_template(self.model, TEMPLATE)
        self.assertFalse(result["injected"])
        kept = json.loads(path.read_text(encoding="utf-8"))
        self.assertEqual(kept["chat_template"], "already-set")

    def test_force_overwrites_template(self) -> None:
        path = self.model / "tokenizer_config.json"
        data = json.loads(path.read_text(encoding="utf-8"))
        data["chat_template"] = "already-set"
        path.write_text(json.dumps(data), encoding="utf-8")
        result = apply_patches.apply_chat_template(self.model, TEMPLATE, force=True)
        self.assertTrue(result["injected"])

    def test_missing_tokenizer_config_fails(self) -> None:
        (self.model / "tokenizer_config.json").unlink()
        with self.assertRaises(apply_patches.PatchError):
            apply_patches.apply_chat_template(self.model, TEMPLATE)

    def test_preview_renders_roles(self) -> None:
        text = apply_patches.preview_default_template(
            [
                {"role": "system", "content": "你是助手"},
                {"role": "user", "content": "你好"},
            ]
        )
        self.assertIn("<|user|>\n你好<|end|>\n", text)
        self.assertTrue(text.endswith("<|assistant|>\n"))
        self.assertIn("message['role']", TEMPLATE)


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