import unittest

from core.contracts import ResolvedModel
from core.path_templates import infer_download_path_templates
from core.settings import calculate_template_subfolder

BASE_MODELS = {
    "base_models": [
        {
            "name": "Pony",
            "aliases": ["pony", "ponyxl", "pony diffusion"],
        },
        {
            "name": "SDXL 1.0",
            "aliases": ["sdxl", "sdxl10"],
        },
        {
            "name": "Illustrious",
            "aliases": ["illustrious", "illustriousxl"],
        },
        {
            "name": "Flux.1 Krea",
            "aliases": ["Flux.1 Krea", "flux 1 krea"],
        },
        {
            "name": "Krea 2",
            "aliases": ["krea", "krea 2"],
        },
        {
            "name": "Anima",
            "aliases": ["anima"],
        },
        {
            "name": "Ideogram 4.0",
            "aliases": ["Ideogram 4.0", "ideogram", "ideogram 4 0"],
        },
        {
            "name": "Upscaler",
            "aliases": ["upscaler"],
        },
    ]
}


def _typed_models(items):
    return [ResolvedModel.from_mapping(item) for item in items]


class PathTemplateInferenceTests(unittest.TestCase):
    def test_rejects_non_model_collections(self):
        with self.assertRaises(TypeError):
            infer_download_path_templates(0, BASE_MODELS)

    def test_detects_base_model_and_first_tag_layout(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "loras",
                        "relative_path": r"ponyxl\style\model-a.safetensors",
                        "path": r"C:\models\loras\ponyxl\style\model-a.safetensors",
                    },
                    {
                        "category": "lora",
                        "relative_path": r"ponyxl\character\model-b.safetensors",
                        "path": r"C:\models\loras\ponyxl\character\model-b.safetensors",
                    },
                    {
                        "category": "loras",
                        "relative_path": r"sdxl\style\model-c.safetensors",
                        "path": r"C:\models\loras\sdxl\style\model-c.safetensors",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertEqual("{base_model}/{first_tag}", result["templates"]["loras"])
        self.assertEqual("ponyxl", result["base_model_path_mappings"]["Pony"])
        self.assertEqual("sdxl", result["base_model_path_mappings"]["SDXL 1.0"])
        self.assertTrue(result["categories"]["loras"]["apply"])

    def test_detects_flat_layout(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "vae",
                        "relative_path": "vae-a.safetensors",
                        "path": r"C:\models\vae\vae-a.safetensors",
                    },
                    {
                        "category": "vae",
                        "relative_path": "vae-b.safetensors",
                        "path": r"C:\models\vae\vae-b.safetensors",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertEqual("", result["templates"]["vae"])
        self.assertTrue(result["categories"]["vae"]["apply"])

    def test_detects_nested_checkpoint_base_model_mappings(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "checkpoints",
                        "relative_path": r"SDXL\Pony\pony-realism.safetensors",
                        "path": r"C:\models\checkpoints\SDXL\Pony\pony-realism.safetensors",
                    },
                    {
                        "category": "checkpoints",
                        "relative_path": r"SDXL\Illustrious\wai.safetensors",
                        "path": r"C:\models\checkpoints\SDXL\Illustrious\wai.safetensors",
                    },
                    {
                        "category": "checkpoints",
                        "relative_path": r"SDXL\Illustrious\realistic\real.safetensors",
                        "path": r"C:\models\checkpoints\SDXL\Illustrious\realistic\real.safetensors",
                    },
                    {
                        "category": "checkpoints",
                        "relative_path": r"SDXL\SDXL\juggernaut.safetensors",
                        "path": r"C:\models\checkpoints\SDXL\SDXL\juggernaut.safetensors",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertEqual("{base_model}", result["templates"]["checkpoints"])
        self.assertEqual("SDXL/Pony", result["base_model_path_mappings"]["Pony"])
        self.assertEqual(
            "SDXL/Illustrious",
            result["base_model_path_mappings"]["Illustrious"],
        )
        self.assertEqual(
            "SDXL/SDXL",
            result["base_model_path_mappings"]["SDXL 1.0"],
        )

    def test_base_model_mapping_matches_longer_metadata_name(self):
        subfolder = calculate_template_subfolder(
            "checkpoints",
            {"base_model": "Pony Diffusion V6 XL"},
            {
                "download_path_templates": {"checkpoints": "{base_model}"},
                "base_model_path_mappings": {"Pony": "SDXL/Pony"},
            },
        )

        self.assertEqual("SDXL/Pony", subfolder)

    def test_detect_does_not_map_krea2_to_flux_krea_folder(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "loras",
                        "relative_path": r"FLUX\KREA\concept\snofs_krea_v1.safetensors",
                        "path": r"C:\models\loras\FLUX\KREA\concept\snofs_krea_v1.safetensors",
                    },
                    {
                        "category": "loras",
                        "relative_path": r"FLUX\KREA\style\another_flux_krea.safetensors",
                        "path": r"C:\models\loras\FLUX\KREA\style\another_flux_krea.safetensors",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertNotIn("Krea 2", result["base_model_path_mappings"])
        self.assertEqual(
            "FLUX/KREA",
            result["base_model_path_mappings"]["Flux.1 Krea"],
        )

    def test_detect_does_not_map_anima_to_wan_animate_folder(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "checkpoints",
                        "relative_path": r"WAN\WAN2.2ANIMATE\model-a.safetensors",
                        "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-a.safetensors",
                    },
                    {
                        "category": "checkpoints",
                        "relative_path": r"WAN\WAN2.2ANIMATE\model-b.safetensors",
                        "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-b.safetensors",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertNotIn("Anima", result["base_model_path_mappings"])

    def test_detect_accepts_version_and_plural_alias_suffixes(self):
        result = infer_download_path_templates(
            _typed_models(
                [
                    {
                        "category": "diffusion_models",
                        "relative_path": r"IDEOGRAM4\model-a.safetensors",
                        "path": r"C:\models\diffusion_models\IDEOGRAM4\model-a.safetensors",
                    },
                    {
                        "category": "upscale_models",
                        "relative_path": r"ESRGAN\Upscalers\photo\model-a.pth",
                        "path": r"C:\models\upscale_models\ESRGAN\Upscalers\photo\model-a.pth",
                    },
                ]
            ),
            BASE_MODELS,
        )

        self.assertEqual(
            "IDEOGRAM4",
            result["base_model_path_mappings"]["Ideogram 4.0"],
        )
        self.assertEqual(
            "ESRGAN/Upscalers",
            result["base_model_path_mappings"]["Upscaler"],
        )


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