import asyncio
import hashlib
import json
import os
from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast

import httpx
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI

import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import BaseOpenAILLM
from litellm.secret_managers.get_azure_ad_token_provider import (
    get_azure_ad_token_provider,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import _add_path_to_api_base

azure_ad_cache = DualCache()


class AzureOpenAIError(BaseLLMException):
    def __init__(
        self,
        status_code,
        message,
        request: Optional[httpx.Request] = None,
        response: Optional[httpx.Response] = None,
        headers: Optional[Union[httpx.Headers, dict]] = None,
        body: Optional[dict] = None,
    ):
        super().__init__(
            status_code=status_code,
            message=message,
            request=request,
            response=response,
            headers=headers,
            body=body,
        )


def process_azure_headers(headers: Union[httpx.Headers, dict]) -> dict:
    openai_headers = {}
    if "x-ratelimit-limit-requests" in headers:
        openai_headers["x-ratelimit-limit-requests"] = headers[
            "x-ratelimit-limit-requests"
        ]
    if "x-ratelimit-remaining-requests" in headers:
        openai_headers["x-ratelimit-remaining-requests"] = headers[
            "x-ratelimit-remaining-requests"
        ]
    if "x-ratelimit-limit-tokens" in headers:
        openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"]
    if "x-ratelimit-remaining-tokens" in headers:
        openai_headers["x-ratelimit-remaining-tokens"] = headers[
            "x-ratelimit-remaining-tokens"
        ]
    llm_response_headers = {
        "{}-{}".format("llm_provider", k): v for k, v in headers.items()
    }

    return {**llm_response_headers, **openai_headers}


def get_azure_ad_token_from_entra_id(
    tenant_id: str,
    client_id: str,
    client_secret: str,
    scope: str = "https://cognitiveservices.azure.com/.default",
) -> Callable[[], str]:
    """
    Get Azure AD token provider from `client_id`, `client_secret`, and `tenant_id`

    Args:
        tenant_id: str
        client_id: str
        client_secret: str
        scope: str

    Returns:
        callable that returns a bearer token.
    """
    from azure.identity import ClientSecretCredential, get_bearer_token_provider

    verbose_logger.debug("Getting Azure AD Token from Entra ID")

    if tenant_id.startswith("os.environ/"):
        _tenant_id = get_secret_str(tenant_id)
    else:
        _tenant_id = tenant_id

    if client_id.startswith("os.environ/"):
        _client_id = get_secret_str(client_id)
    else:
        _client_id = client_id

    if client_secret.startswith("os.environ/"):
        _client_secret = get_secret_str(client_secret)
    else:
        _client_secret = client_secret

    verbose_logger.debug(
        "tenant_id=%s, client_id=%s, client_secret=[set=%s]",
        _tenant_id,
        _client_id,
        _client_secret is not None,
    )
    if _tenant_id is None or _client_id is None or _client_secret is None:
        raise ValueError("tenant_id, client_id, and client_secret must be provided")
    credential = ClientSecretCredential(_tenant_id, _client_id, _client_secret)

    token_provider = get_bearer_token_provider(credential, scope)

    verbose_logger.debug("token_provider %s", token_provider)

    return token_provider


def get_azure_ad_token_from_username_password(
    client_id: str,
    azure_username: str,
    azure_password: str,
    scope: str = "https://cognitiveservices.azure.com/.default",
) -> Callable[[], str]:
    """
    Get Azure AD token provider from `client_id`, `azure_username`, and `azure_password`

    Args:
        client_id: str
        azure_username: str
        azure_password: str
        scope: str

    Returns:
        callable that returns a bearer token.
    """
    from azure.identity import UsernamePasswordCredential, get_bearer_token_provider

    verbose_logger.debug(
        "client_id=%s, azure_username=[set=%s], azure_password=[set=%s]",
        client_id,
        azure_username is not None,
        azure_password is not None,
    )
    credential = UsernamePasswordCredential(
        client_id=client_id,
        username=azure_username,
        password=azure_password,
    )

    token_provider = get_bearer_token_provider(credential, scope)

    verbose_logger.debug("token_provider %s", token_provider)

    return token_provider


def get_azure_ad_token_from_oidc(
    azure_ad_token: str,
    azure_client_id: Optional[str] = None,
    azure_tenant_id: Optional[str] = None,
    scope: Optional[str] = None,
) -> str:
    """
    Get Azure AD token from OIDC token

    Args:
        azure_ad_token: str
        azure_client_id: Optional[str]
        azure_tenant_id: Optional[str]
        scope: str

    Returns:
        `azure_ad_token_access_token` - str
    """
    if scope is None:
        scope = "https://cognitiveservices.azure.com/.default"
    azure_authority_host = os.getenv(
        "AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com"
    )
    azure_client_id = azure_client_id or os.getenv("AZURE_CLIENT_ID")
    azure_tenant_id = azure_tenant_id or os.getenv("AZURE_TENANT_ID")
    if azure_client_id is None or azure_tenant_id is None:
        raise AzureOpenAIError(
            status_code=422,
            message="AZURE_CLIENT_ID and AZURE_TENANT_ID must be set",
        )

    oidc_token = get_secret_str(azure_ad_token)

    if oidc_token is None:
        raise AzureOpenAIError(
            status_code=401,
            message="OIDC token could not be retrieved from secret manager.",
        )

    azure_ad_token_cache_key = json.dumps(
        {
            "azure_client_id": azure_client_id,
            "azure_tenant_id": azure_tenant_id,
            "azure_authority_host": azure_authority_host,
            "oidc_token": oidc_token,
        }
    )

    azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key)
    if azure_ad_token_access_token is not None:
        return azure_ad_token_access_token

    client = litellm.module_level_client

    req_token = client.post(
        f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
        data={
            "client_id": azure_client_id,
            "grant_type": "client_credentials",
            "scope": scope,
            "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
            "client_assertion": oidc_token,
        },
    )

    if req_token.status_code != 200:
        raise AzureOpenAIError(
            status_code=req_token.status_code,
            message=req_token.text,
        )

    azure_ad_token_json = req_token.json()
    azure_ad_token_access_token = azure_ad_token_json.get("access_token", None)
    azure_ad_token_expires_in = azure_ad_token_json.get("expires_in", None)

    if azure_ad_token_access_token is None:
        raise AzureOpenAIError(
            status_code=422, message="Azure AD Token access_token not returned"
        )

    if azure_ad_token_expires_in is None:
        raise AzureOpenAIError(
            status_code=422, message="Azure AD Token expires_in not returned"
        )

    azure_ad_cache.set_cache(
        key=azure_ad_token_cache_key,
        value=azure_ad_token_access_token,
        ttl=azure_ad_token_expires_in,
    )

    return azure_ad_token_access_token


def select_azure_base_url_or_endpoint(azure_client_params: dict):
    azure_endpoint = azure_client_params.get("azure_endpoint", None)
    if azure_endpoint is not None:
        # see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
        if "/openai/deployments" in azure_endpoint:
            # this is base_url, not an azure_endpoint
            azure_client_params["base_url"] = azure_endpoint
            azure_client_params.pop("azure_endpoint")

    return azure_client_params


def get_azure_ad_token(
    litellm_params: GenericLiteLLMParams,
) -> Optional[str]:
    """
    Get Azure AD token from various authentication methods.

    This function tries different methods to obtain an Azure AD token:
    1. From an existing token provider
    2. From Entra ID using tenant_id, client_id, and client_secret
    3. From username and password
    4. From OIDC token
    5. From a service principal with secret workflow
    6. From DefaultAzureCredential

    Args:
        litellm_params: Dictionary containing authentication parameters
            - azure_ad_token_provider: Optional callable that returns a token
            - azure_ad_token: Optional existing token
            - tenant_id: Optional Azure tenant ID
            - client_id: Optional Azure client ID
            - client_secret: Optional Azure client secret
            - azure_username: Optional Azure username
            - azure_password: Optional Azure password

    Returns:
        Azure AD token as string if successful, None otherwise
    """
    # Extract parameters
    # Use `or` instead of default parameter to handle cases where key exists but value is None
    azure_ad_token_provider = litellm_params.get("azure_ad_token_provider")
    azure_ad_token = litellm_params.get("azure_ad_token") or get_secret_str(
        "AZURE_AD_TOKEN"
    )
    tenant_id = litellm_params.get("tenant_id") or os.getenv("AZURE_TENANT_ID")
    client_id = litellm_params.get("client_id") or os.getenv("AZURE_CLIENT_ID")
    client_secret = litellm_params.get("client_secret") or os.getenv(
        "AZURE_CLIENT_SECRET"
    )
    azure_username = litellm_params.get("azure_username") or os.getenv("AZURE_USERNAME")
    azure_password = litellm_params.get("azure_password") or os.getenv("AZURE_PASSWORD")
    scope = litellm_params.get("azure_scope") or os.getenv(
        "AZURE_SCOPE", "https://cognitiveservices.azure.com/.default"
    )
    if scope is None:
        scope = "https://cognitiveservices.azure.com/.default"

    # Try to get token provider from Entra ID
    if azure_ad_token_provider is None and tenant_id and client_id and client_secret:
        verbose_logger.debug(
            "Using Azure AD Token Provider from Entra ID for Azure Auth"
        )
        azure_ad_token_provider = get_azure_ad_token_from_entra_id(
            tenant_id=tenant_id,
            client_id=client_id,
            client_secret=client_secret,
            scope=scope,
        )

    # Try to get token provider from username and password
    if (
        azure_ad_token_provider is None
        and azure_username
        and azure_password
        and client_id
    ):
        verbose_logger.debug("Using Azure Username and Password for Azure Auth")
        azure_ad_token_provider = get_azure_ad_token_from_username_password(
            azure_username=azure_username,
            azure_password=azure_password,
            client_id=client_id,
            scope=scope,
        )

    # Try to get token from OIDC
    if (
        client_id
        and tenant_id
        and azure_ad_token
        and azure_ad_token.startswith("oidc/")
    ):
        verbose_logger.debug("Using Azure OIDC Token for Azure Auth")
        azure_ad_token = get_azure_ad_token_from_oidc(
            azure_ad_token=azure_ad_token,
            azure_client_id=client_id,
            azure_tenant_id=tenant_id,
            scope=scope,
        )
    # Try to get token provider from service principal or DefaultAzureCredential
    elif (
        azure_ad_token_provider is None
        and litellm.enable_azure_ad_token_refresh is True
    ):
        verbose_logger.debug(
            "Using Azure AD token provider based on Service Principal with Secret workflow or DefaultAzureCredential for Azure Auth"
        )
        try:
            azure_ad_token_provider = get_azure_ad_token_provider(azure_scope=scope)
        except ValueError:
            verbose_logger.debug("Azure AD Token Provider could not be used.")
        except Exception as e:
            verbose_logger.error(
                f"Error calling Azure AD token provider: {str(e)}. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential"
            )
            raise e

        #########################################################
        # If litellm.enable_azure_ad_token_refresh is True and no other token provider is available,
        # try to get DefaultAzureCredential provider
        #########################################################
        if azure_ad_token_provider is None and azure_ad_token is None:
            azure_ad_token_provider = (
                BaseAzureLLM._try_get_default_azure_credential_provider(
                    scope=scope,
                )
            )

    # Execute the token provider to get the token if available
    if azure_ad_token_provider and callable(azure_ad_token_provider):
        try:
            token = azure_ad_token_provider()
            if not isinstance(token, str):
                verbose_logger.error(
                    f"Azure AD token provider returned non-string value: {type(token)}"
                )
                raise TypeError(f"Azure AD token must be a string, got {type(token)}")
            else:
                azure_ad_token = token
        except TypeError:
            # Re-raise TypeError directly
            raise
        except Exception as e:
            verbose_logger.error(f"Error calling Azure AD token provider: {str(e)}")
            raise RuntimeError(f"Failed to get Azure AD token: {str(e)}") from e

    return azure_ad_token


class BaseAzureLLM(BaseOpenAILLM):
    @staticmethod
    def _try_get_default_azure_credential_provider(
        scope: str,
    ) -> Optional[Callable[[], str]]:
        """
        Try to get DefaultAzureCredential provider

        Args:
            scope: Azure scope for the token

        Returns:
            Token provider callable if DefaultAzureCredential is enabled and available, None otherwise
        """
        from litellm.types.secret_managers.get_azure_ad_token_provider import (
            AzureCredentialType,
        )

        verbose_logger.debug("Attempting to use DefaultAzureCredential for Azure Auth")

        try:
            azure_ad_token_provider = get_azure_ad_token_provider(
                azure_scope=scope,
                azure_credential=AzureCredentialType.DefaultAzureCredential,
            )
            verbose_logger.debug(
                "Successfully obtained Azure AD token provider using DefaultAzureCredential"
            )
            return azure_ad_token_provider
        except Exception as e:
            verbose_logger.debug(f"DefaultAzureCredential failed: {str(e)}")
            return None

    def get_azure_openai_client(
        self,
        api_key: Optional[str],
        api_base: Optional[str],
        api_version: Optional[str] = None,
        client: Optional[
            Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]
        ] = None,
        litellm_params: Optional[dict] = None,
        _is_async: bool = False,
        model: Optional[str] = None,
    ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]]:
        openai_client: Optional[
            Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]
        ] = None
        client_initialization_params: dict = locals()
        client_initialization_params["is_async"] = _is_async
        _lp = litellm_params or {}
        _ad_provider = _lp.get("azure_ad_token_provider")
        _ad_token = _lp.get("azure_ad_token")
        _client_secret = _lp.get("client_secret")
        _azure_password = _lp.get("azure_password")
        client_initialization_params["azure_ad_token"] = (
            hashlib.sha256(_ad_token.encode()).hexdigest()
            if isinstance(_ad_token, str)
            else None
        )
        client_initialization_params["azure_ad_token_provider"] = (
            f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}"
            f"|tenant_id={_lp.get('tenant_id')}"
            f"|client_id={_lp.get('client_id')}"
            f"|client_secret={hashlib.sha256(_client_secret.encode()).hexdigest() if isinstance(_client_secret, str) else None}"
            f"|azure_username={_lp.get('azure_username')}"
            f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}"
            f"|azure_scope={_lp.get('azure_scope')}"
        )
        if client is None:
            cached_client = self.get_cached_openai_client(
                client_initialization_params=client_initialization_params,
                client_type="azure",
            )
            if cached_client:
                if isinstance(
                    cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)
                ):
                    return cached_client

            azure_client_params = self.initialize_azure_sdk_client(
                litellm_params=litellm_params or {},
                api_key=api_key,
                api_base=api_base,
                model_name=model,
                api_version=api_version,
                is_async=_is_async,
            )

            # For Azure v1 API, use standard OpenAI client instead of AzureOpenAI
            # See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs
            if self._is_azure_v1_api_version(api_version):
                # Extract only params that OpenAI client accepts
                # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview"
                # The OpenAI client accepts a callable for `api_key` and re-invokes it
                # on every request (via `_refresh_api_key`), so passing
                # `azure_ad_token_provider` directly preserves Azure AD token refresh
                # behavior that the regular AzureOpenAI client provides.
                v1_api_key: Optional[Union[str, Callable[[], Any]]] = (
                    azure_client_params.get("api_key")
                    or azure_client_params.get("azure_ad_token_provider")
                    or azure_client_params.get("azure_ad_token")
                )
                if _is_async is True and callable(v1_api_key):
                    # AsyncOpenAI expects an async provider; wrap the sync provider
                    # returned by azure-identity. Offload to a thread so a token
                    # refresh (blocking HTTP call to AAD on cache miss) does not
                    # stall the event loop.
                    _sync_provider = v1_api_key

                    async def _async_v1_api_key() -> str:
                        return await asyncio.to_thread(_sync_provider)

                    v1_api_key = _async_v1_api_key

                v1_params: Dict[str, Any] = {
                    "api_key": v1_api_key,
                    "base_url": f"{api_base}/openai/v1/",
                }
                if "timeout" in azure_client_params:
                    v1_params["timeout"] = azure_client_params["timeout"]
                if "max_retries" in azure_client_params:
                    v1_params["max_retries"] = azure_client_params["max_retries"]
                if "http_client" in azure_client_params:
                    v1_params["http_client"] = azure_client_params["http_client"]

                verbose_logger.debug(
                    f"Using Azure v1 API with base_url: {v1_params['base_url']}"
                )

                if _is_async is True:
                    openai_client = AsyncOpenAI(**v1_params)  # type: ignore
                else:
                    openai_client = OpenAI(**v1_params)  # type: ignore
            else:
                # Traditional Azure API uses AzureOpenAI client
                if _is_async is True:
                    openai_client = AsyncAzureOpenAI(**azure_client_params)
                else:
                    openai_client = AzureOpenAI(**azure_client_params)  # type: ignore
        else:
            openai_client = client
            if (
                api_version is not None
                and isinstance(openai_client, (AzureOpenAI, AsyncAzureOpenAI))
                and isinstance(openai_client._custom_query, dict)
            ):
                # set api_version to version passed by user
                openai_client._custom_query.setdefault("api-version", api_version)

        # save client in-memory cache
        self.set_cached_openai_client(
            openai_client=openai_client,
            client_initialization_params=client_initialization_params,
            client_type="azure",
        )
        return openai_client

    def initialize_azure_sdk_client(
        self,
        litellm_params: dict,
        api_key: Optional[str],
        api_base: Optional[str],
        model_name: Optional[str],
        api_version: Optional[str],
        is_async: bool,
    ) -> dict:
        azure_ad_token_provider = litellm_params.get("azure_ad_token_provider")
        # If we have api_key, then we have higher priority
        azure_ad_token = litellm_params.get("azure_ad_token")

        # litellm_params sometimes contains the key, but the value is None
        # We should respect environment variables in this case
        tenant_id = self._resolve_env_var(
            litellm_params, "tenant_id", "AZURE_TENANT_ID"
        )
        client_id = self._resolve_env_var(
            litellm_params, "client_id", "AZURE_CLIENT_ID"
        )
        client_secret = self._resolve_env_var(
            litellm_params, "client_secret", "AZURE_CLIENT_SECRET"
        )
        azure_username = self._resolve_env_var(
            litellm_params, "azure_username", "AZURE_USERNAME"
        )
        azure_password = self._resolve_env_var(
            litellm_params, "azure_password", "AZURE_PASSWORD"
        )
        scope = self._resolve_env_var(litellm_params, "azure_scope", "AZURE_SCOPE")
        if scope is None:
            scope = "https://cognitiveservices.azure.com/.default"

        max_retries = litellm_params.get("max_retries")
        timeout = litellm_params.get("timeout")
        if (
            not api_key
            and azure_ad_token_provider is None
            and tenant_id
            and client_id
            and client_secret
        ):
            verbose_logger.debug(
                "Using Azure AD Token Provider from Entra ID for Azure Auth"
            )
            azure_ad_token_provider = get_azure_ad_token_from_entra_id(
                tenant_id=tenant_id,
                client_id=client_id,
                client_secret=client_secret,
                scope=scope,
            )
        if (
            azure_ad_token_provider is None
            and azure_username
            and azure_password
            and client_id
        ):
            verbose_logger.debug("Using Azure Username and Password for Azure Auth")
            azure_ad_token_provider = get_azure_ad_token_from_username_password(
                azure_username=azure_username,
                azure_password=azure_password,
                client_id=client_id,
                scope=scope,
            )

        if azure_ad_token is not None and azure_ad_token.startswith("oidc/"):
            verbose_logger.debug("Using Azure OIDC Token for Azure Auth")
            azure_ad_token = get_azure_ad_token_from_oidc(
                azure_ad_token=azure_ad_token,
                azure_client_id=client_id,
                azure_tenant_id=tenant_id,
                scope=scope,
            )
        elif (
            not api_key
            and azure_ad_token_provider is None
            and litellm.enable_azure_ad_token_refresh is True
        ):
            verbose_logger.debug(
                "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth"
            )
            try:
                azure_ad_token_provider = get_azure_ad_token_provider(
                    azure_scope=scope,
                )
            except ValueError:
                verbose_logger.debug("Azure AD Token Provider could not be used.")
        if api_version is None:
            api_version = os.getenv(
                "AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION
            )

        _api_key = api_key
        if _api_key is not None and isinstance(_api_key, str):
            # only show first 5 chars of api_key
            _api_key = _api_key[:8] + "*" * 15
        verbose_logger.debug(
            f"Initializing Azure OpenAI Client for {model_name}, Api Base: {str(api_base)}, Api Key:{_api_key}"
        )
        azure_client_params = {
            "api_key": api_key,
            "azure_endpoint": api_base,
            "api_version": api_version,
            "azure_ad_token": azure_ad_token,
            "azure_ad_token_provider": azure_ad_token_provider,
        }
        # init http client + SSL Verification settings
        if is_async is True:
            azure_client_params["http_client"] = self._get_async_http_client()
        else:
            azure_client_params["http_client"] = self._get_sync_http_client()

        if max_retries is not None:
            azure_client_params["max_retries"] = max_retries
        if timeout is not None:
            azure_client_params["timeout"] = timeout

        if azure_ad_token_provider is not None:
            azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider
        # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client
        # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client

        azure_client_params = select_azure_base_url_or_endpoint(
            azure_client_params=azure_client_params
        )

        return azure_client_params

    def _init_azure_client_for_cloudflare_ai_gateway(
        self,
        api_base: str,
        model: str,
        api_version: str,
        max_retries: int,
        timeout: Union[float, httpx.Timeout],
        litellm_params: dict,
        api_key: Optional[str],
        azure_ad_token: Optional[str],
        azure_ad_token_provider: Optional[Callable[[], str]],
        acompletion: bool,
        client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None,
    ) -> Union[AzureOpenAI, AsyncAzureOpenAI]:
        ## build base url - assume api base includes resource name
        tenant_id = litellm_params.get("tenant_id", os.getenv("AZURE_TENANT_ID"))
        client_id = litellm_params.get("client_id", os.getenv("AZURE_CLIENT_ID"))
        scope = litellm_params.get(
            "azure_scope",
            os.getenv("AZURE_SCOPE", "https://cognitiveservices.azure.com/.default"),
        )
        if client is None:
            if not api_base.endswith("/"):
                api_base += "/"
            api_base += f"{model}"

            azure_client_params: Dict[str, Any] = {
                "api_version": api_version,
                "base_url": f"{api_base}",
                "http_client": litellm.client_session,
                "max_retries": max_retries,
                "timeout": timeout,
            }
            if api_key is not None:
                azure_client_params["api_key"] = api_key
            elif azure_ad_token is not None:
                if azure_ad_token.startswith("oidc/"):
                    azure_ad_token = get_azure_ad_token_from_oidc(
                        azure_ad_token=azure_ad_token,
                        azure_client_id=client_id,
                        azure_tenant_id=tenant_id,
                        scope=scope,
                    )

                azure_client_params["azure_ad_token"] = azure_ad_token
            if azure_ad_token_provider is not None:
                azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider

            if acompletion is True:
                client = AsyncAzureOpenAI(**azure_client_params)  # type: ignore
            else:
                client = AzureOpenAI(**azure_client_params)  # type: ignore
        return client

    @staticmethod
    def _base_validate_azure_environment(
        headers: dict, litellm_params: Optional[GenericLiteLLMParams]
    ) -> dict:
        litellm_params = litellm_params or GenericLiteLLMParams()

        # Check if api-key is already in headers; if so, use it
        if "api-key" in headers:
            return headers

        api_key = (
            litellm_params.api_key
            or litellm.api_key
            or litellm.azure_key
            or get_secret_str("AZURE_OPENAI_API_KEY")
            or get_secret_str("AZURE_API_KEY")
        )

        if api_key:
            headers["api-key"] = api_key
            return headers

        ### Fallback to Azure AD token-based authentication if no API key is available
        ### Retrieves Azure AD token and adds it to the Authorization header
        azure_ad_token = get_azure_ad_token(litellm_params)
        if azure_ad_token:
            headers["Authorization"] = f"Bearer {azure_ad_token}"

        return headers

    @staticmethod
    def _get_base_azure_url(
        api_base: Optional[str],
        litellm_params: Optional[Union[GenericLiteLLMParams, Dict[str, Any]]],
        route: Union[Literal["/openai/responses", "/openai/vector_stores"], str],
        default_api_version: Optional[Union[str, Literal["latest", "preview"]]] = None,
    ) -> str:
        """
        Get the base Azure URL for the given route and API version.

        Args:
            api_base: The base URL of the Azure API.
            litellm_params: The litellm parameters.
            route: The route to the API.
            default_api_version: The default API version to use if no api_version is provided. If 'latest', it will use `openai/v1/...` route.
        """

        api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
        if api_base is None:
            raise ValueError(
                f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`"
            )
        original_url = httpx.URL(api_base)

        # Extract api_version or use default
        litellm_params = litellm_params or {}
        api_version = (
            cast(Optional[str], litellm_params.get("api_version"))
            or default_api_version
        )

        # Create a new dictionary with existing params
        query_params = dict(original_url.params)

        # Add api_version if needed
        if "api-version" not in query_params and api_version:
            query_params["api-version"] = api_version

        # Add the path to the base URL
        if route not in api_base:
            new_url = _add_path_to_api_base(api_base=api_base, ending_path=route)
        else:
            new_url = api_base

        if BaseAzureLLM._is_azure_v1_api_version(api_version):
            # ensure the request go to /openai/v1 and not just /openai
            if "/openai/v1" not in new_url:
                parsed_url = httpx.URL(new_url)
                new_url = str(
                    parsed_url.copy_with(
                        path=parsed_url.path.replace("/openai", "/openai/v1")
                    )
                )

        # Use the new query_params dictionary
        final_url = httpx.URL(new_url).copy_with(params=query_params)

        return str(final_url)

    @staticmethod
    def _is_azure_v1_api_version(api_version: Optional[str]) -> bool:
        if api_version is None:
            return False
        return api_version in {"preview", "latest", "v1"}

    def _resolve_env_var(
        self, litellm_params: Dict[str, Any], param_key: str, env_var_key: str
    ) -> Optional[str]:
        """Resolve the environment variable for a given parameter key.

        The logic here is different from `params.get(key, os.getenv(env_var))` because
        litellm_params may contain the key with a None value, in which case we want
        to fallback to the environment variable.
        """
        param_value = litellm_params.get(param_key)
        if param_value is not None:
            return param_value
        return os.getenv(env_var_key)


class AzureCredentials(NamedTuple):
    api_base: Optional[str]
    api_key: Optional[str]
    api_version: Optional[str]


def get_azure_credentials(
    api_base: Optional[str] = None,
    api_key: Optional[str] = None,
    api_version: Optional[str] = None,
) -> AzureCredentials:
    """Resolve Azure credentials from params, litellm globals, and env vars."""
    resolved_api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
    resolved_api_version = (
        api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
    )
    resolved_api_key = (
        api_key
        or litellm.api_key
        or litellm.azure_key
        or get_secret_str("AZURE_OPENAI_API_KEY")
        or get_secret_str("AZURE_API_KEY")
    )
    return AzureCredentials(
        api_base=resolved_api_base,
        api_key=resolved_api_key,
        api_version=resolved_api_version,
    )
