from typing import List, Set, Optional, Dict
import logging
import inspect

from autobyteus.llm.autobyteus_provider import AutobyteusModelProvider
from autobyteus.llm.models import LLMModel, ModelInfo, ProviderModelGroup
from autobyteus.llm.providers import LLMProvider
from autobyteus.llm.runtimes import LLMRuntime
from autobyteus.llm.utils.llm_config import LLMConfig, TokenPricingConfig
from autobyteus.llm.base_llm import BaseLLM
from autobyteus.utils.parameter_schema import ParameterSchema, ParameterDefinition, ParameterType

from autobyteus.llm.api.claude_llm import ClaudeLLM
from autobyteus.llm.api.mistral_llm import MistralLLM
from autobyteus.llm.api.openai_llm import OpenAILLM
from autobyteus.llm.api.deepseek_llm import DeepSeekLLM
from autobyteus.llm.api.gemini_llm import GeminiLLM
from autobyteus.llm.api.grok_llm import GrokLLM
from autobyteus.llm.api.kimi_llm import KimiLLM
from autobyteus.llm.api.qwen_llm import QwenLLM
from autobyteus.llm.api.zhipu_llm import ZhipuLLM
from autobyteus.llm.api.minimax_llm import MinimaxLLM
from autobyteus.llm.ollama_provider import OllamaModelProvider
from autobyteus.llm.lmstudio_provider import LMStudioModelProvider
from autobyteus.utils.singleton import SingletonMeta

logger = logging.getLogger(__name__)

class LLMFactory(metaclass=SingletonMeta):
    _models_by_provider: Dict[LLMProvider, List[LLMModel]] = {}
    _models_by_identifier: Dict[str, LLMModel] = {}
    _initialized = False

    @staticmethod
    def ensure_initialized():
        """Ensures the factory is initialized before use."""
        if not LLMFactory._initialized:
            LLMFactory._initialize_registry()
            LLMFactory._initialized = True

    @staticmethod
    def reinitialize():
        """Reinitializes the model registry."""
        logger.info("Reinitializing LLM model registry...")
        LLMFactory._initialized = False
        LLMFactory._models_by_provider.clear()
        LLMFactory._models_by_identifier.clear()
        LLMFactory.ensure_initialized()
        logger.info("LLM model registry reinitialized successfully.")

    @staticmethod
    def _initialize_registry():
        """Initializes the registry with built-in models and discovers runtime models."""
        # Hardcoded direct-API models. Runtime defaults to API.
        supported_models = [
            # OPENAI Provider Models (latest generation only)
            LLMModel(
                name="gpt-5.2",
                value="gpt-5.2",
                provider=LLMProvider.OPENAI,
                llm_class=OpenAILLM,
                canonical_name="gpt-5.2",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(1.75, 14.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="reasoning_effort",
                        param_type=ParameterType.ENUM,
                        description="Controls how hard the model thinks. Higher effort improves quality but can increase latency and cost.",
                        required=False,
                        default_value="none",
                        enum_values=["none", "low", "medium", "high", "xhigh"]
                    ),
                    ParameterDefinition(
                        name="reasoning_summary",
                        param_type=ParameterType.ENUM,
                        description="Include a reasoning summary in the response when supported.",
                        required=False,
                        default_value="none",
                        enum_values=["none", "auto", "concise", "detailed"]
                    )
                ])
            ),
            LLMModel(
                name="gpt-5.2-chat-latest",
                value="gpt-5.2-chat-latest",
                provider=LLMProvider.OPENAI,
                llm_class=OpenAILLM,
                canonical_name="gpt-5.2-chat-latest",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(1.75, 14.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="reasoning_effort",
                        param_type=ParameterType.ENUM,
                        description="Controls how hard the model thinks. Higher effort improves quality but can increase latency and cost.",
                        required=False,
                        default_value="none",
                        enum_values=["none", "low", "medium", "high", "xhigh"]
                    ),
                    ParameterDefinition(
                        name="reasoning_summary",
                        param_type=ParameterType.ENUM,
                        description="Include a reasoning summary in the response when supported.",
                        required=False,
                        default_value="none",
                        enum_values=["none", "auto", "concise", "detailed"]
                    )
                ])
            ),
            # MISTRAL Provider Models
            LLMModel(
                name="mistral-large",
                value="mistral-large-latest",
                provider=LLMProvider.MISTRAL,
                llm_class=MistralLLM,
                canonical_name="mistral-large",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(2.00, 6.00)
                )
            ),
            LLMModel(
                name="devstral-2",
                value="devstral-2512",
                provider=LLMProvider.MISTRAL,
                llm_class=MistralLLM,
                canonical_name="devstral-2",
                default_config=LLMConfig(
                    # Pricing from Mistral launch: $0.40 input / $2.00 output per MTokens.
                    pricing_config=TokenPricingConfig(0.40, 2.00)
                )
            ),
            # GROK Provider Models (latest flagship + fast)
            LLMModel(
                name="grok-4",
                value="grok-4",
                provider=LLMProvider.GROK,
                llm_class=GrokLLM,
                canonical_name="grok-4",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(3.00, 15.00)
                )
            ),
            LLMModel(
                name="grok-4-1-fast-reasoning",
                value="grok-4-1-fast-reasoning",
                provider=LLMProvider.GROK,
                llm_class=GrokLLM,
                canonical_name="grok-4-1-fast-reasoning",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.20, 0.50)
                )
            ),
            LLMModel(
                name="grok-4-1-fast-non-reasoning",
                value="grok-4-1-fast-non-reasoning",
                provider=LLMProvider.GROK,
                llm_class=GrokLLM,
                canonical_name="grok-4-1-fast-non-reasoning",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.20, 0.50)
                )
            ),
            LLMModel(
                name="grok-code-fast-1",
                value="grok-code-fast-1",
                provider=LLMProvider.GROK,
                llm_class=GrokLLM,
                canonical_name="grok-code-fast-1",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.20, 1.50)
                )
            ),
            # ANTHROPIC Provider Models
            LLMModel(
                name="claude-4.5-opus",
                value="claude-opus-4-5-20251101",
                provider=LLMProvider.ANTHROPIC,
                llm_class=ClaudeLLM,
                canonical_name="claude-4.5-opus",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(5.00, 25.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_enabled",
                        param_type=ParameterType.BOOLEAN,
                        description="Enable extended thinking summaries in Claude responses",
                        required=False,
                        default_value=False
                    ),
                    ParameterDefinition(
                        name="thinking_budget_tokens",
                        param_type=ParameterType.INTEGER,
                        description="Token budget for extended thinking (min 1024)",
                        required=False,
                        default_value=1024,
                        min_value=1024
                    )
                ])
            ),
            LLMModel(
                name="claude-4.5-sonnet",
                value="claude-sonnet-4-5-20250929",
                provider=LLMProvider.ANTHROPIC,
                llm_class=ClaudeLLM,
                canonical_name="claude-4.5-sonnet",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(3.00, 15.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_enabled",
                        param_type=ParameterType.BOOLEAN,
                        description="Enable extended thinking summaries in Claude responses",
                        required=False,
                        default_value=False
                    ),
                    ParameterDefinition(
                        name="thinking_budget_tokens",
                        param_type=ParameterType.INTEGER,
                        description="Token budget for extended thinking (min 1024)",
                        required=False,
                        default_value=1024,
                        min_value=1024
                    )
                ])
            ),
            LLMModel(
                name="claude-4.5-haiku",
                value="claude-haiku-4-5-20251001",
                provider=LLMProvider.ANTHROPIC,
                llm_class=ClaudeLLM,
                canonical_name="claude-4.5-haiku",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(1.00, 5.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_enabled",
                        param_type=ParameterType.BOOLEAN,
                        description="Enable extended thinking summaries in Claude responses",
                        required=False,
                        default_value=False
                    ),
                    ParameterDefinition(
                        name="thinking_budget_tokens",
                        param_type=ParameterType.INTEGER,
                        description="Token budget for extended thinking (min 1024)",
                        required=False,
                        default_value=1024,
                        min_value=1024
                    )
                ])
            ),
            # DEEPSEEK Provider Models
            LLMModel(
                name="deepseek-chat",
                value="deepseek-chat",
                provider=LLMProvider.DEEPSEEK,
                llm_class=DeepSeekLLM,
                canonical_name="deepseek-chat",
                default_config=LLMConfig(
                    rate_limit=60,
                    token_limit=8000,
                    pricing_config=TokenPricingConfig(0.014, 0.28)
                )
            ),
            # Adding deepseek-reasoner support
            LLMModel(
                name="deepseek-reasoner",
                value="deepseek-reasoner",
                provider=LLMProvider.DEEPSEEK,
                llm_class=DeepSeekLLM,
                canonical_name="deepseek-reasoner",
                default_config=LLMConfig(
                    rate_limit=60,
                    token_limit=8000,
                    pricing_config=TokenPricingConfig(0.14, 2.19)
                )
            ),
            # GEMINI Provider Models
            LLMModel(
                name="gemini-3-pro-preview",
                value="gemini-3-pro-preview",
                provider=LLMProvider.GEMINI,
                llm_class=GeminiLLM,
                canonical_name="gemini-3-pro",
                default_config=LLMConfig(
                    # Pricing from Gemini 3 Pro preview launch (per 1M tokens).
                    pricing_config=TokenPricingConfig(2.00, 12.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_level",
                        param_type=ParameterType.ENUM,
                        description="How deeply the model should reason before responding",
                        required=False,
                        default_value="minimal",
                        enum_values=["minimal", "low", "medium", "high"]
                    ),
                    ParameterDefinition(
                        name="include_thoughts",
                        param_type=ParameterType.BOOLEAN,
                        description="Include model thought summaries in responses",
                        required=False,
                        default_value=False
                    )
                ])
            ),
            LLMModel(
                name="gemini-3-flash-preview",
                value="gemini-3-flash-preview",
                provider=LLMProvider.GEMINI,
                llm_class=GeminiLLM,
                canonical_name="gemini-3-flash",
                default_config=LLMConfig(
                    # Pricing from Gemini 3 Flash preview launch (per 1M tokens).
                    pricing_config=TokenPricingConfig(0.50, 3.00)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_level",
                        param_type=ParameterType.ENUM,
                        description="How deeply the model should reason before responding",
                        required=False,
                        default_value="minimal",
                        enum_values=["minimal", "low", "medium", "high"]
                    ),
                    ParameterDefinition(
                        name="include_thoughts",
                        param_type=ParameterType.BOOLEAN,
                        description="Include model thought summaries in responses",
                        required=False,
                        default_value=False
                    )
                ])
            ),
            # KIMI Provider Models
            LLMModel(
                name="kimi-k2-0711-preview",
                value="kimi-k2-0711-preview",
                provider=LLMProvider.KIMI,
                llm_class=KimiLLM,
                canonical_name="kimi-k2-0711-preview",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.55, 2.21)
                )
            ),
            LLMModel(
                name="kimi-k2-0905-preview",
                value="kimi-k2-0905-preview",
                provider=LLMProvider.KIMI,
                llm_class=KimiLLM,
                canonical_name="kimi-k2-0905-preview",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.55, 2.21)
                )
            ),
            LLMModel(
                name="kimi-k2-turbo-preview",
                value="kimi-k2-turbo-preview",
                provider=LLMProvider.KIMI,
                llm_class=KimiLLM,
                canonical_name="kimi-k2-turbo-preview",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(2.76, 2.76)
                )
            ),
            LLMModel(
                name="kimi-latest",
                value="kimi-latest",
                provider=LLMProvider.KIMI,
                llm_class=KimiLLM,
                canonical_name="kimi-latest",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(1.38, 4.14)
                )
            ),
            LLMModel(
                name="kimi-thinking-preview",
                value="kimi-thinking-preview",
                provider=LLMProvider.KIMI,
                llm_class=KimiLLM,
                canonical_name="kimi-thinking-preview",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(27.59, 27.59)
                )
            ),
            # QWEN Provider Models
            LLMModel(
                name="qwen3-max",
                value="qwen-max",
                provider=LLMProvider.QWEN,
                llm_class=QwenLLM,
                canonical_name="qwen3-max",
                default_config=LLMConfig(
                    token_limit=262144,
                    pricing_config=TokenPricingConfig(
                        input_token_pricing=2.4,
                        output_token_pricing=12.0
                    )
                )
            ),
            # ZHIPU Provider Models
            LLMModel(
                name="glm-4.7",
                value="glm-4.7",
                provider=LLMProvider.ZHIPU,
                llm_class=ZhipuLLM,
                canonical_name="glm-4.7",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(13.8, 13.8)
                ),
                config_schema=ParameterSchema(parameters=[
                    ParameterDefinition(
                        name="thinking_type",
                        param_type=ParameterType.ENUM,
                        description="Enable or disable deep thinking",
                        required=False,
                        default_value="enabled",
                        enum_values=["enabled", "disabled"]
                    )
                ])
            ),
            # MINIMAX Provider Models
            LLMModel(
                name="minimax-m2.1",
                value="MiniMax-M2.1",
                provider=LLMProvider.MINIMAX,
                llm_class=MinimaxLLM,
                canonical_name="minimax-m2.1",
                default_config=LLMConfig(
                    pricing_config=TokenPricingConfig(0.15, 0.45)
                )
            ),
        ]
        for model in supported_models:
            LLMFactory.register_model(model)

        # Discover models from runtimes
        OllamaModelProvider.discover_and_register()
        LMStudioModelProvider.discover_and_register()
        AutobyteusModelProvider.discover_and_register()

    @staticmethod
    def register_model(model: LLMModel):
        """Registers a new LLM model."""
        identifier = model.model_identifier
        if identifier in LLMFactory._models_by_identifier:
            logger.debug(f"Redefining model with identifier '{identifier}'.")
            # Remove old model from provider group to replace it
            old_model = LLMFactory._models_by_identifier[identifier]
            if old_model.provider in LLMFactory._models_by_provider:
                # This check is needed because a model might be in _models_by_identifier but not yet in _models_by_provider if re-registering
                if old_model in LLMFactory._models_by_provider[old_model.provider]:
                    LLMFactory._models_by_provider[old_model.provider].remove(old_model)

        LLMFactory._models_by_identifier[identifier] = model
        LLMFactory._models_by_provider.setdefault(model.provider, []).append(model)

    @staticmethod
    def create_llm(model_identifier: str, llm_config: Optional[LLMConfig] = None) -> BaseLLM:
        """
        Creates an LLM instance for the specified unique model identifier.
        Raises an error if the identifier is not found or if a non-unique name is provided.
        """
        LLMFactory.ensure_initialized()
        
        # First, try a direct lookup by the unique model_identifier
        model = LLMFactory._models_by_identifier.get(model_identifier)
        if model:
            return model.create_llm(llm_config)

        # If not found, check if the user provided a non-unique name by mistake
        found_by_name = [m for m in LLMFactory._models_by_identifier.values() if m.name == model_identifier]
        if len(found_by_name) > 1:
            identifiers = [m.model_identifier for m in found_by_name]
            raise ValueError(
                f"The model name '{model_identifier}' is ambiguous. Please use one of the unique "
                f"model identifiers: {identifiers}"
            )
        
        raise ValueError(f"Model with identifier '{model_identifier}' not found.")

    # --- New Public API ---

    @staticmethod
    def list_available_models() -> List[ModelInfo]:
        """Returns a list of all available models with their detailed info."""
        LLMFactory.ensure_initialized()
        models = sorted(LLMFactory._models_by_identifier.values(), key=lambda m: m.model_identifier)
        return [m.to_model_info() for m in models]

    @staticmethod
    def list_models_by_provider(provider: LLMProvider) -> List[ModelInfo]:
        """Returns a list of available models for a specific provider."""
        LLMFactory.ensure_initialized()
        provider_models = sorted(
            [m for m in LLMFactory._models_by_identifier.values() if m.provider == provider],
            key=lambda m: m.model_identifier
        )
        return [m.to_model_info() for m in provider_models]

    @staticmethod
    def list_models_by_runtime(runtime: LLMRuntime) -> List[ModelInfo]:
        """Returns a list of available models for a specific runtime."""
        LLMFactory.ensure_initialized()
        runtime_models = sorted(
            [m for m in LLMFactory._models_by_identifier.values() if m.runtime == runtime],
            key=lambda m: m.model_identifier
        )
        return [m.to_model_info() for m in runtime_models]

    @staticmethod
    def get_canonical_name(model_identifier: str) -> Optional[str]:
        """
        Retrieves the canonical name for a given model identifier.
        """
        LLMFactory.ensure_initialized()
        model = LLMFactory._models_by_identifier.get(model_identifier)
        if model:
            return model.canonical_name
        
        logger.warning(f"Could not find model with identifier '{model_identifier}' to get its canonical name.")
        return None

    @staticmethod
    def get_provider(model_identifier: str) -> Optional[LLMProvider]:
        """
        Retrieves the provider for a given model identifier.

        Returns:
            Optional[LLMProvider]: The provider when a single match exists; otherwise None.

        Raises:
            ValueError: If `model_identifier` resolves to multiple models by name.
        """
        LLMFactory.ensure_initialized()

        # Prefer exact identifier lookup.
        model = LLMFactory._models_by_identifier.get(model_identifier)
        if model:
            return model.provider

        # Fallback to name lookup for convenience.
        found_by_name = [m for m in LLMFactory._models_by_identifier.values() if m.name == model_identifier]
        if len(found_by_name) == 1:
            return found_by_name[0].provider
        if len(found_by_name) > 1:
            identifiers = [m.model_identifier for m in found_by_name]
            raise ValueError(
                f"The model name '{model_identifier}' is ambiguous. Please use one of the unique "
                f"model identifiers: {identifiers}"
            )

        logger.warning(f"Could not find model with identifier '{model_identifier}' to get its provider.")
        return None

    @staticmethod
    def reload_models(provider: LLMProvider) -> int:
        """
        Reloads models for a specific provider.
        Strategy: Clear-Then-Discover (Fail Fast).
        
        1. Clears all existing models for the provider.
        2. Fetches fresh models.
        3. Registers them.
        
        If step 2 fails, the registry for this provider remains empty, correctly reflecting 
        that the server is unreachable or returning no models.
        """
        LLMFactory.ensure_initialized()
        
        provider_handlers = {
            LLMProvider.LMSTUDIO: LMStudioModelProvider,
            LLMProvider.AUTOBYTEUS: AutobyteusModelProvider,
            LLMProvider.OLLAMA: OllamaModelProvider,
        }

        handler = provider_handlers.get(provider)
        if not handler:
            logger.warning(f"Reloading is not supported for provider: {provider}")
            return len(LLMFactory._models_by_provider.get(provider, []))

        # 1. Clear old models for this provider immediately
        current_provider_models = LLMFactory._models_by_provider.get(provider, [])
        ids_to_remove = [m.model_identifier for m in current_provider_models]
        
        logger.info(f"Clearing {len(ids_to_remove)} models for provider {provider} before discovery.")
        
        for mid in ids_to_remove:
            if mid in LLMFactory._models_by_identifier:
                del LLMFactory._models_by_identifier[mid]
        
        if provider in LLMFactory._models_by_provider:
            del LLMFactory._models_by_provider[provider]

        # 2. Fetch new models
        try:
            # We assume the handler has a static .get_models() method
            if not hasattr(handler, 'get_models'):
                logger.error(f"Provider handler {handler} does not implement get_models()")
                return 0 # We already cleared everything

            new_models: List[LLMModel] = handler.get_models()
        except Exception as e:
            logger.error(f"Failed to fetch models for {provider} during reload. Registry for this provider is now empty. Error: {e}")
            return 0

        # 3. Register new models
        logger.info(f"Registering {len(new_models)} new models for provider {provider}.")
        for model in new_models:
            LLMFactory.register_model(model)

        return len(new_models)

default_llm_factory = LLMFactory()
