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.openai import OpenAIChatModel
from pydantic_ai.models.openrouter import OpenRouterModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.providers.openrouter import OpenRouterProvider

from 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 'azure_endpoint' and 'azure_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 'openrouter_model' in config.")
        api_key = os.getenv("OPENROUTER_API_KEY", "")
        if not api_key:
            raise ValueError("OpenRouter requires OPENROUTER_API_KEY to be set in .env.")
        return OpenRouterModel(
            config.openrouter_model,
            provider=OpenRouterProvider(api_key=api_key),
        )

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