"""
AWS Cognito authentication for AiCippy.

Implements direct CLI authentication against Cognito (auth.aivibe.cloud)
with secure token storage and automatic refresh.

This module provides:
- Direct username/password authentication via Cognito USER_PASSWORD_AUTH
- OTP-based authentication via Cognito CUSTOM_AUTH
- Secure credential storage via OS keychain
- JWT validation against Cognito JWKS
- Automatic token refresh
- Rate-limit (HTTP 429) handling with Retry-After support
- Connectivity pre-check via TCP probe
- Structured error handling
"""

from __future__ import annotations

import asyncio
import contextlib
import threading
import time as _time
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any, Final

import httpx
from jose import JWTError, jwt

from aicippy.auth.keychain import KeychainStorage
from aicippy.auth.models import AuthResult, TokenInfo, UserInfo
from aicippy.config import get_settings
from aicippy.utils.correlation import get_correlation_id
from aicippy.utils.logging import get_logger
from aicippy.utils.retry import async_retry

if TYPE_CHECKING:
    from types import TracebackType

logger = get_logger(__name__)

# ============================================================================
# Constants
# ============================================================================

HTTP_TIMEOUT_SECONDS: Final[float] = 30.0
TOKEN_EXPIRY_BUFFER_SECONDS: Final[int] = 300  # Refresh 5 min before expiry
_JWKS_TTL_SECONDS: Final[int] = 86400  # 24 hours
_DEFAULT_RATE_LIMIT_WAIT_SECONDS: Final[int] = 60
_DEFAULT_TOKEN_EXPIRY_SECONDS: Final[int] = 3600


# ============================================================================
# Cognito API Error
# ============================================================================


class CognitoAPIError(Exception):
    """Error from Cognito REST API.

    Raised when the Cognito IDP HTTP endpoint returns a non-200 response.

    Attributes:
        error_type: The Cognito error code (e.g. ``NotAuthorizedException``).
        message: Human-readable error description from Cognito.
        retry_after: Seconds to wait before retrying, or ``None`` if not a
            rate-limit error.  Populated from the ``Retry-After`` header when
            the server returns HTTP 429 / ``TooManyRequestsException``.
    """

    __slots__ = ("error_type", "message", "retry_after")

    def __init__(
        self,
        error_type: str,
        message: str,
        retry_after: int | None = None,
    ) -> None:
        self.error_type = error_type
        self.message = message
        self.retry_after = retry_after
        super().__init__(f"{error_type}: {message}")

    @property
    def is_rate_limited(self) -> bool:
        """Return ``True`` if this error represents a rate-limit (429)."""
        return self.error_type == "TooManyRequestsException"


# ============================================================================
# Cognito Authentication Handler
# ============================================================================


class CognitoAuth:
    """
    AWS Cognito authentication handler.

    Authenticates directly against Cognito via USER_PASSWORD_AUTH.
    No local servers, no browser callbacks. Pure CLI flow.

    Example:
        >>> auth = CognitoAuth()
        >>> if not auth.is_authenticated():
        ...     result = await auth.login("dev@aivibe.in", "password")
        ...     if result.success:
        ...         print(f"Logged in as {result.user_email}")
        >>> tokens = auth.get_current_tokens()
    """

    __slots__ = (
        "_client_id",
        "_client_loop",
        "_cognito_url",
        "_http_client",
        "_jwks_cache",
        "_jwks_fetched_at",
        "_jwks_url",
        "_keyring_available",
        "_refresh_lock",
        "_region",
        "_settings",
        "_storage",
        "_user_pool_id",
        "_userinfo_url",
    )

    def __init__(self) -> None:
        """Initialize Cognito auth handler."""
        self._settings = get_settings()
        self._storage = KeychainStorage()
        self._http_client: httpx.AsyncClient | None = None
        self._client_loop: asyncio.AbstractEventLoop | None = None

        # Cognito config
        self._client_id = self._settings.cognito_client_id
        self._user_pool_id = self._settings.cognito_user_pool_id
        self._region = self._settings.aws_region

        # Cognito IDP REST endpoint (unsigned public operations)
        self._cognito_url = f"https://cognito-idp.{self._region}.amazonaws.com/"

        # Endpoints for token operations
        domain = self._settings.cognito_domain
        self._userinfo_url = f"{domain}/oauth2/userInfo"
        self._jwks_url = (
            f"https://cognito-idp.{self._region}.amazonaws.com/"
            f"{self._user_pool_id}/.well-known/jwks.json"
        )

        # JWKS cache with TTL tracking
        self._jwks_cache: dict[str, Any] | None = None
        self._jwks_fetched_at: float | None = None

        # Keyring availability flag - set to False if keyring storage fails
        self._keyring_available: bool = True

        # Thread safety for token refresh
        self._refresh_lock = threading.Lock()

        logger.debug(
            "cognito_auth_initialized",
            user_pool_id=self._user_pool_id,
            region=self._region,
        )

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create HTTP client, recreating if the event loop changed."""
        current_loop = asyncio.get_running_loop()
        if (
            self._http_client is None
            or self._http_client.is_closed
            or self._client_loop is not current_loop
        ):
            # Close stale client if it exists (may fail on dead loop)
            if self._http_client is not None and not self._http_client.is_closed:
                with contextlib.suppress(Exception):
                    await self._http_client.aclose()
            self._http_client = httpx.AsyncClient(
                timeout=HTTP_TIMEOUT_SECONDS,
                follow_redirects=True,
                headers={"User-Agent": "AiCippy-CLI/1.0"},
            )
            self._client_loop = current_loop
        return self._http_client

    async def _cognito_api_call(
        self,
        action: str,
        payload: dict[str, Any],
    ) -> dict[str, Any]:
        """Make a direct HTTP call to Cognito IDP API.

        Uses the public Cognito REST API which does not require IAM signing
        for public client operations like ``InitiateAuth`` on a client with
        no secret.

        Args:
            action: The Cognito API action (e.g. ``InitiateAuth``).
            payload: The JSON request body.

        Returns:
            Parsed JSON response from Cognito.

        Raises:
            CognitoAPIError: When Cognito returns a non-200 status.  If the
                response is HTTP 429 (``TooManyRequestsException``), the
                ``retry_after`` attribute is populated from the
                ``Retry-After`` header when present.
        """
        headers = {
            "Content-Type": "application/x-amz-json-1.1",
            "X-Amz-Target": (f"AWSCognitoIdentityProviderService.{action}"),
        }
        client = await self._get_client()
        response = await client.post(
            self._cognito_url,
            json=payload,
            headers=headers,
        )

        if response.status_code != 200:
            try:
                error_data = response.json()
            except Exception:
                logger.warning(
                    "cognito_error_json_parse_failed",
                    status_code=response.status_code,
                    exc_info=True,
                )
                error_data = {}
            error_type = error_data.get("__type", "UnknownError")
            # Cognito returns the error type as a fully qualified class
            # name; strip the prefix to get just the error code.
            if "#" in error_type:
                error_type = error_type.rsplit("#", 1)[-1]
            error_message = error_data.get("message", "Unknown error")

            # --- Rate-limit (HTTP 429) handling ---
            retry_after: int | None = None
            if response.status_code == 429 or error_type == "TooManyRequestsException":
                # Normalise error_type for callers regardless of how
                # Cognito signals the throttle.
                error_type = "TooManyRequestsException"

                # Parse Retry-After header (seconds) if present.
                raw_retry = response.headers.get("Retry-After")
                if raw_retry is not None:
                    try:
                        retry_after = int(raw_retry)
                    except (ValueError, TypeError):
                        retry_after = _DEFAULT_RATE_LIMIT_WAIT_SECONDS
                else:
                    retry_after = _DEFAULT_RATE_LIMIT_WAIT_SECONDS

                logger.warning(
                    "cognito_rate_limited",
                    action=action,
                    retry_after=retry_after,
                    status_code=response.status_code,
                )

            raise CognitoAPIError(
                error_type,
                error_message,
                retry_after=retry_after,
            )

        return response.json()

    # ------------------------------------------------------------------ #
    # Connectivity pre-check
    # ------------------------------------------------------------------ #

    async def check_connectivity(self) -> bool:
        """Check whether the Cognito endpoint is reachable.

        Uses a lightweight TCP probe (no HTTP/TLS handshake) via
        :func:`aicippy.utils.network.check_cognito_reachable`.

        Returns:
            ``True`` if the Cognito IDP endpoint for the configured
            region is reachable, ``False`` otherwise.
        """
        from aicippy.utils.network import check_cognito_reachable

        reachable = check_cognito_reachable(region=self._region)
        logger.debug(
            "cognito_connectivity_check",
            region=self._region,
            reachable=reachable,
        )
        return reachable

    async def close(self) -> None:
        """Close HTTP client and cleanup resources."""
        if self._http_client is not None:
            with contextlib.suppress(Exception):
                await self._http_client.aclose()
            self._http_client = None
            self._client_loop = None
        logger.debug("cognito_auth_closed")

    def is_authenticated(self) -> bool:
        """Check if user is authenticated with valid tokens."""
        try:
            return self._storage.has_valid_tokens()
        except Exception as e:
            logger.warning("auth_check_failed", error=str(e))
            return False

    def _is_token_near_expiry(self, tokens: TokenInfo) -> bool:
        """Check if tokens are expired or within the buffer window.

        Returns True if the token should be refreshed proactively.
        """
        buffer_dt = timedelta(seconds=TOKEN_EXPIRY_BUFFER_SECONDS)
        return datetime.now(UTC) >= (tokens.expires_at - buffer_dt)

    def get_current_tokens(self) -> TokenInfo | None:
        """Get current tokens, attempting refresh if expired or near expiry.

        Proactively refreshes tokens within TOKEN_EXPIRY_BUFFER_SECONDS
        (5 minutes) of expiry to prevent request failures.
        """
        try:
            tokens = self._storage.get_tokens()
        except Exception as e:
            logger.warning("token_retrieval_failed", error=str(e))
            return None

        if tokens is None:
            return None

        # Proactively refresh tokens that are expired or near expiry
        if self._is_token_near_expiry(tokens):
            logger.debug(
                "tokens_near_expiry_refreshing",
                expired=tokens.is_expired(),
                buffer_seconds=TOKEN_EXPIRY_BUFFER_SECONDS,
            )
            try:
                refreshed = self._refresh_tokens_sync(
                    tokens.refresh_token,
                )
                if refreshed:
                    return refreshed
            except Exception as e:
                logger.warning(
                    "token_refresh_failed",
                    error=str(e),
                )
            # If refresh failed but token is not yet expired, return it
            if not tokens.is_expired():
                return tokens
            return None

        return tokens

    # ------------------------------------------------------------------ #
    # Rate-limit helper
    # ------------------------------------------------------------------ #

    @staticmethod
    def _rate_limit_failure_message(
        error: CognitoAPIError,
    ) -> str:
        """Build a user-friendly rate-limit failure message."""
        wait = error.retry_after or _DEFAULT_RATE_LIMIT_WAIT_SECONDS
        return f"Too many login attempts. Please wait {wait} seconds before trying again."

    async def login(self, email: str, password: str) -> AuthResult:
        """
        Authenticate directly with Cognito using email and password.

        Uses USER_PASSWORD_AUTH flow - no browser, no localhost, no
        callbacks.  Authenticates against auth.aivibe.cloud Cognito
        server.

        Args:
            email: User's email address.
            password: User's password.

        Returns:
            AuthResult with success/failure status.
        """
        correlation_id = get_correlation_id()
        logger.info(
            "direct_login_started",
            correlation_id=correlation_id,
        )

        try:
            response = await self._cognito_api_call(
                "InitiateAuth",
                {
                    "ClientId": self._client_id,
                    "AuthFlow": "USER_PASSWORD_AUTH",
                    "AuthParameters": {
                        "USERNAME": email,
                        "PASSWORD": password,
                    },
                },
            )
        except CognitoAPIError as e:
            logger.warning(
                "login_failed",
                error_code=e.error_type,
                correlation_id=correlation_id,
            )

            if e.is_rate_limited:
                return AuthResult.failure(
                    self._rate_limit_failure_message(e),
                )
            if e.error_type in (
                "NotAuthorizedException",
                "UserNotFoundException",
            ):
                return AuthResult.failure("Invalid email or password")
            if e.error_type == "UserNotConfirmedException":
                return AuthResult.failure(
                    "Account not verified. Check your email.",
                )
            if e.error_type == "PasswordResetRequiredException":
                return AuthResult.failure(
                    "Password reset required. Use Email OTP login or contact support@aivibe.in"
                )
            return AuthResult.failure(
                f"Authentication failed: {e.message}",
            )
        except Exception as e:
            logger.exception(
                "login_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return AuthResult.failure(f"Login error: {e}")

        # Handle challenges (NEW_PASSWORD_REQUIRED, MFA, etc.)
        if "ChallengeName" in response:
            challenge = response["ChallengeName"]
            logger.warning(
                "auth_challenge_received",
                challenge=challenge,
                correlation_id=correlation_id,
            )
            if challenge == "NEW_PASSWORD_REQUIRED":
                return AuthResult.failure(
                    "New password required. Use Email OTP login or contact support@aivibe.in"
                )
            return AuthResult.failure(
                f"Auth challenge not supported in CLI: {challenge}",
            )

        # Extract tokens from successful auth via shared helper
        auth_result_data = response.get("AuthenticationResult", {})
        return self._process_auth_tokens(
            auth_result_data,
            email,
            correlation_id,
        )

    def _is_token_expired(self) -> bool:
        """Check if the currently stored tokens are expired."""
        try:
            tokens = self._storage.get_tokens()
            if tokens is None:
                return True
            return tokens.is_expired()
        except Exception as e:
            logger.warning(
                "token_expiry_check_failed",
                error=str(e),
            )
            return True

    def _refresh_tokens_sync(
        self,
        refresh_token: str,
    ) -> TokenInfo | None:
        """Refresh tokens synchronously using Cognito API.

        Thread-safe: uses _refresh_lock to prevent concurrent refresh
        attempts. Double-checks token near-expiry after acquiring the
        lock to avoid redundant refreshes.
        """
        with self._refresh_lock:
            # Double-check: another thread may have already refreshed.
            # Must use _is_token_near_expiry (not _is_token_expired) to
            # match the caller's refresh trigger condition.
            try:
                current_tokens = self._storage.get_tokens()
            except Exception as e:
                logger.warning(
                    "token_retrieval_after_recheck_failed",
                    error=str(e),
                )
                current_tokens = None

            if current_tokens is not None and not self._is_token_near_expiry(current_tokens):
                return current_tokens

            correlation_id = get_correlation_id()

            try:
                http_response = httpx.post(
                    self._cognito_url,
                    json={
                        "ClientId": self._client_id,
                        "AuthFlow": "REFRESH_TOKEN_AUTH",
                        "AuthParameters": {
                            "REFRESH_TOKEN": refresh_token,
                        },
                    },
                    headers={
                        "Content-Type": "application/x-amz-json-1.1",
                        "X-Amz-Target": ("AWSCognitoIdentityProviderService.InitiateAuth"),
                    },
                    timeout=HTTP_TIMEOUT_SECONDS,
                )

                if http_response.status_code != 200:
                    try:
                        error_data = http_response.json()
                        error_code = error_data.get(
                            "__type",
                            "UnknownError",
                        )
                        if "#" in error_code:
                            error_code = error_code.rsplit("#", 1)[-1]
                    except Exception:
                        logger.warning(
                            "token_refresh_json_parse_failed",
                            status_code=http_response.status_code,
                            exc_info=True,
                        )
                        error_code = f"HTTP_{http_response.status_code}"

                    # Handle rate limiting (HTTP 429 or Cognito TooManyRequests)
                    is_rate_limited = (
                        http_response.status_code == 429 or error_code == "TooManyRequestsException"
                    )
                    if is_rate_limited:
                        retry_after = float(http_response.headers.get("Retry-After", "5"))
                        logger.warning(
                            "token_refresh_rate_limited",
                            retry_after=retry_after,
                            correlation_id=correlation_id,
                        )
                        import time

                        time.sleep(min(retry_after, 10.0))
                        # Retry once after backoff
                        retry_resp = httpx.post(
                            self._cognito_url,
                            json={
                                "ClientId": self._client_id,
                                "AuthFlow": "REFRESH_TOKEN_AUTH",
                                "AuthParameters": {
                                    "REFRESH_TOKEN": refresh_token,
                                },
                            },
                            headers={
                                "Content-Type": "application/x-amz-json-1.1",
                                "X-Amz-Target": ("AWSCognitoIdentityProviderService.InitiateAuth"),
                            },
                            timeout=HTTP_TIMEOUT_SECONDS,
                        )
                        if retry_resp.status_code == 200:
                            response = retry_resp.json()
                            # Fall through to token processing below
                        else:
                            logger.warning(
                                "token_refresh_rate_limit_retry_failed",
                                status_code=retry_resp.status_code,
                                correlation_id=correlation_id,
                            )
                            return None
                    else:
                        logger.warning(
                            "token_refresh_failed",
                            error_code=error_code,
                            correlation_id=correlation_id,
                        )
                        return None
                else:
                    response = http_response.json()
            except Exception as e:
                logger.warning(
                    "token_refresh_error",
                    error=str(e),
                    correlation_id=correlation_id,
                )
                return None

            auth_result = response.get("AuthenticationResult", {})
            access_token = auth_result.get("AccessToken")
            id_token = auth_result.get("IdToken")
            expires_in = auth_result.get("ExpiresIn", _DEFAULT_TOKEN_EXPIRY_SECONDS)

            if not access_token or not id_token:
                return None

            expires_at = datetime.now(UTC) + timedelta(seconds=expires_in)

            tokens = TokenInfo(
                access_token=access_token,
                id_token=id_token,
                refresh_token=auth_result.get(
                    "RefreshToken",
                    refresh_token,
                ),
                expires_at=expires_at,
                token_type=auth_result.get("TokenType", "Bearer"),
            )

            if self._keyring_available:
                try:
                    self._storage.store_tokens(tokens)
                except Exception as e:
                    logger.warning(
                        "token_storage_after_refresh_failed",
                        error=str(e),
                        correlation_id=correlation_id,
                    )
                    self._keyring_available = False

            logger.info(
                "tokens_refreshed",
                expires_at=expires_at.isoformat(),
                correlation_id=correlation_id,
            )

            return tokens

    @async_retry(max_attempts=2, min_wait=0.5, max_wait=2.0)
    async def _get_user_info(
        self,
        access_token: str,
    ) -> UserInfo | None:
        """Get user information from Cognito userInfo endpoint."""
        correlation_id = get_correlation_id()
        client = await self._get_client()

        try:
            response = await client.get(
                self._userinfo_url,
                headers={
                    "Authorization": f"Bearer {access_token}",
                },
            )
        except httpx.HTTPError as e:
            logger.error(
                "userinfo_http_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return None

        if response.status_code != 200:
            logger.error(
                "userinfo_failed",
                status=response.status_code,
                correlation_id=correlation_id,
            )
            return None

        try:
            data = response.json()
        except ValueError as e:
            logger.error(
                "userinfo_json_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return None

        sub = data.get("sub")
        if not sub:
            logger.error(
                "userinfo_missing_sub",
                correlation_id=correlation_id,
            )
            return None

        return UserInfo(
            sub=sub,
            email=data.get("email", ""),
            email_verified=data.get("email_verified", False),
            name=data.get("name"),
            phone_number=data.get("phone_number"),
            phone_number_verified=data.get(
                "phone_number_verified",
                False,
            ),
        )

    def logout(self) -> None:
        """Log out and clear all stored credentials."""
        correlation_id = get_correlation_id()
        try:
            self._storage.clear_all()
            self._jwks_cache = None
            self._jwks_fetched_at = None
            logger.info("logged_out", correlation_id=correlation_id)
        except Exception as e:
            logger.warning(
                "logout_storage_error",
                error=str(e),
                correlation_id=correlation_id,
            )

    def _is_jwks_cache_expired(self) -> bool:
        """Check if the JWKS cache has exceeded its TTL."""
        if self._jwks_cache is None or self._jwks_fetched_at is None:
            return True
        return (_time.monotonic() - self._jwks_fetched_at) >= _JWKS_TTL_SECONDS

    async def validate_token(self, token: str) -> bool:
        """Validate a JWT token against Cognito's JWKS."""
        correlation_id = get_correlation_id()

        try:
            if self._is_jwks_cache_expired():
                # Clear stale cache so we re-fetch
                self._jwks_cache = None
                self._jwks_fetched_at = None
                client = await self._get_client()
                try:
                    response = await client.get(self._jwks_url)
                    if response.status_code != 200:
                        logger.error(
                            "jwks_fetch_failed",
                            status=response.status_code,
                            correlation_id=correlation_id,
                        )
                        return False
                    self._jwks_cache = response.json()
                    self._jwks_fetched_at = _time.monotonic()
                except httpx.HTTPError as e:
                    logger.error(
                        "jwks_fetch_http_error",
                        error=str(e),
                        correlation_id=correlation_id,
                    )
                    return False

            if self._jwks_cache is None:
                logger.error(
                    "jwks_cache_unavailable",
                    correlation_id=correlation_id,
                )
                return False

            unverified = jwt.get_unverified_header(token)
            kid = unverified.get("kid")

            if not kid:
                logger.warning(
                    "token_missing_kid",
                    correlation_id=correlation_id,
                )
                return False

            key = None
            for k in self._jwks_cache.get("keys", []):
                if k.get("kid") == kid:
                    key = k
                    break

            if not key:
                # Key not found -- may be a key rotation. Re-fetch
                # JWKS once and retry the lookup before giving up.
                logger.info(
                    "jwks_key_not_found_refetching",
                    kid=kid,
                    correlation_id=correlation_id,
                )
                self._jwks_cache = None
                self._jwks_fetched_at = None
                try:
                    client = await self._get_client()
                    response = await client.get(self._jwks_url)
                    if response.status_code != 200:
                        logger.error(
                            "jwks_refetch_failed",
                            status=response.status_code,
                            correlation_id=correlation_id,
                        )
                        return False
                    self._jwks_cache = response.json()
                    self._jwks_fetched_at = _time.monotonic()
                except httpx.HTTPError as e:
                    logger.error(
                        "jwks_refetch_http_error",
                        error=str(e),
                        correlation_id=correlation_id,
                    )
                    return False

                # Retry key lookup after re-fetch
                for k in self._jwks_cache.get("keys", []):
                    if k.get("kid") == kid:
                        key = k
                        break

                if not key:
                    logger.warning(
                        "jwks_key_not_found_after_refetch",
                        kid=kid,
                        correlation_id=correlation_id,
                    )
                    return False

            jwt.decode(
                token,
                key,
                algorithms=["RS256"],
                audience=self._client_id,
                issuer=(f"https://cognito-idp.{self._region}.amazonaws.com/{self._user_pool_id}"),
            )

            return True

        except JWTError as e:
            logger.warning(
                "token_validation_failed",
                error=str(e),
                correlation_id=correlation_id,
            )
            return False
        except Exception as e:
            logger.exception(
                "token_validation_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return False

    async def otp_login(self, email: str) -> AuthResult:
        """Initiate OTP-based login via Cognito CUSTOM_AUTH flow.

        Sends a one-time password to the user's email address via
        Cognito Lambda triggers (Define/Create/Verify Auth Challenge).

        Args:
            email: User's registered email address.

        Returns:
            AuthResult with ``session`` stored in ``session`` field on
            challenge, or a failure AuthResult if the initiation fails.
            On success (CUSTOM_CHALLENGE returned), the result has:
              - ``success=False`` (login not yet complete)
              - ``session`` containing the session token for
                ``verify_otp``
              - ``user_email`` set to the email
        """
        correlation_id = get_correlation_id()
        logger.info(
            "otp_login_initiated",
            email=email,
            correlation_id=correlation_id,
        )

        try:
            response = await self._cognito_api_call(
                "InitiateAuth",
                {
                    "ClientId": self._client_id,
                    "AuthFlow": "CUSTOM_AUTH",
                    "AuthParameters": {
                        "USERNAME": email,
                    },
                },
            )
        except CognitoAPIError as e:
            logger.warning(
                "otp_login_initiation_failed",
                error_code=e.error_type,
                correlation_id=correlation_id,
            )
            if e.is_rate_limited:
                return AuthResult.failure(
                    self._rate_limit_failure_message(e),
                )
            if e.error_type == "UserNotFoundException":
                return AuthResult.failure(
                    "No account found for this email address.",
                )
            if e.error_type == "NotAuthorizedException":
                return AuthResult.failure(
                    "OTP login is not enabled for this account.",
                )
            return AuthResult.failure(
                f"OTP initiation failed: {e.message}",
            )
        except Exception as e:
            logger.exception(
                "otp_login_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return AuthResult.failure(f"OTP login error: {e}")

        # Cognito should return CUSTOM_CHALLENGE requiring the OTP code
        challenge_name = response.get("ChallengeName")
        session = response.get("Session")

        if challenge_name == "CUSTOM_CHALLENGE" and session:
            logger.info(
                "otp_challenge_received",
                email=email,
                correlation_id=correlation_id,
            )
            # Return the session token via a partial result; caller
            # uses verify_otp() with this session to complete login.
            return AuthResult(
                success=False,
                user_email=email,
                session=session,
            )

        # If we got tokens directly (unlikely for CUSTOM_AUTH), handle
        auth_result_data = response.get("AuthenticationResult")
        if auth_result_data:
            return self._process_auth_tokens(
                auth_result_data,
                email,
                correlation_id,
            )

        return AuthResult.failure(f"Unexpected challenge from Cognito: {challenge_name or 'none'}")

    async def verify_otp(
        self,
        email: str,
        otp_code: str,
        session: str,
    ) -> AuthResult:
        """Verify the OTP code to complete CUSTOM_AUTH login.

        Args:
            email: User's email address (same as in ``otp_login``).
            otp_code: The 6-digit OTP code the user received via email.
            session: The session token returned by ``otp_login``.

        Returns:
            AuthResult with tokens on success, or failure details.
        """
        correlation_id = get_correlation_id()
        logger.info(
            "otp_verification_started",
            email=email,
            correlation_id=correlation_id,
        )

        try:
            response = await self._cognito_api_call(
                "RespondToAuthChallenge",
                {
                    "ClientId": self._client_id,
                    "ChallengeName": "CUSTOM_CHALLENGE",
                    "Session": session,
                    "ChallengeResponses": {
                        "USERNAME": email,
                        "ANSWER": otp_code,
                    },
                },
            )
        except CognitoAPIError as e:
            logger.warning(
                "otp_verification_failed",
                error_code=e.error_type,
                correlation_id=correlation_id,
            )
            if e.is_rate_limited:
                return AuthResult.failure(
                    self._rate_limit_failure_message(e),
                )
            if e.error_type == "NotAuthorizedException":
                return AuthResult.failure(
                    "Invalid or expired OTP code.",
                )
            if e.error_type == "CodeMismatchException":
                return AuthResult.failure(
                    "Incorrect OTP code. Please try again.",
                )
            if e.error_type == "ExpiredCodeException":
                return AuthResult.failure(
                    "OTP code has expired. Please request a new one.",
                )
            return AuthResult.failure(
                f"OTP verification failed: {e.message}",
            )
        except Exception as e:
            logger.exception(
                "otp_verification_error",
                error=str(e),
                correlation_id=correlation_id,
            )
            return AuthResult.failure(
                f"OTP verification error: {e}",
            )

        # Check for additional challenges
        if "ChallengeName" in response:
            challenge = response["ChallengeName"]
            logger.warning(
                "otp_additional_challenge",
                challenge=challenge,
                correlation_id=correlation_id,
            )
            return AuthResult.failure(
                f"Additional auth challenge not supported in CLI: {challenge}"
            )

        # Extract tokens from successful verification
        auth_result_data = response.get(
            "AuthenticationResult",
            {},
        )
        return self._process_auth_tokens(
            auth_result_data,
            email,
            correlation_id,
        )

    def _process_auth_tokens(
        self,
        auth_result_data: dict[str, Any],
        email: str,
        correlation_id: str,
    ) -> AuthResult:
        """Process authentication tokens from a Cognito response.

        Shared by ``login``, ``otp_login``, and ``verify_otp`` to
        avoid duplicated token extraction and storage logic.

        Args:
            auth_result_data: The ``AuthenticationResult`` dict from
                Cognito.
            email: User's email address.
            correlation_id: Request correlation ID for logging.

        Returns:
            AuthResult with tokens and user info.
        """
        access_token = auth_result_data.get("AccessToken")
        id_token = auth_result_data.get("IdToken")
        refresh_token = auth_result_data.get("RefreshToken")
        expires_in = auth_result_data.get("ExpiresIn", _DEFAULT_TOKEN_EXPIRY_SECONDS)

        if not access_token or not id_token:
            return AuthResult.failure(
                "Missing tokens in auth response",
            )

        expires_at = datetime.now(UTC) + timedelta(seconds=expires_in)

        tokens = TokenInfo(
            access_token=access_token,
            id_token=id_token,
            refresh_token=refresh_token or "",
            expires_at=expires_at,
            token_type=auth_result_data.get("TokenType", "Bearer"),
        )

        # Extract user info from ID token claims
        user_email = email
        user_id = "unknown"
        try:
            claims = jwt.get_unverified_claims(id_token)
            user_id = claims.get("sub", "unknown")
            user_email = claims.get("email", email)
        except Exception as e:
            logger.warning("id_token_decode_failed", error=str(e))

        # Store credentials securely in OS keychain
        if self._keyring_available:
            try:
                self._storage.store_tokens(tokens)
                self._storage.store_user_info(user_id, user_email)
            except Exception as e:
                logger.warning(
                    "credential_storage_failed_keyring_unavailable",
                    error=str(e),
                    correlation_id=correlation_id,
                )
                self._keyring_available = False

        logger.info(
            "login_success",
            user_email=user_email,
            correlation_id=correlation_id,
        )

        return AuthResult.success_result(
            user_email=user_email,
            user_id=user_id,
            expires_at=expires_at,
            tokens=tokens,
        )

    async def __aenter__(self) -> CognitoAuth:
        """Async context manager entry."""
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Async context manager exit with cleanup."""
        await self.close()
