#!/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 predictor 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
from vi.inference.predictors.base_predictor import BasePredictor


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

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

    """

    module: str
    class_name: str


class PredictorRegistry:
    """Registry for model predictors with plugin-style architecture.

    Maps loader types to appropriate predictor classes. New predictors can be
    added by decorating classes with @PredictorRegistry.register().

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

    Example:
        ```python
        from vi.inference.predictors import BasePredictor, PredictorRegistry


        @PredictorRegistry.register(
            predictor_key="my_model", loader_types=["MyModelLoader"]
        )
        class MyModelPredictor(BasePredictor):
            def __call__(self, image_path, **kwargs):
                # Model-specific inference logic
                pass
        ```

    """

    _registry: dict[str, type[BasePredictor]] = {}
    _loader_mapping: dict[str, str] = {}
    _metadata: dict[str, PredictorMetadata] = {}
    _auto_discovered: bool = False

    @classmethod
    def register(
        cls,
        predictor_key: str,
        loader_types: list[str] | None = None,
    ) -> Callable[[type[BasePredictor]], type[BasePredictor]]:
        """Register a model predictor.

        Registers a predictor class with the registry, associating it with
        specific loader types for automatic selection.

        Args:
            predictor_key: Unique identifier for this predictor (e.g., "qwen25vl").
            loader_types: List of loader class names that this predictor handles
                (e.g., ["Qwen25VLLoader"]).

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

        Example:
            ```python
            @PredictorRegistry.register(
                predictor_key="qwenvl", loader_types=["Qwen25VLLoader"]
            )
            class QwenVLPredictor(BasePredictor):
                pass
            ```

        """

        def decorator(predictor_class: type[BasePredictor]) -> type[BasePredictor]:
            # Register the predictor
            cls._registry[predictor_key] = predictor_class

            # Map loader types to predictor key
            if loader_types:
                for loader_type in loader_types:
                    cls._loader_mapping[loader_type] = predictor_key

            # Store metadata for lazy loading
            cls._metadata[predictor_key] = PredictorMetadata(
                module=predictor_class.__module__,
                class_name=predictor_class.__name__,
            )

            return predictor_class

        return decorator

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

        Scans the architectures directory and imports all predictor 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
        predictors_module = importlib.import_module(
            "vi.inference.predictors.architectures"
        )
        if predictors_module.__file__ is None:
            # Module is a namespace package or __file__ is not available
            # Try to get path from __path__ attribute
            if hasattr(predictors_module, "__path__"):
                architectures_path = Path(predictors_module.__path__[0])
            else:
                # Cannot determine path, skip auto-discovery
                return
        else:
            architectures_path = Path(predictors_module.__file__).parent

        # Import all Python files in the architectures directory
        for item in architectures_path.iterdir():
            if item.is_file() and item.suffix == ".py" and item.stem != "__init__":
                # Import architecture module (e.g., qwenvl, deepseekocr, nvila)
                module_name = f"vi.inference.predictors.architectures.{item.stem}"
                try:
                    importlib.import_module(module_name)
                except ImportError:
                    # Silently skip modules with missing dependencies
                    pass

    @classmethod
    def _lazy_import_predictor(cls, predictor_key: str) -> type[BasePredictor]:
        """Lazily import a predictor class by key.

        Args:
            predictor_key: The predictor key to import.

        Returns:
            The predictor class.

        Raises:
            ValueError: If the predictor key is unknown.

        """
        if predictor_key not in cls._metadata:
            raise ValueError(
                f"Unknown predictor: '{predictor_key}'. "
                f"Available predictors: {cls.list_predictors()}"
            )

        metadata = cls._metadata[predictor_key]

        # Import the predictor
        module_path = metadata.module
        class_name = metadata.class_name
        module = importlib.import_module(module_path)
        predictor_class = getattr(module, class_name)

        # Cache in registry for subsequent calls
        cls._registry[predictor_key] = predictor_class

        return predictor_class

    @classmethod
    def get_predictor(
        cls,
        predictor_key: str | None = None,
        loader: BaseLoader | None = None,
    ) -> type[BasePredictor]:
        """Get predictor class by key or from loader instance.

        Retrieves the appropriate predictor class based on the provided identifier
        or by detecting the loader type. Predictors are imported lazily only when
        requested.

        Args:
            predictor_key: Direct predictor key (e.g., "qwen25vl").
            loader: Loader instance to find matching predictor for.

        Returns:
            Predictor class that can be instantiated with a loader.

        Raises:
            ValueError: If no arguments provided, both arguments provided,
                or if the specified predictor/loader is not found.

        Example:
            ```python
            # Get by predictor key
            predictor_class = PredictorRegistry.get_predictor(predictor_key="qwen25vl")

            # Get by loader instance
            loader = Qwen25VLLoader(model_meta)
            predictor_class = PredictorRegistry.get_predictor(loader=loader)
            ```

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

        if args_provided == 0:
            raise ValueError("Must provide either predictor_key or loader")

        if args_provided > 1:
            raise ValueError("Provide either predictor_key or loader, not both")

        cls._discover_architectures()

        # Resolve to predictor_key
        if loader:
            loader_type = type(loader).__name__
            if loader_type not in cls._loader_mapping:
                raise ValueError(
                    f"No predictor registered for loader type '{loader_type}'. "
                    f"Supported loader types: {cls.list_loader_types()}"
                )
            predictor_key = cls._loader_mapping[loader_type]

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

        return cls._lazy_import_predictor(predictor_key)

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

        Returns:
            List of predictor keys that can be used with get_predictor().

        Example:
            ```python
            available = PredictorRegistry.list_predictors()
            print(f"Available predictors: {available}")
            ```

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

    @classmethod
    def list_loader_types(cls) -> list[str]:
        """List all loader types that have registered predictors.

        Returns:
            List of loader class names that have associated predictors.

        Example:
            ```python
            types = PredictorRegistry.list_loader_types()
            print(f"Supported loader types: {types}")
            # Output: ['Qwen25VLLoader', 'NVILALoader', ...]
            ```

        """
        # Return all registered loader types
        return sorted(cls._loader_mapping.keys())

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

        Args:
            predictor_key: Predictor key to check.
            loader: Loader instance to check.

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

        Example:
            ```python
            if PredictorRegistry.is_registered(predictor_key="qwen25vl"):
                print("Qwen2.5-VL predictor is available")
            ```

        """
        if predictor_key:
            return predictor_key in cls._registry or predictor_key in cls._metadata
        if loader:
            loader_type = type(loader).__name__
            return loader_type in cls._loader_mapping
        return False
