"""
Agent Factory for creating PydanticAI agents with different model providers.
"""

import os
from enum import Enum
from typing import Any, Dict, List, Optional, Type, TypeVar

import logfire
from pydantic import BaseModel
from pydantic_ai import Agent

from quber.settings import get_settings

_logfire_configured = False

OutputT = TypeVar("OutputT", bound=BaseModel)

#: Model families that reject sampling parameters. Claude Sonnet 5 and the
#: Opus 4.7 line onward return a 400 when a request carries a non-default
#: temperature, so agents that pin greedy decoding must omit the setting there.
NO_SAMPLING_MODEL_MARKERS = (
    "sonnet-5",
    "opus-4-7",
    "opus-4-8",
    "opus-5",
    "fable",
    "mythos",
)


def supports_sampling_temperature(model_id: str) -> bool:
    """Whether the model accepts an explicit temperature setting."""
    lowered = model_id.lower()
    return not any(marker in lowered for marker in NO_SAMPLING_MODEL_MARKERS)


class ModelProvider(str, Enum):
    """Supported model providers."""

    ANTHROPIC = "anthropic"
    OPENAI = "openai"
    GROQ = "groq"
    GEMINI = "gemini"
    VERTEXAI = "vertexai"
    BEDROCK = "bedrock"
    OLLAMA = "ollama"
    XAI = "xai"
    MISTRAL = "mistral"
    DEEPSEEK = "deepseek"
    COHERE = "cohere"
    CEREBRAS = "cerebras"


class ModelConfig:
    """Model configuration mappings."""

    # Default model per provider, used when no model is specified. Prefer the
    # provider's rolling alias over a dated snapshot so the default tracks the
    # latest build and does not rot (Bedrock has no alias, so it pins a dated
    # id; Ollama is local and not in the registry). Every non-local default is
    # asserted to exist in pydantic-ai's registry by test_factory, so a default
    # that goes invalid fails CI rather than at runtime.
    DEFAULT_MODELS = {
        ModelProvider.ANTHROPIC: "claude-haiku-4-5",
        ModelProvider.OPENAI: "gpt-4.1-mini",
        ModelProvider.GROQ: "llama-3.3-70b-versatile",
        ModelProvider.GEMINI: "gemini-2.5-flash",
        ModelProvider.VERTEXAI: "gemini-2.5-flash",
        ModelProvider.BEDROCK: "anthropic.claude-haiku-4-5-20251001-v1:0",
        ModelProvider.OLLAMA: "llama3.2",
        ModelProvider.XAI: "grok-4-fast",
        ModelProvider.MISTRAL: "mistral-small-latest",
        ModelProvider.DEEPSEEK: "deepseek-chat",
        ModelProvider.COHERE: "command-r-plus-08-2024",
        ModelProvider.CEREBRAS: "gpt-oss-120b",
    }

    # Environment variable names for API keys
    API_KEY_VARS = {
        ModelProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
        ModelProvider.OPENAI: "OPENAI_API_KEY",
        ModelProvider.GROQ: "GROQ_API_KEY",
        ModelProvider.GEMINI: "GEMINI_API_KEY",
        ModelProvider.VERTEXAI: "GOOGLE_APPLICATION_CREDENTIALS",
        ModelProvider.BEDROCK: "AWS_ACCESS_KEY_ID",
        ModelProvider.OLLAMA: None,  # No API key needed for local Ollama
        ModelProvider.XAI: "XAI_API_KEY",
        ModelProvider.MISTRAL: "MISTRAL_API_KEY",
        ModelProvider.DEEPSEEK: "DEEPSEEK_API_KEY",
        ModelProvider.COHERE: "CO_API_KEY",
        ModelProvider.CEREBRAS: "CEREBRAS_API_KEY",
    }

    # Maps each provider to its prefix in pydantic-ai's model registry
    # (the "provider:model" ids returned by known_model_names). Gemini's
    # developer API is "google" in the registry and Vertex is "google-cloud";
    # Ollama has no registry prefix because it rides the OpenAI-compatible
    # provider, so it has no entry here and falls back to its default model.
    REGISTRY_PREFIXES = {
        ModelProvider.ANTHROPIC: "anthropic",
        ModelProvider.OPENAI: "openai",
        ModelProvider.GROQ: "groq",
        ModelProvider.GEMINI: "google",
        ModelProvider.VERTEXAI: "google-cloud",
        ModelProvider.BEDROCK: "bedrock",
        ModelProvider.XAI: "xai",
        ModelProvider.MISTRAL: "mistral",
        ModelProvider.DEEPSEEK: "deepseek",
        ModelProvider.COHERE: "cohere",
        ModelProvider.CEREBRAS: "cerebras",
    }

    @classmethod
    def get_known_models(cls) -> Dict[str, List[str]]:
        """Model names per provider, taken from pydantic-ai's model registry.

        pydantic-ai maintains the authoritative set of recognized model ids and
        refreshes it on each release, so this tracks model availability with the
        pinned pydantic-ai version rather than a hand-kept file. The registry
        returns "provider:model" ids; we group them under our own provider
        values via ``REGISTRY_PREFIXES`` and strip the prefix.

        Returns:
            Dictionary mapping provider value to a sorted list of model names.
        """
        from pydantic_ai.models import known_model_names

        names = known_model_names()
        by_provider: Dict[str, List[str]] = {}
        for provider, prefix in cls.REGISTRY_PREFIXES.items():
            head = f"{prefix}:"
            models = sorted(name[len(head) :] for name in names if name.startswith(head))
            if models:
                by_provider[provider.value] = models
        return by_provider


class AgentFactory:
    """
    Factory for creating PydanticAI agents with different model providers.

    This factory provides a unified interface for creating agents that can work
    with various LLM providers, automatically handling API key configuration
    and model selection.
    """

    def __init__(self, enable_logfire: bool = True):
        """
        Initialize the Agent Factory.

        Args:
            enable_logfire: Whether to enable logfire telemetry
        """
        global _logfire_configured
        self.enable_logfire = enable_logfire
        logfire_token = get_settings().obs.logfire_token
        if enable_logfire and logfire_token:
            if not _logfire_configured:
                logfire.configure(
                    service_name="quber",
                    token=logfire_token,
                    console=False,
                )
                logfire.instrument_anthropic()
                logfire.instrument_openai()
                logfire.instrument_pydantic_ai()
                _logfire_configured = True
            self.logfire_configured = True
        else:
            self.logfire_configured = False

    def create_agent(
        self,
        output_type: Type[OutputT],
        system_prompt: str,
        provider: Optional[ModelProvider] = None,
        model: Optional[str] = None,
        **agent_kwargs: Any,
    ) -> Agent[Any, OutputT]:
        """
        Create a PydanticAI agent with the specified configuration.

        Args:
            output_type: The Pydantic model for structured output
            system_prompt: System prompt for the agent
            provider: Model provider to use (auto-detected if not specified)
            model: Specific model to use (uses default if not specified)
            **agent_kwargs: Additional kwargs to pass to Agent constructor

        Returns:
            Configured PydanticAI Agent instance

        Raises:
            ValueError: If no valid API key is found or invalid configuration
        """
        # Handle full model string with provider prefix (e.g., "anthropic:claude-sonnet-4-0")
        if model and ":" in model:
            provider_str, model_name = model.split(":", 1)
            # Try to match provider
            for p in ModelProvider:
                if p.value == provider_str.lower():
                    provider = p
                    model = model_name
                    break

        # Auto-detect provider if not specified
        if provider is None:
            provider = self.detect_provider()
            if provider is None:
                raise ValueError(
                    "No API key found. Please set one of: "
                    + ", ".join([v for v in ModelConfig.API_KEY_VARS.values() if v])
                )

        # Validate API key for provider (if required).
        # These provider key-presence probes intentionally read os.getenv by
        # var name: they check arbitrary provider keys (incl. the
        # boto3 AWS_ACCESS_KEY_ID presence for Bedrock), so a typed Settings
        # field per provider would be overkill. Credentials still resolve via
        # their own chains; this is only a yes/no availability check.
        api_key_var = ModelConfig.API_KEY_VARS.get(provider)
        if api_key_var and not os.getenv(api_key_var):
            raise ValueError(f"API key not found: {api_key_var}")

        # Select model
        if model is None:
            model = ModelConfig.DEFAULT_MODELS[provider]

        # Format model string for PydanticAI
        model_str = self.format_model_string(provider, model)

        # Create agent with logfire instrumentation if enabled
        if self.logfire_configured:
            with logfire.span("create_agent", provider=provider.value, model=model):
                agent = Agent(model_str, output_type=output_type, system_prompt=system_prompt, **agent_kwargs)
        else:
            agent = Agent(model_str, output_type=output_type, system_prompt=system_prompt, **agent_kwargs)

        return agent

    def detect_provider(self) -> Optional[ModelProvider]:
        """
        Auto-detect available provider based on environment variables.

        Returns:
            First available provider or None if no API keys found
        """
        # Preferred order when several keys are present; any remaining provider
        # with a key set is considered after these. Ollama has no key var and is
        # skipped (it is local, not auto-detected).
        priority = [
            ModelProvider.ANTHROPIC,
            ModelProvider.OPENAI,
            ModelProvider.GROQ,
            ModelProvider.GEMINI,
        ]
        ordered = priority + [p for p in ModelProvider if p not in priority]

        for provider in ordered:
            key_var = ModelConfig.API_KEY_VARS.get(provider)
            if key_var and os.getenv(key_var):
                return provider

        return None

    def format_model_string(self, provider: ModelProvider, model: str) -> str:
        """
        Format model string for PydanticAI based on provider.

        Args:
            provider: The model provider
            model: The model name

        Returns:
            Formatted model string for PydanticAI
        """
        # For most providers, PydanticAI expects provider:model format
        # Special case: Anthropic models can be used directly or with prefix
        if provider == ModelProvider.ANTHROPIC:
            # If model already has provider prefix, return as is
            if model.startswith("anthropic:"):
                return model
            # Otherwise, just return the model name (PydanticAI accepts both)
            return model

        # For all other providers, use provider:model format
        return f"{provider.value}:{model}"

    def list_available_models(self, provider: Optional[ModelProvider] = None) -> Dict[str, List[str]]:
        """
        List available models for each provider.

        Args:
            provider: Specific provider to list models for (all if None)

        Returns:
            Dictionary of provider -> list of models
        """
        # Known models come from pydantic-ai's model registry
        known_models = ModelConfig.get_known_models()

        if provider:
            provider_name = provider.value
            models = known_models.get(provider_name, [])
            # If no models found in known_models, show a default
            if not models:
                models = [ModelConfig.DEFAULT_MODELS.get(provider, "No models found")]
            return {provider_name: models}

        # Return all providers with their models
        result = {}
        for p in ModelProvider:
            provider_name = p.value
            models = known_models.get(provider_name, [])
            # If no models found, show the default
            if not models:
                models = [ModelConfig.DEFAULT_MODELS.get(p, "No models found")]
            result[provider_name] = models

        return result

    def get_available_providers(self) -> list[str]:
        """
        Get list of providers with configured API keys.

        Returns:
            List of available provider names
        """
        available = []
        for provider in ModelProvider:
            api_key_var = ModelConfig.API_KEY_VARS.get(provider)
            # Skip providers without API key requirements (like Ollama)
            if api_key_var is None:
                # Ollama is always available if installed locally
                if provider == ModelProvider.OLLAMA:
                    available.append(provider.value)
            elif os.getenv(api_key_var):
                available.append(provider.value)
        return available