import asyncio
import hashlib
import json
import os
from collections.abc import Callable, Mapping
from functools import lru_cache
from types import MappingProxyType
from typing import Any, Final, Literal, NamedTuple, cast

import httpx
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
from typing_extensions import ReadOnly, TypedDict

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: Final = DualCache()


class _AzureAdTokenJson(TypedDict, total=False):
    access_token: ReadOnly[str]
    expires_in: ReadOnly[int]


class _AzureV1ClientParams(TypedDict, total=False, extra_items=object):
    base_url: ReadOnly[str]


class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object):
    api_version: ReadOnly[str]
    base_url: ReadOnly[str]
    max_retries: ReadOnly[int]
    timeout: ReadOnly[float | httpx.Timeout]


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


def process_azure_headers(headers: httpx.Headers | dict) -> dict:
    openai_headers: Final = {}
    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: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()}

    return {**llm_response_headers, **openai_headers}


@lru_cache(maxsize=128)
def _cached_entra_id_token_provider(
    tenant_id: str,
    client_id: str,
    client_secret: str,
    scope: str,
) -> Callable[[], str]:
    """Build (once per credential set) a bearer token provider backed by a `ClientSecretCredential`.

    The credential caches the access token internally and only talks to Entra ID when it is close
    to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one
    per request.
    """
    from azure.identity import ClientSecretCredential, get_bearer_token_provider

    return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope)


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.
    """
    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")

    token_provider: Final = _cached_entra_id_token_provider(
        tenant_id=_tenant_id,
        client_id=_client_id,
        client_secret=_client_secret,
        scope=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: Final = UsernamePasswordCredential(
        client_id=client_id,
        username=azure_username,
        password=azure_password,
    )

    token_provider: Final = 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: str | None = None,
    azure_tenant_id: str | None = None,
    scope: str | None = 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: Final = 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: Final = 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: Final = 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: Final = litellm.module_level_client

    req_token: Final = 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: Final[_AzureAdTokenJson] = req_token.json()
    azure_ad_token_access_token = azure_ad_token_json.get("access_token", None)
    azure_ad_token_expires_in: Final = 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: Final = 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,
) -> str | None:
    """
    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: Final = litellm_params.get("tenant_id") or os.getenv("AZURE_TENANT_ID")
    client_id: Final = litellm_params.get("client_id") or os.getenv("AZURE_CLIENT_ID")
    client_secret: Final = litellm_params.get("client_secret") or os.getenv("AZURE_CLIENT_SECRET")
    azure_username: Final = litellm_params.get("azure_username") or os.getenv("AZURE_USERNAME")
    azure_password: Final = 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(
                "Error calling Azure AD token provider: %s. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential",
                e,
            )
            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: Final = azure_ad_token_provider()
            if not isinstance(token, str):
                verbose_logger.error("Azure AD token provider returned non-string value: %s", 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("Error calling Azure AD token provider: %s", e)
            raise RuntimeError(f"Failed to get Azure AD token: {e}") from e

    return azure_ad_token


class BaseAzureLLM(BaseOpenAILLM):
    @staticmethod
    def _try_get_default_azure_credential_provider(
        scope: str,
    ) -> Callable[[], str] | None:
        """
        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: Final = 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("DefaultAzureCredential failed: %s", e)
            return None

    def get_azure_openai_client(
        self,
        api_key: str | None,
        api_base: str | None,
        api_version: str | None = None,
        client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None,
        litellm_params: dict | None = None,
        _is_async: bool = False,
        model: str | None = None,
    ) -> AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None:
        openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None
        client_initialization_params: Final[dict] = locals()
        client_initialization_params["is_async"] = _is_async
        _lp: Final = litellm_params or {}
        _ad_provider: Final = _lp.get("azure_ad_token_provider")
        _ad_token: Final = _lp.get("azure_ad_token")
        _client_secret: Final = _lp.get("client_secret")
        _azure_password: Final = _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 not None:
            if (
                api_version is not None
                and isinstance(client, (AzureOpenAI, AsyncAzureOpenAI))
                and isinstance(client._custom_query, dict)
            ):
                # set api_version to version passed by user
                client._custom_query.setdefault("api-version", api_version)
            self.set_cached_openai_client(
                openai_client=client,
                client_initialization_params=client_initialization_params,
                client_type="azure",
                litellm_owned_client=False,
            )
            return client

        cached_client: Final = 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: Final = 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: str | Callable[[], Any] | None = (
                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: Final = 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: Final[_AzureV1ClientParams] = {
                "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("Using Azure v1 API with base_url: %s", v1_params["base_url"])

            if _is_async is True:
                openai_client = AsyncOpenAI(**v1_params)
            else:
                openai_client = OpenAI(**v1_params)
        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)

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

    def initialize_azure_sdk_client(
        self,
        litellm_params: dict,
        api_key: str | None,
        api_base: str | None,
        model_name: str | None,
        api_version: str | None,
        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: Final = self._resolve_env_var(litellm_params, "tenant_id", "AZURE_TENANT_ID")
        client_id: Final = self._resolve_env_var(litellm_params, "client_id", "AZURE_CLIENT_ID")
        client_secret: Final = self._resolve_env_var(litellm_params, "client_secret", "AZURE_CLIENT_SECRET")
        azure_username: Final = self._resolve_env_var(litellm_params, "azure_username", "AZURE_USERNAME")
        azure_password: Final = 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: Final = litellm_params.get("max_retries")
        timeout: Final = 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(
            "Initializing Azure OpenAI Client for %s, Api Base: %s, Api Key:%s", model_name, api_base, _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: float | httpx.Timeout,
        litellm_params: dict,
        api_key: str | None,
        azure_ad_token: str | None,
        azure_ad_token_provider: Callable[[], str] | None,
        acompletion: bool,
        client: AzureOpenAI | AsyncAzureOpenAI | None = None,
    ) -> AzureOpenAI | AsyncAzureOpenAI:
        ## build base url - assume api base includes resource name
        tenant_id: Final = litellm_params.get("tenant_id", os.getenv("AZURE_TENANT_ID"))
        client_id: Final = litellm_params.get("client_id", os.getenv("AZURE_CLIENT_ID"))
        scope: Final = 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: Final[_AzureGatewayClientParams] = {
                "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)
            else:
                client = AzureOpenAI(**azure_client_params)
        return client

    @staticmethod
    def _base_validate_azure_environment(headers: dict, litellm_params: GenericLiteLLMParams | None) -> 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: Final = (
            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: Final = 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: str | None,
        litellm_params: GenericLiteLLMParams | Mapping[str, object] | None,
        route: Literal["/openai/responses", "/openai/vector_stores"] | str,
        default_api_version: str | Literal["latest", "preview"] | None = 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: Final = httpx.URL(api_base)

        # Extract api_version or use default
        litellm_params = litellm_params or {}
        api_version: Final = cast(str | None, litellm_params.get("api_version")) or default_api_version

        # Create a new dictionary with existing params
        query_params: Final = 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: Final = 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: Final = httpx.URL(new_url).copy_with(params=query_params)

        return str(final_url)

    @staticmethod
    def get_azure_v1_image_url(api_base: str, api_version: str | None, route: str) -> str | None:
        """
        Azure's v1 surface serves images at ``/openai/v1/images/{generations,edits}`` and routes by
        ``model`` in the request body, so any deployment path and stale ``api-version`` in
        ``api_base`` have to be dropped.

        Returns None when ``api_version`` is a dated one, which still uses the deployment route.
        """
        if not BaseAzureLLM._is_azure_v1_api_version(api_version):
            return None

        base_url: Final = httpx.URL(api_base)
        openai_path_start: Final = base_url.path.find("/openai")
        resource_base: Final = str(
            base_url.copy_with(
                path=base_url.path if openai_path_start == -1 else base_url.path[:openai_path_start],
                params=httpx.QueryParams(tuple((k, v) for k, v in base_url.params.multi_items() if k != "api-version")),
            )
        )
        return BaseAzureLLM._get_base_azure_url(
            api_base=resource_base,
            litellm_params=MappingProxyType({"api_version": api_version}),
            route=route,
        )

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

    def _resolve_env_var(
        self, litellm_params: Mapping[str, str | None], param_key: str, env_var_key: str
    ) -> str | None:
        """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: Final = 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: str | None
    api_key: str | None
    api_version: str | None


def get_azure_credentials(
    api_base: str | None = None,
    api_key: str | None = None,
    api_version: str | None = None,
) -> AzureCredentials:
    """Resolve Azure credentials from params, litellm globals, and env vars."""
    resolved_api_base: Final = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
    resolved_api_version: Final = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
    resolved_api_key: Final = (
        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,
    )
