import ast
import importlib.util
import sys
import types
import unittest
from pathlib import Path
from unittest import mock


ROOT = Path(__file__).resolve().parent
EGOSPEAK_ROOT = ROOT / "third_party" / "EgoSpeak"
MINIROAD_MODEL = EGOSPEAK_ROOT / "MiniROAD" / "model"


class _Registry(dict):
    def register(self, name, module=None):
        if module is not None:
            self[name] = module
            return None

        def decorator(value):
            self[name] = value
            return value

        return decorator


def _load_model_builder():
    package_name = "_egospeak_optional_mamba_test_model"
    module_name = f"{package_name}.model_builder"

    package = types.ModuleType(package_name)
    package.__path__ = [str(MINIROAD_MODEL)]
    fake_utils = types.ModuleType("utils")
    fake_utils.Registry = _Registry

    spec = importlib.util.spec_from_file_location(
        module_name, MINIROAD_MODEL / "model_builder.py"
    )
    module = importlib.util.module_from_spec(spec)
    with mock.patch.dict(
        sys.modules,
        {package_name: package, module_name: module, "utils": fake_utils},
    ):
        assert spec.loader is not None
        spec.loader.exec_module(module)
    return module


class EgoSpeakOptionalMambaTest(unittest.TestCase):
    def test_default_requirements_do_not_install_mamba(self):
        requirements = (EGOSPEAK_ROOT / "requirements.txt").read_text()
        packages = {
            line.partition("==")[0].strip().lower()
            for line in requirements.splitlines()
            if line.strip() and not line.lstrip().startswith(("#", "-"))
        }
        self.assertNotIn("mamba-ssm", packages)

    def test_model_package_does_not_eagerly_import_mamba(self):
        tree = ast.parse((MINIROAD_MODEL / "__init__.py").read_text())
        imports = {
            node.module
            for node in ast.walk(tree)
            if isinstance(node, ast.ImportFrom)
        }
        self.assertNotIn("mamba", imports)

    def test_miniroad_build_does_not_load_mamba(self):
        module = _load_model_builder()

        class DummyMiniROAD:
            def __init__(self, cfg):
                self.cfg = cfg

            def to(self, device):
                self.device = device
                return self

        module.META_ARCHITECTURES["MiniROADA"] = DummyMiniROAD
        with mock.patch.object(module.importlib, "import_module") as import_module:
            model = module.build_model({"model": "MiniROADA"}, "cuda")

        import_module.assert_not_called()
        self.assertEqual(model.device, "cuda")

    def test_mamba_model_reports_optional_dependency(self):
        module = _load_model_builder()
        missing = ModuleNotFoundError(
            "No module named 'mamba_ssm'", name="mamba_ssm"
        )
        with mock.patch.object(
            module.importlib, "import_module", side_effect=missing
        ):
            with self.assertRaisesRegex(
                ModuleNotFoundError, "optional Mamba backend"
            ):
                module.build_model({"model": "MambaA"}, "cuda")


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