import os

from azure.identity import AzureCliCredential, get_bearer_token_provider
from dotenv import load_dotenv
from openai import AsyncAzureOpenAI
from pydantic_ai.models import Model
from pydantic_ai.models.bedrock import BedrockConverseModel
from pydantic_ai.models.cerebras import CerebrasModel
from pydantic_ai.models.gemini import GeminiModel
from pydantic_ai.models.mistral import MistralModel
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.models.openrouter import OpenRouterModel
from pydantic_ai.providers.bedrock import BedrockProvider
from pydantic_ai.providers.cerebras import CerebrasProvider
from pydantic_ai.providers.google import GoogleProvider
from pydantic_ai.providers.mistral import MistralProvider
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.providers.openrouter import OpenRouterProvider

from probegpt.core.config import ProviderConfig

load_dotenv()


def create_model(config: ProviderConfig) -> Model:
    if config.provider_type == "azure_openai":
        if not config.azure_endpoint or not config.azure_deployment:
            raise ValueError("Azure OpenAI requires endpoint and deployment.")
        credential = AzureCliCredential()
        token_provider = get_bearer_token_provider(
            credential, "https://cognitiveservices.azure.com/.default"
        )
        client = AsyncAzureOpenAI(
            azure_endpoint=config.azure_endpoint,
            api_version=config.azure_api_version,
            azure_ad_token_provider=token_provider,
        )
        return OpenAIChatModel(
            config.azure_deployment,
            provider=OpenAIProvider(openai_client=client),
        )

    elif config.provider_type == "openrouter":
        if not config.openrouter_model:
            raise ValueError("OpenRouter requires a model to be selected.")
        api_key = os.getenv("OPENROUTER_API_KEY", "")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY is not set.")
        return OpenRouterModel(
            config.openrouter_model,
            provider=OpenRouterProvider(api_key=api_key),
        )

    elif config.provider_type == "vertex_ai":
        if not config.vertex_project:
            raise ValueError("Vertex AI requires a GCP project ID.")
        if not config.vertex_model:
            raise ValueError("Vertex AI requires a model name.")
        return GeminiModel(
            config.vertex_model,
            provider=GoogleProvider(  # type: ignore[arg-type]
                project=config.vertex_project,
                location=config.vertex_region,
                vertexai=True,
            ),
        )

    elif config.provider_type == "bedrock":
        if not config.bedrock_model:
            raise ValueError("Bedrock requires a model ID.")
        return BedrockConverseModel(
            config.bedrock_model,
            provider=BedrockProvider(
                region_name=config.bedrock_region or "us-east-1",
                profile_name=config.bedrock_profile or None,
            ),
        )

    elif config.provider_type == "cerebras":
        if not config.cerebras_model:
            raise ValueError("Cerebras requires a model name.")
        api_key = os.getenv("CEREBRAS_API_KEY", "")
        if not api_key:
            raise ValueError("CEREBRAS_API_KEY is not set.")
        return CerebrasModel(
            config.cerebras_model,
            provider=CerebrasProvider(api_key=api_key),
        )

    elif config.provider_type == "mistral":
        if not config.mistral_model:
            raise ValueError("Mistral requires a model name.")
        api_key = os.getenv("MISTRAL_API_KEY", "")
        if not api_key:
            raise ValueError("MISTRAL_API_KEY is not set.")
        return MistralModel(
            config.mistral_model,
            provider=MistralProvider(api_key=api_key),
        )

    else:
        raise ValueError(f"Unknown provider type: {config.provider_type!r}")
