GPT Image MCP Server

"""Provider registry for managing multiple LLM providers."""

import asyncio
import logging
from typing import Any

from .base import LLMProvider, ProviderError

logger = logging.getLogger(__name__)


class ProviderRegistry:
    """Registry for managing LLM providers and model routing."""

    def __init__(self):
        self._providers: dict[str, LLMProvider] = {}
        self._model_to_provider: dict[str, str] = {}
        self._registry_lock = asyncio.Lock()

    async def register_provider(self, provider: LLMProvider) -> None:
        """Register a new provider.

        Args:
            provider: The provider instance to register

        Raises:
            ProviderError: If provider registration fails
        """
        if not provider.is_available():
            logger.warning(
                f"Provider {provider.name} is not available, skipping registration"
            )
            return

        async with self._registry_lock:
            # Check for conflicting provider names
            if provider.name in self._providers:
                raise ProviderError(
                    f"Provider '{provider.name}' is already registered",
                    provider_name=provider.name,
                    error_code="DUPLICATE_PROVIDER",
                )

            # Register the provider
            self._providers[provider.name] = provider

            # Map all supported models to this provider
            supported_models = provider.get_supported_models()
            for model_id in supported_models:
                if model_id in self._model_to_provider:
                    existing_provider = self._model_to_provider[model_id]
                    logger.warning(
                        f"Model '{model_id}' is already registered with provider "
                        f"'{existing_provider}'. "
                        f"Overriding with '{provider.name}'"
                    )

                self._model_to_provider[model_id] = provider.name

            logger.info(
                f"Registered provider '{provider.name}' with "
                f"{len(supported_models)} models"
            )

    async def unregister_provider(self, provider_name: str) -> None:
        """Unregister a provider.

        Args:
            provider_name: Name of the provider to unregister
        """
        async with self._registry_lock:
            if provider_name not in self._providers:
                logger.warning(f"Provider '{provider_name}' is not registered")
                return

            self._providers.pop(provider_name)

            # Remove model mappings
            models_to_remove = []
            for model_id, mapped_provider in self._model_to_provider.items():
                if mapped_provider == provider_name:
                    models_to_remove.append(model_id)

            for model_id in models_to_remove:
                del self._model_to_provider[model_id]

            logger.info(
                f"Unregistered provider '{provider_name}' and "
                f"{len(models_to_remove)} models"
            )

    def get_provider(self, provider_name: str) -> LLMProvider | None:
        """Get a provider by name.

        Args:
            provider_name: Name of the provider

        Returns:
            Provider instance or None if not found
        """
        return self._providers.get(provider_name)

    def get_provider_for_model(self, model_id: str) -> LLMProvider | None:
        """Get the provider that supports a specific model.

        Args:
            model_id: The model ID to look up

        Returns:
            Provider instance or None if model is not supported
        """
        provider_name = self._model_to_provider.get(model_id)
        if provider_name:
            return self._providers.get(provider_name)
        return None

    def get_all_providers(self) -> list[LLMProvider]:
        """Get all registered providers.

        Returns:
            List of all registered providers
        """
        return list(self._providers.values())

    def get_available_providers(self) -> list[LLMProvider]:
        """Get all available (enabled and properly configured) providers.

        Returns:
            List of available providers
        """
        return [
            provider for provider in self._providers.values() if provider.is_available()
        ]

    def get_supported_models(self) -> set[str]:
        """Get all supported models across all providers.

        Returns:
            Set of all supported model IDs
        """
        return set(self._model_to_provider.keys())

    def get_models_by_provider(self) -> dict[str, set[str]]:
        """Get models grouped by provider.

        Returns:
            Dict mapping provider names to their supported models
        """
        result = {}
        for provider_name, provider in self._providers.items():
            result[provider_name] = provider.get_supported_models()
        return result

    def is_model_supported(self, model_id: str) -> bool:
        """Check if a model is supported by any provider.

        Args:
            model_id: The model ID to check

        Returns:
            True if the model is supported
        """
        return model_id in self._model_to_provider

    def get_model_info(self, model_id: str) -> dict[str, Any] | None:
        """Get detailed information about a model.

        Args:
            model_id: The model ID

        Returns:
            Dict containing model information or None if not found
        """
        provider = self.get_provider_for_model(model_id)
        if not provider:
            return None

        capabilities = provider.get_model_capabilities(model_id)
        if not capabilities:
            return None

        return {
            "model_id": model_id,
            "provider": provider.name,
            "capabilities": capabilities,
            "is_available": provider.is_available(),
        }

    def validate_model_request(
        self, model_id: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Validate a model request and normalize parameters.

        Args:
            model_id: The model ID
            params: Request parameters

        Returns:
            Normalized parameters

        Raises:
            ProviderError: If validation fails
        """
        provider = self.get_provider_for_model(model_id)
        if not provider:
            raise ProviderError(
                f"Model '{model_id}' is not supported by any registered provider",
                provider_name="registry",
                error_code="UNSUPPORTED_MODEL",
            )

        if not provider.is_available():
            raise ProviderError(
                f"Provider '{provider.name}' for model '{model_id}' is not available",
                provider_name=provider.name,
                error_code="PROVIDER_UNAVAILABLE",
            )

        return provider.validate_model_params(model_id, params)

    def get_default_model(self, provider_name: str | None = None) -> str | None:
        """Get a default model, optionally from a specific provider.

        Args:
            provider_name: Optional provider name to get default from

        Returns:
            Default model ID or None if no models available
        """
        if provider_name:
            provider = self.get_provider(provider_name)
            if provider and provider.is_available():
                models = provider.get_supported_models()
                return next(iter(models)) if models else None
        else:
            # Get default from any available provider
            for provider in self.get_available_providers():
                models = provider.get_supported_models()
                if models:
                    return next(iter(models))

        return None

    def get_registry_stats(self) -> dict[str, Any]:
        """Get registry statistics.

        Returns:
            Dict containing registry statistics
        """
        available_providers = self.get_available_providers()

        return {
            "total_providers": len(self._providers),
            "available_providers": len(available_providers),
            "total_models": len(self._model_to_provider),
            "providers": {
                name: {
                    "available": provider.is_available(),
                    "models": list(provider.get_supported_models()),
                }
                for name, provider in self._providers.items()
            },
        }

    def __str__(self) -> str:
        return (
            f"ProviderRegistry(providers={len(self._providers)}, "
            f"models={len(self._model_to_provider)})"
        )

    def __repr__(self) -> str:
        return (
            f"ProviderRegistry(providers={list(self._providers.keys())}, "
            f"models={len(self._model_to_provider)})"
        )