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"
WHITELIST = apply_patches.load_whitelist(ROOT / "patches" / "architectures" / "whitelist.txt")


class ArchitectureTest(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 write_architectures(self, names: list[str]) -> None:
        path = self.model / "config.json"
        data = json.loads(path.read_text(encoding="utf-8"))
        data["architectures"] = names
        path.write_text(json.dumps(data), encoding="utf-8")

    def test_whitelist_passes(self) -> None:
        result = apply_patches.apply_architectures(self.model, WHITELIST, {}, "fail")
        self.assertEqual(result["architectures"], ["LlamaForCausalLM"])
        self.assertFalse(result["changed"])

    def test_unknown_architecture_fails(self) -> None:
        self.write_architectures(["ExampleForCausalLM"])
        with self.assertRaises(apply_patches.PatchError) as ctx:
            apply_patches.apply_architectures(self.model, WHITELIST, {}, "fail")
        self.assertIn("ExampleForCausalLM", str(ctx.exception))

    def test_unknown_architecture_skip_keeps_name(self) -> None:
        self.write_architectures(["ExampleForCausalLM"])
        result = apply_patches.apply_architectures(self.model, WHITELIST, {}, "skip")
        self.assertEqual(result["skipped"], ["ExampleForCausalLM"])
        data = json.loads((self.model / "config.json").read_text(encoding="utf-8"))
        self.assertEqual(data["architectures"], ["ExampleForCausalLM"])

    def test_alias_rewrites_to_whitelist(self) -> None:
        self.write_architectures(["LlamaForCausalLMAlias"])
        result = apply_patches.apply_architectures(
            self.model,
            WHITELIST,
            {"LlamaForCausalLMAlias": "LlamaForCausalLM"},
            "fail",
        )
        self.assertEqual(result["rewritten"], ["LlamaForCausalLMAlias->LlamaForCausalLM"])
        data = json.loads((self.model / "config.json").read_text(encoding="utf-8"))
        self.assertEqual(data["architectures"], ["LlamaForCausalLM"])

    def test_shipped_aliases_are_empty(self) -> None:
        aliases = apply_patches.load_aliases(ROOT / "patches" / "architectures" / "aliases.json")
        self.assertEqual(aliases, {})


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