import logging

import certifi
from azure.ai.projects.aio import AIProjectClient
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.pipeline.transport import AsyncioRequestsTransport
from azure.identity.aio import DefaultAzureCredential, WorkloadIdentityCredential

from unique_web_search.settings import env_settings

_LOGGER = logging.getLogger(__name__)


def _get_workload_identity_credentials(
    with_request_transport: bool = False,
) -> WorkloadIdentityCredential:
    if with_request_transport:
        transport = AsyncioRequestsTransport(connection_verify=certifi.where())
        credentials = WorkloadIdentityCredential(transport=transport)
    else:
        credentials = WorkloadIdentityCredential()

    return credentials


def get_credentials():
    match env_settings.azure_identity_credential_type:
        case "workload":
            return _get_workload_identity_credentials(
                env_settings.use_unique_private_endpoint_transport
            )
        case "default":
            return DefaultAzureCredential()
        case _:
            raise ValueError(
                f"Invalid Azure identity credential type: {env_settings.azure_identity_credential_type}"
            )


async def credentials_are_valid(credentials: AsyncTokenCredential) -> bool:
    _LOGGER.info("Validating Azure identity credentials")
    try:
        await credentials.get_token(
            env_settings.azure_identity_credentials_validate_token_url
        )
        _LOGGER.info("Azure identity credentials are valid")
        return True
    except Exception as e:
        _LOGGER.error(f"Azure identity credentials are invalid: {e}")
        return False


def get_project_client(
    credentials: AsyncTokenCredential, endpoint: str
) -> AIProjectClient:
    endpoint = env_settings.azure_ai_project_endpoint or endpoint

    if not endpoint:
        raise ValueError(
            "Azure AI project endpoint is not set from environment variables or configuration"
        )

    if env_settings.use_unique_private_endpoint_transport:
        transport = AsyncioRequestsTransport(connection_verify=certifi.where())
        return AIProjectClient(
            credential=credentials,
            endpoint=endpoint,
            transport=transport,
        )
    else:
        return AIProjectClient(
            credential=credentials,
            endpoint=endpoint,
        )
