"""Azure SDK client helpers."""

from __future__ import annotations

import asyncio
from typing import Any

from azure.identity import (
    AzureCliCredential,
    ChainedTokenCredential,
    DefaultAzureCredential,
)
from azure.identity._constants import KnownAuthorities
from azure.mgmt.resourcegraph import ResourceGraphClient
from azure.mgmt.resourcegraph.models import QueryRequest
from azure.mgmt.subscription import SubscriptionClient

from ..adt_types import (
    AzureClientError,
    AzureDiscoveryRequest,
    AzureEnvironment,
    AzureEnvironmentConfig,
)
from ..adt_types.models import ScaleControlConfig
from ..utils.rate_control import AdaptiveRateController, RateLimitConfig
from .logging import get_logger

LOGGER = get_logger()

_ENVIRONMENT_MAP: dict[AzureEnvironment, AzureEnvironmentConfig] = {
    AzureEnvironment.AZURE_PUBLIC: AzureEnvironmentConfig(
        name=AzureEnvironment.AZURE_PUBLIC,
        authority_host=KnownAuthorities.AZURE_PUBLIC_CLOUD,
        resource_manager="https://management.azure.com/",
        graph_endpoint="https://graph.microsoft.com",
        storage_suffix="core.windows.net",
    ),
    AzureEnvironment.AZURE_GOV: AzureEnvironmentConfig(
        name=AzureEnvironment.AZURE_GOV,
        authority_host=KnownAuthorities.AZURE_GOVERNMENT,
        resource_manager="https://management.usgovcloudapi.net/",
        graph_endpoint="https://graph.microsoft.us",
        storage_suffix="core.usgovcloudapi.net",
    ),
    AzureEnvironment.AZURE_CHINA: AzureEnvironmentConfig(
        name=AzureEnvironment.AZURE_CHINA,
        authority_host=KnownAuthorities.AZURE_CHINA,
        resource_manager="https://management.chinacloudapi.cn/",
        graph_endpoint="https://microsoftgraph.chinacloudapi.cn",
        storage_suffix="core.chinacloudapi.cn",
    ),
    AzureEnvironment.AZURE_GERMANY: AzureEnvironmentConfig(
        name=AzureEnvironment.AZURE_GERMANY,
        authority_host=KnownAuthorities.AZURE_GERMANY,
        resource_manager="https://management.microsoftazure.de/",
        graph_endpoint="https://graph.microsoft.de",
        storage_suffix="core.cloudapi.de",
    ),
    AzureEnvironment.AZURE_STACK: AzureEnvironmentConfig(
        name=AzureEnvironment.AZURE_STACK,
        authority_host=KnownAuthorities.AZURE_PUBLIC_CLOUD,
        resource_manager="https://management.local.azurestack.external/",
        graph_endpoint="https://graph.microsoft.com",
        storage_suffix="local.azurestack.external",
    ),
}


def build_environment_config(
    environment: AzureEnvironment,
) -> AzureEnvironmentConfig:
    """Return endpoints for the requested cloud."""

    if environment not in _ENVIRONMENT_MAP:
        raise AzureClientError(f"Unsupported environment: {environment}")
    return _ENVIRONMENT_MAP[environment]


def get_credential(
    request: AzureDiscoveryRequest,
    config: AzureEnvironmentConfig,
) -> ChainedTokenCredential:
    """Return credential chain honoring CLI preference."""

    credential_chain = []
    if request.prefer_cli_credentials:
        credential_chain.append(
            AzureCliCredential(
                authority=config.authority_host,
                tenant_id=request.tenant_id,
            )
        )
    credential_chain.append(
        DefaultAzureCredential(
            authority=config.authority_host,
            exclude_interactive_browser_credential=False,
        )
    )
    return ChainedTokenCredential(*credential_chain)


async def ensure_subscription_ids(
    request: AzureDiscoveryRequest,
    credential: ChainedTokenCredential,
    base_url: str | None = None,
) -> list[str]:
    """Resolve the subscription list, querying Azure if needed."""

    if request.subscriptions:
        return sorted({sub.strip() for sub in request.subscriptions if sub})

    loop = asyncio.get_running_loop()

    # Add explicit credential scopes for sovereign cloud support
    credential_scopes = [base_url + ".default"] if base_url else None
    client = SubscriptionClient(
        credential=credential,
        base_url=base_url,
        credential_scopes=credential_scopes,
    )

    def _list_subscriptions() -> list[str]:
        return [sub.subscription_id for sub in client.subscriptions.list()]

    try:
        results: list[str] = await loop.run_in_executor(
            None, _list_subscriptions
        )
    except Exception as exc:  # noqa: BLE001
        LOGGER.error(
            "Failed to enumerate subscriptions",
            extra={"context": {"error": str(exc)}},
        )
        raise AzureClientError("Unable to enumerate subscriptions") from exc
    if not results:
        raise AzureClientError("No subscriptions available for discovery")
    return sorted(results)


async def query_resource_graph(
    credential: ChainedTokenCredential,
    query: str,
    subscriptions: list[str],
    batch_size: int,
    base_url: str | None = None,
    max_objects: int = 0,
    scale_controls: ScaleControlConfig | None = None,
) -> list[dict[str, Any]]:
    """Execute a Resource Graph query with pagination and adaptive rate control."""

    if not query or not query.strip():
        raise AzureClientError("Resource Graph query cannot be empty")

    # Initialize adaptive rate controller when scale controls are enabled
    rate_controller: AdaptiveRateController | None = None
    if scale_controls and scale_controls.enabled:
        rate_controller = AdaptiveRateController(
            RateLimitConfig(
                initial_rps=scale_controls.initial_rps,
                min_rps=scale_controls.min_rps,
                max_rps=scale_controls.max_rps,
            )
        )
        batch_size = scale_controls.initial_batch_size

    # Add explicit credential scopes for sovereign cloud support
    credential_scopes = [base_url + ".default"] if base_url else None
    client = ResourceGraphClient(
        credential=credential,
        base_url=base_url,
        credential_scopes=credential_scopes,
    )
    payload = QueryRequest(
        subscriptions=subscriptions,
        query=query,
        options={"resultFormat": "objectArray", "top": batch_size},
    )
    loop = asyncio.get_running_loop()
    results: list[dict[str, Any]] = []
    skip_token: str | None = None

    while True:
        if max_objects and len(results) >= max_objects:
            break

        if skip_token:
            payload.options["skipToken"] = skip_token

        # Acquire rate limit token before making request
        if rate_controller:
            await rate_controller.acquire()

        def _run_query() -> Any:
            return client.resources(payload)

        try:
            response = await loop.run_in_executor(None, _run_query)
            if rate_controller:
                await rate_controller.record_success()
        except Exception as exc:  # noqa: BLE001
            error_str = str(exc).lower()
            if rate_controller and ("429" in error_str or "throttl" in error_str):
                await rate_controller.record_throttle()
                # Retry once after throttle
                await rate_controller.acquire()
                try:
                    response = await loop.run_in_executor(None, _run_query)
                    await rate_controller.record_success()
                except Exception as retry_exc:  # noqa: BLE001
                    LOGGER.error(
                        "Resource Graph query failed after retry",
                        extra={"context": {"error": str(retry_exc)}},
                    )
                    raise AzureClientError("Resource Graph query failure") from retry_exc
            else:
                LOGGER.error(
                    "Resource Graph query failed",
                    extra={"context": {"error": str(exc)}},
                )
                raise AzureClientError("Resource Graph query failure") from exc

        results.extend(response.data or [])
        if max_objects:
            results = results[:max_objects]

        skip_token = getattr(response, "skip_token", None)
        if not skip_token:
            break

    return results
