#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   registry.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK loader registry module.
"""

import importlib
from collections.abc import Callable
from pathlib import Path

from msgspec import Struct
from vi.inference.loaders.base_loader import BaseLoader


class LoaderMetadata(Struct, kw_only=True):
    """Metadata for lazy-loading a loader.

    Attributes:
        module: Python module path containing the loader class.
        class_name: Name of the loader class.

    """

    module: str
    class_name: str


class LoaderRegistry:
    """Registry for model loaders with plugin-style architecture.

    Allows registration and retrieval of model-specific loaders without tight
    coupling to specific implementations. New model loaders can be added by
    decorating classes with @LoaderRegistry.register().

    Supports lazy loading - loaders are only imported when actually needed,
    preventing unnecessary dependency checks at import time.

    Example:
        ```python
        from vi.inference.loaders import BaseLoader, LoaderRegistry


        @LoaderRegistry.register(
            loader_key="my_model",
            model_types=["my_model_type"],
            architectures=["MyModelForVision"],
        )
        class MyModelLoader(BaseLoader):
            def __init__(self, model_meta, **kwargs):
                # Model-specific loading logic
                pass
        ```

    """

    _registry: dict[str, type[BaseLoader]] = {}
    _model_type_mapping: dict[str, str] = {}
    _architecture_mapping: dict[str, str] = {}
    _metadata: dict[str, LoaderMetadata] = {}
    _auto_discovered: bool = False

    @classmethod
    def register(
        cls,
        loader_key: str,
        model_types: list[str] | None = None,
        architectures: list[str] | None = None,
    ) -> Callable[[type[BaseLoader]], type[BaseLoader]]:
        """Register a model loader.

        Registers a loader class with the registry, associating it with specific
        model types and architectures for automatic detection.

        Args:
            loader_key: Unique identifier for this loader (e.g., "qwen25vl").
            model_types: List of model_type values from config.json that this
                loader handles (e.g., ["qwen2_5_vl"]).
            architectures: List of architecture names from config.json that this
                loader handles (e.g., ["Qwen2_5_VLForConditionalGeneration"]).

        Returns:
            Decorator function that registers the class and returns it unchanged.

        Example:
            ```python
            @LoaderRegistry.register(
                loader_key="qwen25vl",
                model_types=["qwen2_5_vl"],
                architectures=["Qwen2_5_VLForConditionalGeneration"],
            )
            class Qwen25VLLoader(BaseLoader):
                pass
            ```

        """

        def decorator(loader_class: type[BaseLoader]) -> type[BaseLoader]:
            # Register the loader
            cls._registry[loader_key] = loader_class

            # Map model types to loader key
            if model_types:
                for model_type in model_types:
                    cls._model_type_mapping[model_type] = loader_key

            # Map architectures to loader key
            if architectures:
                for arch in architectures:
                    cls._architecture_mapping[arch] = loader_key

            # Store metadata for lazy loading
            cls._metadata[loader_key] = LoaderMetadata(
                module=loader_class.__module__,
                class_name=loader_class.__name__,
            )

            return loader_class

        return decorator

    @classmethod
    def _discover_architectures(cls) -> None:
        """Auto-discover and import all architecture modules.

        Scans the architectures directory and imports all loader modules to
        trigger their @register decorators. This is called lazily before
        accessing the registry.

        """
        if cls._auto_discovered:
            return

        cls._auto_discovered = True

        # Get the architectures directory
        loaders_module = importlib.import_module("vi.inference.loaders.architectures")
        if loaders_module.__file__ is None:
            # Module is a namespace package or __file__ is not available
            # Try to get path from __path__ attribute
            if hasattr(loaders_module, "__path__"):
                architectures_path = Path(loaders_module.__path__[0])
            else:
                # Cannot determine path, skip auto-discovery
                return
        else:
            architectures_path = Path(loaders_module.__file__).parent

        # Import all Python files in the architectures directory
        # Handles single files and directory structures
        for item in architectures_path.iterdir():
            if item.is_file() and item.suffix == ".py" and item.stem != "__init__":
                module_name = f"vi.inference.loaders.architectures.{item.stem}"
                try:
                    importlib.import_module(module_name)
                except ImportError:
                    # Silently skip modules with missing dependencies
                    pass
            elif item.is_dir() and not item.name.startswith("_"):
                module_name = f"vi.inference.loaders.architectures.{item.name}"
                try:
                    importlib.import_module(module_name)
                except ImportError:
                    # Silently skip modules with missing dependencies
                    pass

    @classmethod
    def _lazy_import_loader(cls, loader_key: str) -> type[BaseLoader]:
        """Lazily import a loader class by key.

        Args:
            loader_key: The loader key to import.

        Returns:
            The loader class.

        Raises:
            ValueError: If the loader key is unknown.

        """
        if loader_key not in cls._metadata:
            raise ValueError(
                f"Unknown loader: '{loader_key}'. "
                f"Available loaders: {cls.list_loaders()}"
            )

        metadata = cls._metadata[loader_key]

        # Import the loader
        module_path = metadata.module
        class_name = metadata.class_name
        module = importlib.import_module(module_path)
        loader_class = getattr(module, class_name)

        # Cache in registry for subsequent calls
        cls._registry[loader_key] = loader_class

        return loader_class

    @classmethod
    def get_loader(
        cls,
        loader_key: str | None = None,
        model_type: str | None = None,
        architecture: str | None = None,
    ) -> type[BaseLoader]:
        """Get a loader class by key, model type, or architecture.

        Retrieves the appropriate loader class based on the provided identifier.
        Useful for both explicit loader selection and automatic detection.
        Loaders are imported lazily only when requested.

        Args:
            loader_key: Direct loader key (e.g., "qwen25vl").
            model_type: Model type from config.json (e.g., "qwen2_5_vl").
            architecture: Architecture name from config.json (e.g.,
                "Qwen2_5_VLForConditionalGeneration").

        Returns:
            Loader class that can be instantiated with model metadata.

        Raises:
            ValueError: If no arguments provided, multiple arguments provided,
                or if the specified loader/model/architecture is not found.

        Example:
            ```python
            # Get by loader key
            loader_class = LoaderRegistry.get_loader(loader_key="qwen25vl")

            # Get by model type
            loader_class = LoaderRegistry.get_loader(model_type="qwen2_5_vl")

            # Get by architecture
            loader_class = LoaderRegistry.get_loader(
                architecture="Qwen2_5_VLForConditionalGeneration"
            )
            ```

        """
        # Count how many arguments were provided
        args_provided = sum(
            x is not None for x in [loader_key, model_type, architecture]
        )

        if args_provided == 0:
            raise ValueError(
                "Must provide one of: loader_key, model_type, or architecture"
            )

        if args_provided > 1:
            raise ValueError(
                "Provide only one of: loader_key, model_type, or architecture"
            )

        cls._discover_architectures()

        # Resolve to loader_key
        if model_type:
            if model_type not in cls._model_type_mapping:
                raise ValueError(
                    f"No loader registered for model_type '{model_type}'. "
                    f"Supported model types: {cls.list_model_types()}"
                )
            loader_key = cls._model_type_mapping[model_type]

        if architecture:
            if architecture not in cls._architecture_mapping:
                raise ValueError(
                    f"No loader registered for architecture '{architecture}'. "
                    f"Supported architectures: {cls.list_architectures()}"
                )
            loader_key = cls._architecture_mapping[architecture]

        # Get from cache or lazy import
        if loader_key in cls._registry:
            return cls._registry[loader_key]

        return cls._lazy_import_loader(loader_key)

    @classmethod
    def list_loaders(cls) -> list[str]:
        """List all registered loader keys.

        Returns:
            List of loader keys that can be used with get_loader().

        Example:
            ```python
            available = LoaderRegistry.list_loaders()
            print(f"Available loaders: {available}")
            ```

        """
        # Include both registered and metadata-defined loaders
        all_loaders = set(cls._registry.keys()) | set(cls._metadata.keys())
        return sorted(all_loaders)

    @classmethod
    def list_model_types(cls) -> list[str]:
        """List all supported model types from config.json.

        Returns:
            List of model_type values that can be automatically detected.

        Example:
            ```python
            types = LoaderRegistry.list_model_types()
            print(f"Supported model types: {types}")
            # Output: Supported model types: ['qwen2_5_vl', 'nvila', ...]
            ```

        """
        return sorted(cls._model_type_mapping.keys())

    @classmethod
    def list_architectures(cls) -> list[str]:
        """List all supported architecture names from config.json.

        Returns:
            List of architecture names that can be automatically detected.

        Example:
            ```python
            archs = LoaderRegistry.list_architectures()
            print(f"Supported architectures: {archs}")
            # Output: ['Qwen2_5_VLForConditionalGeneration', ...]
            ```

        """
        return sorted(cls._architecture_mapping.keys())

    @classmethod
    def is_registered(
        cls,
        loader_key: str | None = None,
        model_type: str | None = None,
        architecture: str | None = None,
    ) -> bool:
        """Check if a loader is registered for the given identifier.

        Args:
            loader_key: Loader key to check.
            model_type: Model type to check.
            architecture: Architecture name to check.

        Returns:
            True if a loader is registered, False otherwise.

        Example:
            ```python
            if LoaderRegistry.is_registered(model_type="qwen2_5_vl"):
                print("Qwen2.5-VL is supported")
            ```

        """
        if loader_key:
            return loader_key in cls._metadata
        if model_type:
            return model_type in cls._model_type_mapping
        if architecture:
            return architecture in cls._architecture_mapping
        return False

    @classmethod
    def get_loader_key(cls, loader_class: type[BaseLoader]) -> str | None:
        """Get the loader key for a given loader class.

        Args:
            loader_class: The loader class to look up.

        Returns:
            The loader key if found, None otherwise.

        Example:
            ```python
            from vi.inference.loaders import Qwen25VLLoader

            key = LoaderRegistry.get_loader_key(Qwen25VLLoader)
            print(key)  # Output: "qwen25vl"
            ```

        """
        # Check registered loaders first
        for key, registered_class in cls._registry.items():
            if registered_class == loader_class:
                return key

        # Check metadata for lazy-loaded loaders by class name
        class_name = loader_class.__name__
        for key, metadata in cls._metadata.items():
            if metadata.class_name == class_name:
                return key

        return None
