from tests.support.paths import PROJECT_ROOT
import importlib.util
import unittest


def load_module(name, relative_path):
    specification = importlib.util.spec_from_file_location(
        name,
        PROJECT_ROOT / relative_path,
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class ModelCatalogTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.models = load_module(
            "anyimage_server_models",
            "src/anyimage/server/model_catalog.py",
        )


    def test_unknown_model_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "Unknown model"):
            self.models.get_downloadable_model("UNKNOWN")

    def test_three_tiers_and_retired_keys(self):
        self.assertEqual(list(self.models.MOGE_MODELS), ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE3_VITL"])
        self.assertEqual(list(self.models.BACKGROUND_MODELS), ["BIREFNET_LITE", "BEN2_BASE", "BIREFNET_HR_MATTING"])
        self.assertEqual(list(self.models.UPSCALE_MODELS), ["REALESRGAN_GENERAL_WDN_X4V3", "REALESRGAN_X4PLUS", "HAT_GAN_X4_SHARPER"])
        self.assertEqual(len(self.models.DOWNLOADABLE_MODELS), 9)
        self.assertEqual(self.models.DEFAULT_MODEL_KEYS, ("MOGE2_VITS_NORMAL", "BIREFNET_LITE", "REALESRGAN_GENERAL_WDN_X4V3"))
        for key in ("DA3_BASE", "DA3_SMALL", "MOGE2_VITL_NORMAL", "SAFMN_REAL_X4", "PISA_SR"):
            with self.assertRaisesRegex(ValueError, "Unknown model"):
                self.models.get_downloadable_model(key)

    def test_ben2_download_only_includes_required_hugging_face_files(self):
        model = self.models.get_downloadable_model("BEN2_BASE")

        self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
        self.assertEqual(model.license_name, "Apache 2.0")
        self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
        self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
        self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")


    def test_ben2_model_directory_uses_repository_name(self):
        self.assertEqual(
            self.models.get_downloadable_model("BEN2_BASE").directory_name,
            "BEN2-ONNX",
        )

    def test_upscale_models_use_verified_onnx_artifacts(self):
        expected = {
            "REALESRGAN_GENERAL_WDN_X4V3": ("realesr-general-wdn-x4v3.onnx", 4867419, "b021a837"),
            "HAT_GAN_X4_SHARPER": ("model.onnx", 45874438, "032d8b03"),
            "REALESRGAN_X4PLUS": ("realesrgan-x4plus.onnx", 67051618, "4ed6a45a"),
        }
        self.assertEqual(set(self.models.UPSCALE_MODELS), set(expected))
        for key, (filename, size, checksum_prefix) in expected.items():
            with self.subTest(model=key):
                model = self.models.get_downloadable_model(key)
                self.assertIn(filename, model.download_patterns)
                self.assertIn(filename, model.ready_patterns)
                self.assertEqual(model.r2_files[0][:2], (filename, size))
                self.assertTrue(model.r2_files[0][2].startswith(checksum_prefix))
                self.assertEqual(self.models.model_record(model)["family"], "upscale")
        wdn = self.models.UPSCALE_MODELS["REALESRGAN_GENERAL_WDN_X4V3"]
        self.assertEqual(wdn.repository, "xinntao/Real-ESRGAN")
        self.assertEqual(wdn.huggingface_repository, "")
        self.assertEqual(
            self.models.UPSCALE_MODELS["HAT_GAN_X4_SHARPER"].huggingface_repository, ""
        )


    def test_moge2_models_use_official_onnx_models(self):
        for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
                             ("MOGE2_VITB_NORMAL", "vitb")):
            with self.subTest(key=key):
                model = self.models.get_downloadable_model(key)
                self.assertEqual(
                    model.repository,
                    f"Ruicheng/moge-2-{variant}-normal-onnx",
                )
                self.assertEqual(model.license_name, "MIT")
                self.assertEqual(model.download_patterns, ("model.onnx",))
                self.assertEqual(model.ready_patterns, ("model.onnx",))
                self.assertEqual(model.r2_directory, model.repository.split("/")[-1])
                self.assertEqual(model.r2_files[0][0], "model.onnx")

    def test_moge2_models_match_uploaded_r2_variants(self):
        self.assertEqual(
            list(self.models.MOGE2_MODELS),
            ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL"],
        )
        self.assertEqual(
            [model.r2_directory for model in self.models.MOGE2_MODELS.values()],
            [
                "moge-2-vits-normal-onnx",
                "moge-2-vitb-normal-onnx",
            ],
        )