from __future__ import annotations

from typing import Callable, Optional, Dict

from diona.project.agent.model import BaseModelProvider


class ModelProviderInfo:
    def __init__(
        self,
        name: str,
        provider_class: type,
        create_func: Callable[[], BaseModelProvider],
        description: str = ""
    ):
        self.name = name
        self.provider_class = provider_class
        self.create_func = create_func
        self.description = description


class ModelProviderManager:
    _instance: Optional[ModelProviderManager] = None

    def __new__(cls) -> ModelProviderManager:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return

        self._providers: Dict[str, ModelProviderInfo] = {}
        self._current_provider_name: str = "default"
        self._current_provider: Optional[BaseModelProvider] = None
        self._initialized = True

        self._register_default_providers()

    def _register_default_providers(self):
        self.register(
            name="default",
            provider_class=BaseModelProvider,
            create_func=lambda: BaseModelProvider.get_default(),
            description="Default model provider using environment variables"
        )
        self.register(
            name="deepseek-chat",
            provider_class=BaseModelProvider,
            create_func=lambda: BaseModelProvider.get_default(),
            description="Default model provider using environment variables"
        )
        self._current_provider = BaseModelProvider.get_default()

    def register(
        self,
        name: str,
        provider_class: type,
        create_func: Callable[[], BaseModelProvider],
        description: str = ""
    ) -> None:
        self._providers[name] = ModelProviderInfo(
            name=name,
            provider_class=provider_class,
            create_func=create_func,
            description=description
        )

    def unregister(self, name: str) -> bool:
        if name in self._providers:
            del self._providers[name]
            return True
        return False

    def get(self, name: str) -> Optional[ModelProviderInfo]:
        return self._providers.get(name)

    def list_all(self) -> list[ModelProviderInfo]:
        return list(self._providers.values())

    def list_names(self) -> list[str]:
        return list(self._providers.keys())

    def set_current(self, name: str) -> bool:
        if name not in self._providers:
            return False

        self._current_provider_name = name
        self._current_provider = self._providers[name].create_func()
        return True

    def get_current(self) -> BaseModelProvider:
        if self._current_provider is None:
            self._current_provider = BaseModelProvider.get_default()
        return self._current_provider

    def get_current_name(self) -> str:
        return self._current_provider_name


_manager: Optional[ModelProviderManager] = None


def get_model_provider_manager() -> ModelProviderManager:
    global _manager
    if _manager is None:
        _manager = ModelProviderManager()
    return _manager
