# Copyright (c) Microsoft. All rights reserved.

import logging
from typing import TYPE_CHECKING, Any

from azure.core.exceptions import ClientAuthenticationError

from ..exceptions import ServiceInvalidAuthError

if TYPE_CHECKING:
    from azure.core.credentials import TokenCredential
    from azure.core.credentials_async import AsyncTokenCredential

logger: logging.Logger = logging.getLogger(__name__)


def get_entra_auth_token(
    credential: "TokenCredential",
    token_endpoint: str,
    **kwargs: Any,
) -> str | None:
    """Retrieve a Microsoft Entra Auth Token for a given token endpoint.

    The token endpoint may be specified as an environment variable, via the .env
    file or as an argument. If the token endpoint is not provided, the default is None.

    Args:
        credential: The Azure credential to use for authentication.
        token_endpoint: The token endpoint to use to retrieve the authentication token.

    Keyword Args:
        **kwargs: Additional keyword arguments to pass to the token retrieval method.

    Returns:
        The Azure token or None if the token could not be retrieved.
    """
    if not token_endpoint:
        raise ServiceInvalidAuthError(
            "A token endpoint must be provided either in settings, as an environment variable, or as an argument."
        )

    try:
        auth_token = credential.get_token(token_endpoint, **kwargs)
    except ClientAuthenticationError as ex:
        logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
        return None

    return auth_token.token if auth_token else None


async def get_entra_auth_token_async(
    credential: "AsyncTokenCredential", token_endpoint: str, **kwargs: Any
) -> str | None:
    """Retrieve a async Microsoft Entra Auth Token for a given token endpoint.

    The token endpoint may be specified as an environment variable, via the .env
    file or as an argument. If the token endpoint is not provided, the default is None.

    Args:
        credential: The async Azure credential to use for authentication.
        token_endpoint: The token endpoint to use to retrieve the authentication token.

    Keyword Args:
        **kwargs: Additional keyword arguments to pass to the token retrieval method.

    Returns:
        The Azure token or None if the token could not be retrieved.
    """
    if not token_endpoint:
        raise ServiceInvalidAuthError(
            "A token endpoint must be provided either in settings, as an environment variable, or as an argument."
        )

    try:
        auth_token = await credential.get_token(token_endpoint, **kwargs)
    except ClientAuthenticationError as ex:
        logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
        return None

    return auth_token.token if auth_token else None
