"""
Model Registry - Maps model types to classifier classes.

Provides:
    - Registration of model classes
    - Factory function to create models from config
    - Model type detection from model names
"""

from typing import Dict, Type, Optional

from badas.core.base import VideoClassifierBase


# Registry mapping model_type -> classifier class
_MODEL_REGISTRY: Dict[str, Type[VideoClassifierBase]] = {}


def register_model(cls: Type[VideoClassifierBase]) -> Type[VideoClassifierBase]:
    """
    Register a model class in the registry.

    Can be used as a decorator:
        @register_model
        class MyClassifier(VideoClassifierBase):
            MODEL_TYPE = "my_model"
    """
    if cls.MODEL_TYPE is None:
        raise ValueError(f"{cls.__name__} must define MODEL_TYPE")

    _MODEL_REGISTRY[cls.MODEL_TYPE] = cls
    return cls


def get_model_class(model_type: str) -> Type[VideoClassifierBase]:
    """
    Get the classifier class for a model type.

    Args:
        model_type: Model type identifier (e.g., "vjepa2", "videomae")

    Returns:
        The classifier class

    Raises:
        KeyError: If model_type is not registered
    """
    if model_type not in _MODEL_REGISTRY:
        available = list(_MODEL_REGISTRY.keys())
        raise KeyError(
            f"Unknown model_type: '{model_type}'. "
            f"Available types: {available}"
        )
    return _MODEL_REGISTRY[model_type]


def create_model(config, model_type: Optional[str] = None) -> VideoClassifierBase:
    """
    Create a model instance from config.

    Args:
        config: Model configuration
        model_type: Explicit model type. If None, detected from config.model.model_name

    Returns:
        Instantiated model
    """
    if model_type is None:
        model_type = detect_model_type(config.model.model_name)

    cls = get_model_class(model_type)
    return cls(config)


def detect_model_type(model_name: str) -> str:
    """
    Detect model type from model name.

    Args:
        model_name: HuggingFace model name or path

    Returns:
        Model type string

    Raises:
        ValueError: If model type cannot be detected
    """
    name_lower = model_name.lower()

    if "vjepa2" in name_lower:
        return "vjepa2"

    if "videomae" in name_lower:
        return "videomae"

    raise ValueError(
        f"Cannot detect model type from name: '{model_name}'. "
        f"Please specify model_type explicitly."
    )


def list_registered_models() -> list:
    """Return list of registered model types."""
    return list(_MODEL_REGISTRY.keys())


# Auto-register built-in models
def _register_builtin_models():
    """Register built-in model classes."""
    from badas.core.vjepa2 import VJEPA2Classifier
    from badas.core.videomae import VideoMAEv2Classifier

    register_model(VJEPA2Classifier)
    register_model(VideoMAEv2Classifier)


_register_builtin_models()
