import importlib.util
import unittest
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[1]


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 ModelTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.models = load_module(
            "snapform_models",
            "src/snapform/models.py",
        )

    def test_base_model_is_default_friendly_apache_model(self):
        model = self.models.get_model("BASE")

        self.assertEqual(model.repository, "depth-anything/DA3-BASE")
        self.assertEqual(model.license_name, "Apache 2.0")

    def test_large_model_exposes_noncommercial_license(self):
        model = self.models.get_model("LARGE_1_1")

        self.assertEqual(model.license_name, "CC BY-NC 4.0")
        self.assertFalse(model.enabled)

    def test_giant_models_expose_gaussian_head(self):
        for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
            with self.subTest(model_key=model_key):
                self.assertTrue(self.models.get_model(model_key).supports_gs)
                self.assertFalse(self.models.get_model(model_key).enabled)

    def test_only_uploaded_da3_models_are_enabled(self):
        enabled = [key for key, model in self.models.MODELS.items() if model.enabled]

        self.assertEqual(enabled, ["SMALL", "BASE"])

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

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

        self.assertEqual(model.repository, "PramaLLC/BEN2")
        self.assertEqual(model.license_name, "MIT")
        self.assertEqual(
            model.download_patterns,
            ("config.json", "model.safetensors"),
        )
        self.assertNotIn("BEN2_BASE", self.models.MODELS)

    def test_da3_model_directories_share_the_da3_prefix(self):
        for model_key, model in self.models.MODELS.items():
            with self.subTest(model_key=model_key):
                self.assertTrue(
                    model.directory_name.startswith("da3_"),
                    model.directory_name,
                )
                self.assertEqual(
                    self.models.model_directory_name(model_key),
                    model.directory_name,
                )

    def test_ben2_model_directory_uses_ben2_prefix(self):
        self.assertEqual(
            self.models.model_directory_name("BEN2_BASE"),
            "ben2_base",
        )

    def test_moge2_uses_single_checkpoint_and_mit_license(self):
        model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")

        self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal")
        self.assertEqual(model.license_name, "MIT")
        self.assertEqual(model.download_patterns, ("model.pt",))
        self.assertEqual(model.ready_patterns, ("model.pt",))
        self.assertFalse(model.supports_gs)

    def test_moge2_models_match_uploaded_r2_variants(self):
        self.assertEqual(
            list(self.models.MOGE2_MODELS),
            ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE2_VITL_NORMAL"],
        )
        self.assertEqual(
            [model.r2_directory for model in self.models.MOGE2_MODELS.values()],
            ["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
        )