"""Azure AI authentication handler.

Resolves Azure OpenAI authentication mode (API key vs DefaultAzureCredential)
and constructs the appropriate client configuration.
"""

import os
from dataclasses import dataclass, field
from typing import List, Optional

from azure.ai.evaluation import AzureOpenAIModelConfiguration
from azure.core.exceptions import ClientAuthenticationError

from common import (
    ENV_AZURE_AI_OPENAI_ENDPOINT,
    ENV_AZURE_AI_API_KEY,
    ENV_AZURE_AI_API_VERSION,
    ENV_AZURE_AI_MODEL_NAME,
)
from cli_logging.cli_logger import emit_structured_log
from cli_logging.logging_utils import Operation

from azure.identity import CredentialUnavailableError, DefaultAzureCredential

VALID_AUTH_MODES = ("key", "default-credential")

_CREDENTIAL_UNAVAILABLE_MESSAGE = (
    "DefaultAzureCredential could not find a valid credential. "
    "No Azure CLI, environment, or managed identity credential is available."
)

_CREDENTIAL_UNAVAILABLE_REMEDIATION = [
    "Sign in with Azure CLI: az login --tenant <your-tenant-id>",
    "Ensure the signed-in account has 'Cognitive Services OpenAI User' role "
    "on the target Azure OpenAI resource.",
    "Or switch to API key auth: --azure-ai-auth-mode key",
]

_PERMISSION_DENIED_REMEDIATION_ENTRA = [
    "Verify the signed-in identity has 'Cognitive Services OpenAI User' role "
    "on the target Azure OpenAI resource.",
    "Check role assignment: az role assignment list --assignee <your-identity> "
    "--scope /subscriptions/<sub>/resourceGroups/<rg>/providers/"
    "Microsoft.CognitiveServices/accounts/<resource>",
    "If using Azure CLI, try: az login --tenant <your-tenant-id> to refresh credentials.",
]

_PERMISSION_DENIED_REMEDIATION_KEY = [
    "Verify AZURE_AI_API_KEY is a valid key for the target Azure OpenAI resource.",
    "Regenerate the key in the Azure portal if it may have been rotated or revoked.",
    "Or switch to Entra auth: --azure-ai-auth-mode default-credential",
]


@dataclass(eq=False)
class AuthFailureOutcome(Exception):
    """Normalized user-facing auth error payload."""

    selected_mode: str
    selection_source: str  # "explicit-flag" or "auto-detect"
    category: str  # "missing-key" | "credential-unavailable" | "permission-denied" | "service-auth-failure" | "invalid-flag"
    message: str
    remediation: List[str] = field(default_factory=list)

    def __post_init__(self):
        super().__init__(self.message)

    def __str__(self):
        lines = [
            self.message,
            f"  Mode: {self.selected_mode} ({self.selection_source})",
        ]
        if self.remediation:
            lines.append("Remediation:")
            for step in self.remediation:
                lines.append(f"  - {step}")
        return "\n".join(lines)


def resolve_auth_mode(
    api_key_raw: Optional[str], auth_mode_flag: Optional[str]
) -> tuple[str, str]:
    """Resolve Azure AI authentication mode.

    Args:
        api_key_raw: Raw value of AZURE_AI_API_KEY (may be None, empty, or whitespace).
        auth_mode_flag: Value of --azure-ai-auth-mode flag (None if not provided).

    Returns:
        Tuple of (selected_mode, selection_source).
        selected_mode: "key" or "default-credential".
        selection_source: "explicit-flag" or "auto-detect".

    Raises:
        AuthFailureOutcome: If explicit mode validation fails.
    """
    api_key_usable = bool(api_key_raw and api_key_raw.strip())
    selection_source = "explicit-flag" if auth_mode_flag else "auto-detect"

    # Validate flag value if provided
    if auth_mode_flag and auth_mode_flag not in VALID_AUTH_MODES:
        raise AuthFailureOutcome(
            selected_mode="unknown",
            selection_source="explicit-flag",
            category="invalid-flag",
            message=(
                f"Invalid --azure-ai-auth-mode value: '{auth_mode_flag}'. "
                f"Accepted values: {', '.join(VALID_AUTH_MODES)}"
            ),
            remediation=[
                f"Use one of: {', '.join(VALID_AUTH_MODES)}",
            ],
        )

    # Explicit key mode
    if auth_mode_flag == "key":
        if not api_key_usable:
            raise AuthFailureOutcome(
                selected_mode="key",
                selection_source="explicit-flag",
                category="missing-key",
                message=(
                    "Authentication mode 'key' was explicitly selected but no usable "
                    "API key is available (AZURE_AI_API_KEY is missing, empty, or whitespace-only)."
                ),
                remediation=[
                    "Set AZURE_AI_API_KEY to a valid Azure OpenAI API key.",
                    "Or use --azure-ai-auth-mode default-credential to authenticate via Entra.",
                ],
            )
        selected_mode = "key"

    # Explicit default-credential mode
    elif auth_mode_flag == "default-credential":
        selected_mode = "default-credential"

    # Auto-detect
    else:
        selected_mode = "key" if api_key_usable else "default-credential"

    emit_structured_log(
        "info",
        f"Azure AI auth mode resolved: {selected_mode} (source: {selection_source})",
        operation=Operation.SETUP,
    )

    return selected_mode, selection_source


def verify_credential(selected_mode: str, selection_source: str = "auto-detect") -> None:
    """Fail-fast credential check for default-credential mode.

    Verifies that DefaultAzureCredential can be instantiated (i.e., at least one
    credential in the chain is available). Does NOT acquire a token — scope
    selection and token management are delegated to the Azure AI Evaluation SDK
    (FR-014, FR-016). No-op when mode is "key".

    Raises:
        AuthFailureOutcome: If no credential is available.
    """
    if selected_mode != "default-credential":
        return

    try:
        DefaultAzureCredential()
        emit_structured_log(
            "info",
            "DefaultAzureCredential initialized; token acquisition will be handled by the Azure AI Evaluation SDK.",
            operation=Operation.AUTHENTICATE,
        )
    except CredentialUnavailableError as exc:
        raise AuthFailureOutcome(
            selected_mode=selected_mode,
            selection_source=selection_source,
            category="credential-unavailable",
            message=_CREDENTIAL_UNAVAILABLE_MESSAGE,
            remediation=list(_CREDENTIAL_UNAVAILABLE_REMEDIATION),
        ) from exc


def has_azure_openai() -> bool:
    """Determine whether Azure LLM-based evaluators can be configured.

    Returns True when endpoint, api_version, and model_name env vars are present.
    Auth validity (API key or Entra credential) is NOT checked here — it is
    verified at request time by the SDK to avoid eager pre-validation (FR-016).
    """
    endpoint = os.environ.get(ENV_AZURE_AI_OPENAI_ENDPOINT)
    api_version = os.environ.get(ENV_AZURE_AI_API_VERSION)
    model_name = os.environ.get(ENV_AZURE_AI_MODEL_NAME)

    if not (endpoint and api_version and model_name):
        return False

    return True


def build_azure_openai_client(selected_mode: str) -> AzureOpenAIModelConfiguration:
    """Construct AzureOpenAIModelConfiguration based on auth decision.

    Args:
        selected_mode: "key" or "default-credential".

    Returns:
        Configured AzureOpenAIModelConfiguration instance.
    """
    endpoint = os.environ.get(ENV_AZURE_AI_OPENAI_ENDPOINT)
    api_version = os.environ.get(ENV_AZURE_AI_API_VERSION)
    model_name = os.environ.get(ENV_AZURE_AI_MODEL_NAME)

    if selected_mode == "key":
        api_key = os.environ.get(ENV_AZURE_AI_API_KEY)
        return AzureOpenAIModelConfiguration(
            azure_endpoint=endpoint,
            api_key=api_key,
            api_version=api_version,
            azure_deployment=model_name,
        )
    else:
        # Entra-based authentication: omit api_key and let the SDK's built-in
        # token provider handle DefaultAzureCredential automatically.
        # Per FR-016, we do NOT manually call get_token() or manage token refresh.
        return AzureOpenAIModelConfiguration(
            azure_endpoint=endpoint,
            api_version=api_version,
            azure_deployment=model_name,
        )


def wrap_credential_error(
    exc: Exception,
    selected_mode: str,
    selection_source: str = "auto-detect",
) -> Optional[AuthFailureOutcome]:
    """Translate SDK credential/auth exceptions into user-friendly AuthFailureOutcome.

    Returns AuthFailureOutcome if the exception is a recognized credential or auth error,
    or None if the exception is unrelated to authentication.

    Args:
        exc: The caught exception.
        selected_mode: The resolved auth mode ("key" or "default-credential").
        selection_source: How the mode was selected ("explicit-flag" or "auto-detect").
    """
    # CredentialUnavailableError: no valid credential found (not logged in)
    if isinstance(exc, CredentialUnavailableError):
        return AuthFailureOutcome(
            selected_mode=selected_mode,
            selection_source=selection_source,
            category="credential-unavailable",
            message=_CREDENTIAL_UNAVAILABLE_MESSAGE,
            remediation=list(_CREDENTIAL_UNAVAILABLE_REMEDIATION),
        )

    # ClientAuthenticationError: credential exists but is rejected (RBAC / permission)
    if isinstance(exc, ClientAuthenticationError):
        remediation = (
            _PERMISSION_DENIED_REMEDIATION_KEY if selected_mode == "key"
            else _PERMISSION_DENIED_REMEDIATION_ENTRA
        )

        return AuthFailureOutcome(
            selected_mode=selected_mode,
            selection_source=selection_source,
            category="permission-denied",
            message=f"Azure authentication failed: {exc}",
            remediation=list(remediation),
        )

    return None