import json
import unittest
from pathlib import Path

from core.type_utils import CATEGORY_MAP, normalize_download_category, resolve_model_category


class CategoryMappingTests(unittest.TestCase):
    def test_folder_category_aliases_resolve_to_existing_comfyui_keys(self):
        expected = {
            "checkpoint": "checkpoints",
            "lora": "loras",
            "embedding": "embeddings",
            "textual inversion": "embeddings",
            "textualinversion": "embeddings",
            "upscaler": "upscale_models",
            "unet gguf": "diffusion_models",
            "select-safetensors": "diffusion_models",
            "clip": "text_encoders",
            "ip_adapter": "ipadapter",
            "ultralytics_bbox": "ultralytics",
            "background-removal-model": "background_removal",
            "default": "upscale_models",
        }

        for raw_category, canonical_category in expected.items():
            self.assertEqual(
                canonical_category,
                resolve_model_category(raw_category, target_format="folder"),
            )
            self.assertEqual(
                canonical_category,
                normalize_download_category(raw_category),
            )

    def test_unknown_and_empty_categories_keep_current_fallbacks(self):
        self.assertEqual("checkpoints", normalize_download_category(""))
        self.assertEqual("checkpoints", normalize_download_category(None))
        self.assertEqual("new_category", normalize_download_category("new category"))

    def test_frontend_category_alias_artifact_matches_backend_map(self):
        project_root = Path(__file__).resolve().parents[1]
        artifact = (
            project_root / "web" / "resolver" / "utils" / "category_aliases.generated.js"
        ).read_text(encoding="utf-8")
        payload = artifact.split("Object.freeze(", 1)[1].rsplit(");", 1)[0].strip()

        self.assertEqual(CATEGORY_MAP, json.loads(payload))


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