"""
LLM Provider Registry.

Manages registration and instantiation of LLM providers.
Provides a unified entry point for getting LLM instances.
"""

from __future__ import annotations

import logging
from typing import Optional, Any, Type

from argus.core.config import get_config
from argus.core.llm.base import BaseLLM

logger = logging.getLogger(__name__)


class LLMRegistry:
    """
    Registry for LLM providers.
    
    Maintains a mapping of provider names to their implementation classes.
    Supports dynamic registration of custom providers.
    
    Example:
        >>> from argus.core.llm import LLMRegistry, get_llm
        >>> 
        >>> # Get default provider
        >>> llm = get_llm()
        >>> 
        >>> # Get specific provider
        >>> llm = get_llm("anthropic", model="claude-3-5-sonnet-20241022")
        >>> 
        >>> # List available providers
        >>> providers = LLMRegistry.list_providers()
    """
    
    _providers: dict[str, Type[BaseLLM]] = {}
    _instances: dict[str, BaseLLM] = {}
    
    @classmethod
    def register(cls, name: str, provider_class: Type[BaseLLM]) -> None:
        """
        Register an LLM provider.
        
        Args:
            name: Provider name (e.g., "openai", "anthropic")
            provider_class: Provider implementation class
        """
        cls._providers[name.lower()] = provider_class
        logger.debug(f"Registered LLM provider: {name}")
    
    @classmethod
    def get(
        cls,
        provider: Optional[str] = None,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        cache: bool = True,
        **kwargs: Any,
    ) -> BaseLLM:
        """
        Get an LLM provider instance.
        
        Args:
            provider: Provider name (default from config)
            model: Model name (default from config or provider default)
            api_key: API key (default from config)
            cache: Whether to cache and reuse instances
            **kwargs: Additional provider-specific options
            
        Returns:
            BaseLLM instance
            
        Raises:
            ValueError: If provider is not registered
        """
        config = get_config()
        
        # Use default provider if not specified
        provider = (provider or config.default_provider).lower()
        
        # Get API key from config if not provided
        if api_key is None:
            if provider == "openai":
                api_key = config.llm.openai_api_key
            elif provider == "anthropic":
                api_key = config.llm.anthropic_api_key
            elif provider == "gemini":
                api_key = config.llm.google_api_key
        
        # Get model from config if not provided
        if model is None:
            model = config.get_model_for_provider(provider)
        
        # Check cache
        cache_key = f"{provider}:{model}"
        if cache and cache_key in cls._instances:
            return cls._instances[cache_key]
        
        # Get provider class
        if provider not in cls._providers:
            raise ValueError(
                f"Unknown LLM provider: '{provider}'. "
                f"Available providers: {list(cls._providers.keys())}"
            )
        
        provider_class = cls._providers[provider]
        
        # Build kwargs
        init_kwargs: dict[str, Any] = {
            "model": model,
            "temperature": config.temperature,
            "max_tokens": config.max_tokens,
            **kwargs,
        }
        
        # Add API key if available
        if api_key:
            init_kwargs["api_key"] = api_key
        
        # Add provider-specific config
        if provider == "ollama":
            init_kwargs["host"] = config.llm.ollama_host
        
        # Instantiate provider
        instance = provider_class(**init_kwargs)
        
        # Cache if requested
        if cache:
            cls._instances[cache_key] = instance
        
        return instance
    
    @classmethod
    def list_providers(cls) -> list[str]:
        """
        List registered provider names.
        
        Returns:
            List of provider names
        """
        return list(cls._providers.keys())
    
    @classmethod
    def has_provider(cls, name: str) -> bool:
        """
        Check if a provider is registered.
        
        Args:
            name: Provider name
            
        Returns:
            True if provider is registered
        """
        return name.lower() in cls._providers
    
    @classmethod
    def clear_cache(cls) -> None:
        """Clear the instance cache."""
        cls._instances.clear()


# Register default providers
def _register_default_providers() -> None:
    """Register the built-in LLM providers."""
    # Core providers
    from argus.core.llm.openai import OpenAILLM
    from argus.core.llm.anthropic import AnthropicLLM
    from argus.core.llm.gemini import GeminiLLM
    from argus.core.llm.ollama import OllamaLLM
    from argus.core.llm.cohere import CohereLLM
    from argus.core.llm.mistral import MistralLLM
    from argus.core.llm.groq import GroqLLM
    
    # OpenAI-compatible providers
    from argus.core.llm.deepseek import DeepSeekLLM
    from argus.core.llm.xai import XaiLLM
    from argus.core.llm.perplexity import PerplexityLLM
    from argus.core.llm.nvidia import NvidiaLLM
    from argus.core.llm.together import TogetherLLM
    from argus.core.llm.fireworks import FireworksLLM
    
    # Cloud providers
    from argus.core.llm.bedrock import BedrockLLM
    from argus.core.llm.azure_openai import AzureOpenAILLM
    from argus.core.llm.vertex import VertexAILLM
    from argus.core.llm.huggingface import HuggingFaceLLM
    
    # Enterprise providers
    from argus.core.llm.watsonx import WatsonxLLM
    from argus.core.llm.databricks import DatabricksLLM
    from argus.core.llm.snowflake import SnowflakeLLM
    from argus.core.llm.sambanova import SambanovaLLM
    from argus.core.llm.cerebras import CerebrasLLM
    
    # Utility/self-hosted providers
    from argus.core.llm.litellm import LiteLLMLLM
    from argus.core.llm.cloudflare import CloudflareLLM
    from argus.core.llm.replicate import ReplicateLLM
    from argus.core.llm.vllm import VllmLLM
    from argus.core.llm.llamacpp import LlamaCppLLM
    
    # Register core providers
    LLMRegistry.register("openai", OpenAILLM)
    LLMRegistry.register("anthropic", AnthropicLLM)
    LLMRegistry.register("gemini", GeminiLLM)
    LLMRegistry.register("ollama", OllamaLLM)
    LLMRegistry.register("cohere", CohereLLM)
    LLMRegistry.register("mistral", MistralLLM)
    LLMRegistry.register("groq", GroqLLM)
    
    # Register OpenAI-compatible providers
    LLMRegistry.register("deepseek", DeepSeekLLM)
    LLMRegistry.register("xai", XaiLLM)
    LLMRegistry.register("grok", XaiLLM)  # Alias
    LLMRegistry.register("perplexity", PerplexityLLM)
    LLMRegistry.register("nvidia", NvidiaLLM)
    LLMRegistry.register("together", TogetherLLM)
    LLMRegistry.register("fireworks", FireworksLLM)
    
    # Register cloud providers
    LLMRegistry.register("bedrock", BedrockLLM)
    LLMRegistry.register("azure", AzureOpenAILLM)
    LLMRegistry.register("azure_openai", AzureOpenAILLM)
    LLMRegistry.register("vertex", VertexAILLM)
    LLMRegistry.register("vertex_ai", VertexAILLM)
    LLMRegistry.register("huggingface", HuggingFaceLLM)
    LLMRegistry.register("hf", HuggingFaceLLM)  # Alias
    
    # Register enterprise providers
    LLMRegistry.register("watsonx", WatsonxLLM)
    LLMRegistry.register("ibm", WatsonxLLM)  # Alias
    LLMRegistry.register("databricks", DatabricksLLM)
    LLMRegistry.register("snowflake", SnowflakeLLM)
    LLMRegistry.register("sambanova", SambanovaLLM)
    LLMRegistry.register("cerebras", CerebrasLLM)
    
    # Register utility/self-hosted providers
    LLMRegistry.register("litellm", LiteLLMLLM)
    LLMRegistry.register("cloudflare", CloudflareLLM)
    LLMRegistry.register("replicate", ReplicateLLM)
    LLMRegistry.register("vllm", VllmLLM)
    LLMRegistry.register("llamacpp", LlamaCppLLM)
    LLMRegistry.register("llama.cpp", LlamaCppLLM)  # Alias


# Auto-register on import
_register_default_providers()


# Convenience functions
def get_llm(
    provider: Optional[str] = None,
    model: Optional[str] = None,
    **kwargs: Any,
) -> BaseLLM:
    """
    Get an LLM provider instance.
    
    Convenience function that delegates to LLMRegistry.get().
    
    Args:
        provider: Provider name (default from ARGUS_DEFAULT_PROVIDER)
        model: Model name (default from ARGUS_DEFAULT_MODEL)
        **kwargs: Additional options
        
    Returns:
        BaseLLM instance
        
    Example:
        >>> from argus.core.llm import get_llm
        >>> 
        >>> # Use default provider (from config)
        >>> llm = get_llm()
        >>> 
        >>> # Use specific provider and model
        >>> llm = get_llm("openai", model="gpt-4o")
        >>> 
        >>> # Generate response
        >>> response = llm.generate("What is 2+2?")
        >>> print(response.content)
    """
    return LLMRegistry.get(provider=provider, model=model, **kwargs)


def register_provider(name: str, provider_class: Type[BaseLLM]) -> None:
    """
    Register a custom LLM provider.
    
    Args:
        name: Provider name
        provider_class: Provider implementation class
        
    Example:
        >>> from argus.core.llm import register_provider, BaseLLM
        >>> 
        >>> class MyCustomLLM(BaseLLM):
        ...     # Implementation
        ...     pass
        >>> 
        >>> register_provider("custom", MyCustomLLM)
    """
    LLMRegistry.register(name, provider_class)


def list_providers() -> list[str]:
    """
    List available LLM providers.
    
    Returns:
        List of provider names
        
    Example:
        >>> from argus.core.llm import list_providers
        >>> print(list_providers())
        ['openai', 'anthropic', 'gemini', 'ollama']
    """
    return LLMRegistry.list_providers()
