"""Main API module for Kwikset Halo.

This module contains the API class which provides the main interface
for interacting with Kwikset smart lock devices via the cloud API.
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from functools import partial
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

import aiohttp
import attr
import boto3
import botocore
import pycognito
from aiohttp import ClientSession
from aiohttp.client_exceptions import (
    ClientConnectorError,
    ClientPayloadError,
    ClientResponseError,
    ServerConnectionError,
)
from botocore.exceptions import BotoCoreError, ClientError
from pycognito.aws_srp import AWSSRP
from pycognito.exceptions import (
    SMSMFAChallengeException,
    SoftwareTokenMFAChallengeException,
)

from .const import (
    CLIENT_ID,
    CUSTOM_CHALLENGE_GENERATE_ANSWER,
    DEFAULT_HEADER_ACCEPT_ENCODING,
    DEFAULT_HEADER_USER_AGENT,
    DEFAULT_TIMEOUT,
    LOGGER,
    LOGIN_HEADER_USER_AGENT,
    MAX_RETRIES,
    POOL_ID,
    POOL_REGION,
    RETRY_DELAY,
)
from .device import Device
from .errors import (
    ConnectionError as KwiksetConnectionError,
)
from .errors import (
    MFAChallengeRequired,
    PasswordChangeRequired,
    RequestError,
    TokenExpiredError,
    Unauthenticated,
    UnknownError,
    UserExists,
    UserNotConfirmed,
    UserNotFound,
)
from .home_user import HomeUser
from .subscription import SubscriptionManager
from .user import User

if TYPE_CHECKING:
    from types import TracebackType

# Type alias for token update callback
TokenUpdateCallback = Callable[[str, str, str], Awaitable[None]]

# Mapping of AWS Cognito error codes to custom exception classes
AWS_EXCEPTIONS: dict[str, type[Exception]] = {
    "UserNotFoundException": UserNotFound,
    "UserNotConfirmedException": UserNotConfirmed,
    "UsernameExistsException": UserExists,
    "NotAuthorizedException": Unauthenticated,
    "PasswordResetRequiredException": PasswordChangeRequired,
    "ExpiredTokenException": TokenExpiredError,
    "InvalidParameterException": Unauthenticated,
}


def _map_aws_exception(err: ClientError) -> Exception:
    """Map AWS exception to custom exception.

    :param err: The AWS ClientError exception
    :type err: ``ClientError``
    :rtype: ``Exception``
    """
    error_code = err.response.get("Error", {}).get("Code", "")
    error_message = err.response.get("Error", {}).get("Message", str(err))
    exception_class = AWS_EXCEPTIONS.get(error_code, UnknownError)
    return exception_class(error_message)


@attr.s(repr=False)
class API:
    """Main API class for interacting with Kwikset cloud services.

    This class handles authentication via AWS Cognito and provides
    access to device and user endpoints.

    :param websession: Optional aiohttp ClientSession for HTTP requests.
        If provided (e.g., from Home Assistant's async_get_clientsession),
        the session will be reused and NOT closed by this client.
        If not provided, a new session will be created for each request.
    :type websession: ``Optional[ClientSession]``
    :param user_pool_region: AWS Cognito user pool region
    :type user_pool_region: ``str``

    Example usage with Home Assistant:
        >>> from homeassistant.helpers.aiohttp_client import async_get_clientsession
        >>> session = async_get_clientsession(hass)
        >>> api = API(websession=session)

    Example usage standalone:
        >>> api = API()
    """

    # Optional external websession (e.g., from Home Assistant)
    websession: ClientSession | None = attr.ib(default=None)

    # Configurable attributes via attrs
    user_pool_region: str = attr.ib(default=POOL_REGION)
    username: str | None = attr.ib(default=None)
    timeout: int = attr.ib(default=DEFAULT_TIMEOUT)

    # Authentication tokens - these follow industry standards for OAuth/Cognito clients
    # by being exposed as public attributes for serialization and persistence
    id_token: str | None = attr.ib(default=None, repr=False)
    access_token: str | None = attr.ib(default=None, repr=False)
    refresh_token: str | None = attr.ib(default=None, repr=False)

    # Optional callback for token updates (e.g., for Home Assistant config entry updates)
    # Signature: async def callback(id_token: str, access_token: str, refresh_token: str) -> None
    token_update_callback: TokenUpdateCallback | None = attr.ib(default=None, repr=False)

    # Class-level constants
    user_pool_id: str = POOL_ID
    client_id: str = CLIENT_ID
    pool_region: str = POOL_REGION

    # Internal state - using attr.ib for proper instance initialization
    _loop: asyncio.AbstractEventLoop | None = attr.ib(default=None, init=False, repr=False)
    _boto_session: boto3.Session | None = attr.ib(default=None, init=False, repr=False)
    _request_lock: asyncio.Lock = attr.ib(factory=asyncio.Lock, init=False, repr=False)
    _mfa_cognito: pycognito.Cognito | None = attr.ib(default=None, init=False, repr=False)

    # Endpoint handlers
    device: Device | None = attr.ib(default=None, init=False, repr=False)
    user: User | None = attr.ib(default=None, init=False, repr=False)
    home_user: HomeUser | None = attr.ib(default=None, init=False, repr=False)

    # Subscription manager
    _subscription_manager: SubscriptionManager | None = attr.ib(
        default=None, init=False, repr=False
    )
    # Session created for subscription manager when no external websession is provided
    _subscription_session: ClientSession | None = attr.ib(default=None, init=False, repr=False)

    def __repr__(self) -> str:
        """Return a string representation of the API instance.

        This method intentionally omits sensitive token information
        to prevent accidental exposure in logs or debug output.

        :rtype: ``str``
        """
        auth_status = "authenticated" if self.is_authenticated else "unauthenticated"
        return (
            f"<{self.__class__.__name__}("
            f"username={self.username!r}, "
            f"status={auth_status}, "
            f"region={self.user_pool_region!r})>"
        )

    @property
    def is_authenticated(self) -> bool:
        """Check if the API client is currently authenticated.

        This property provides a convenient way to check authentication
        status without directly accessing token attributes.

        :rtype: ``bool``

        Example:
            >>> if api.is_authenticated:
            ...     homes = await api.user.get_homes()
        """
        return (
            self.id_token is not None
            and self.access_token is not None
            and self.refresh_token is not None
        )

    @property
    def subscriptions(self) -> SubscriptionManager:
        """Return the subscription manager, creating it lazily.

        The subscription manager provides real-time event subscriptions
        via AWS AppSync WebSocket connections.

        :rtype: ``SubscriptionManager``
        :raises Unauthenticated: If the API client is not authenticated

        Example:
            >>> api.subscriptions.set_callback(my_handler)
            >>> sub_id = await api.subscriptions.async_subscribe_device(email)
        """
        if not self.is_authenticated:
            raise Unauthenticated("Cannot access subscriptions without authentication.")

        if self._subscription_manager is None:
            # Use injected websession or create a tracked one for subscriptions
            session = self.websession
            if session is None:
                from aiohttp import ClientSession as _CS  # noqa: N811

                session = _CS()
                self._subscription_session = session

            self._subscription_manager = SubscriptionManager(
                credentials_provider=self._get_id_token,
                session=session,
            )

        return self._subscription_manager

    async def _get_id_token(self) -> str:
        """Return the current id_token, refreshing if needed.

        Used as the credentials provider for ``SubscriptionManager``.

        :rtype: ``str``
        :raises Unauthenticated: If not authenticated
        """
        await self.async_check_token()
        if self.id_token is None:
            raise Unauthenticated("No id_token available")
        return self.id_token

    async def __aenter__(self) -> API:
        """Enter the async context manager.

        :rtype: ``API``

        Example:
            >>> async with API() as api:
            ...     await api.async_login('user@example.com', 'password')
            ...     homes = await api.user.get_homes()
        """
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Exit the async context manager and clean up resources.

        This method clears sensitive token data and cleans up
        internal state when the context manager exits.

        :param exc_type: Exception type if an exception occurred
        :param exc_val: Exception value if an exception occurred
        :param exc_tb: Exception traceback if an exception occurred
        """
        await self.async_close()

    async def async_close(self) -> None:
        """Close the API client and clean up resources.

        This method clears authentication tokens and internal state.
        Call this when you're done using the API client to ensure
        proper cleanup of sensitive data.

        Note: If a websession was provided externally (e.g., from
        Home Assistant), it will NOT be closed by this method.
        """
        # Clear tokens
        self.id_token = None
        self.access_token = None
        self.refresh_token = None

        # Close subscription manager
        if self._subscription_manager is not None:
            await self._subscription_manager.async_close()
            self._subscription_manager = None

        # Close subscription session if we created one
        if self._subscription_session is not None:
            await self._subscription_session.close()
            self._subscription_session = None

        # Clear internal state
        self._mfa_cognito = None
        self._boto_session = None
        self._loop = None

        # Clear endpoint handlers
        self.device = None
        self.user = None
        self.home_user = None

    async def _request(
        self,
        method: str,
        url: str,
        retry_count: int = 0,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Make an authenticated HTTP request against the API.

        This method automatically handles token validation and refresh,
        sets required headers, provides timeout support, and includes
        retry logic for transient failures.

        If a websession was provided during initialization, it will be reused.
        Otherwise, a new session is created and closed after the request.

        :param method: HTTP method (get, post, patch, etc.)
        :type method: ``str``
        :param url: The URL to request
        :type url: ``str``
        :param retry_count: Current retry attempt (internal use)
        :type retry_count: ``int``
        :param kwargs: Additional arguments passed to the request
        :rtype: ``Dict[str, Any]``
        :raises RequestError: If the request fails after all retries
        :raises KwiksetConnectionError: If unable to connect to the API
        """
        await self.async_check_token()

        kwargs.setdefault("headers", {})
        kwargs["headers"].update(
            {
                "Host": urlparse(url).netloc,
                "User-Agent": DEFAULT_HEADER_USER_AGENT,
                "Accept-Encoding": DEFAULT_HEADER_ACCEPT_ENCODING,
            }
        )

        if self.id_token:
            kwargs["headers"]["Authorization"] = f"Bearer {self.id_token}"

        # Set timeout if not already specified
        if "timeout" not in kwargs:
            kwargs["timeout"] = aiohttp.ClientTimeout(total=self.timeout)

        # Use provided websession or create a new one
        close_session = self.websession is None
        session = self.websession if self.websession is not None else ClientSession()

        try:
            async with session.request(method, url, **kwargs) as resp:
                data: dict[str, Any] = await resp.json(content_type=None)
                resp.raise_for_status()
                return data

        except asyncio.TimeoutError as err:
            if retry_count < MAX_RETRIES:
                LOGGER.debug(
                    "Request to %s timed out, retrying (%d/%d)",
                    url,
                    retry_count + 1,
                    MAX_RETRIES,
                )
                await asyncio.sleep(RETRY_DELAY * (retry_count + 1))
                return await self._request(method, url, retry_count=retry_count + 1, **kwargs)
            raise KwiksetConnectionError(
                f"Request to {url} timed out after {MAX_RETRIES} retries"
            ) from err
        except ClientConnectorError as err:
            if retry_count < MAX_RETRIES:
                LOGGER.debug(
                    "Connection to %s failed, retrying (%d/%d)",
                    url,
                    retry_count + 1,
                    MAX_RETRIES,
                )
                await asyncio.sleep(RETRY_DELAY * (retry_count + 1))
                return await self._request(method, url, retry_count=retry_count + 1, **kwargs)
            raise KwiksetConnectionError(
                f"Unable to connect to {url} after {MAX_RETRIES} retries: {err}"
            ) from err
        except ServerConnectionError as err:
            if retry_count < MAX_RETRIES:
                LOGGER.debug(
                    "Server connection to %s failed, retrying (%d/%d)",
                    url,
                    retry_count + 1,
                    MAX_RETRIES,
                )
                await asyncio.sleep(RETRY_DELAY * (retry_count + 1))
                return await self._request(method, url, retry_count=retry_count + 1, **kwargs)
            raise KwiksetConnectionError(
                f"Server connection to {url} failed after {MAX_RETRIES} retries: {err}"
            ) from err
        except ClientResponseError as err:
            # Handle 401/403 specifically for Home Assistant reauth flows
            if err.status in (401, 403):
                raise Unauthenticated(f"Authentication failed for {url}: {err.message}") from err
            raise RequestError(f"Response error while requesting {url}: {err}") from err
        except ClientPayloadError as err:
            raise RequestError(f"Payload error while requesting {url}: {err}") from err
        except Exception as err:
            raise RequestError(f"Unexpected error while requesting {url}: {err}") from err
        finally:
            # Only close the session if we created it ourselves
            if close_session:
                await session.close()

    async def _update_token(
        self,
        id_token: str,
        access_token: str,
        refresh_token: str | None = None,
    ) -> None:
        """Update authentication tokens and initialize endpoint handlers.

        This method updates the internal token storage and optionally
        notifies a callback when tokens are refreshed. This is useful
        for integrations that need to persist tokens (e.g., Home Assistant).

        :param id_token: The JWT ID token
        :type id_token: ``str``
        :param access_token: The JWT access token
        :type access_token: ``str``
        :param refresh_token: Optional refresh token
        :type refresh_token: ``Optional[str]``
        """
        self.id_token = id_token
        self.access_token = access_token
        if refresh_token is not None:
            self.refresh_token = refresh_token

        # Notify callback if tokens were updated (for persistence)
        if self.token_update_callback is not None and self.refresh_token is not None:
            try:
                await self.token_update_callback(
                    self.id_token,
                    self.access_token,
                    self.refresh_token,
                )
            except Exception as err:  # pylint: disable=broad-except
                LOGGER.warning("Token update callback failed: %s", err)

        # Initialize endpoint handlers if not already done
        if self.device is None:
            self.device = Device(self._request)

        if self.user is None:
            self.user = User(self._request)

        if self.home_user is None:
            self.home_user = HomeUser(self._request)

    async def async_authenticate_with_tokens(
        self,
        id_token: str,
        access_token: str,
        refresh_token: str,
        username: str | None = None,
    ) -> None:
        """Restore authentication from previously saved tokens.

        This method allows Home Assistant and other integrations to restore
        a session from tokens stored in config entries without requiring
        the user to re-enter their password.

        :param id_token: The JWT ID token from a previous session
        :type id_token: ``str``
        :param access_token: The JWT access token from a previous session
        :type access_token: ``str``
        :param refresh_token: The refresh token from a previous session
        :type refresh_token: ``str``
        :param username: Optional username (email) for the account
        :type username: ``Optional[str]``
        :raises Unauthenticated: If the tokens are invalid or expired
        :raises TokenExpiredError: If token refresh fails

        Example with Home Assistant:
            >>> # Restore from config entry
            >>> api = API(websession=async_get_clientsession(hass))
            >>> await api.async_authenticate_with_tokens(
            ...     id_token=entry.data["id_token"],
            ...     access_token=entry.data["access_token"],
            ...     refresh_token=entry.data["refresh_token"],
            ... )
        """
        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        if username:
            self.username = username

        # Set tokens directly
        await self._update_token(id_token, access_token, refresh_token)

        # Validate tokens by attempting a refresh
        # This ensures the tokens are still valid
        try:
            await self.async_renew_access_token()
        except (Unauthenticated, UserNotFound) as err:
            # Clear tokens on failure
            self.id_token = None
            self.access_token = None
            self.refresh_token = None
            raise TokenExpiredError(
                "Stored tokens are no longer valid. Re-authentication required."
            ) from err

    async def async_login(
        self,
        email: str,
        password: str | None = None,
    ) -> None:
        """Authenticate with the Kwikset cloud using email and password.

        If MFA is enabled on the account, this method will raise
        MFAChallengeRequired exception with the necessary tokens to
        complete authentication via async_respond_to_mfa_challenge.

        For Kwikset's CUSTOM_CHALLENGE flow, the exception will include
        mfa_type="CUSTOM_CHALLENGE". The caller should then:
        1. Call async_request_custom_challenge_code() to trigger code delivery
        2. Call async_respond_to_mfa_challenge() with the verification code

        :param email: User's email address
        :type email: ``str``
        :param password: User's password
        :type password: ``Optional[str]``
        :raises MFAChallengeRequired: If MFA verification is required
        :raises Unauthenticated: If credentials are invalid
        :raises UnknownError: If an unknown error occurs
        """
        self._loop = asyncio.get_event_loop()
        self.username = email

        # Create cognito client outside try block so it's available in except handlers
        cognito = await self._loop.run_in_executor(
            None,
            partial(self._create_cognito_client, username=email),
        )

        try:
            # Use AWSSRP directly to capture raw challenge responses.
            # This bypasses cognito.authenticate() which can't handle
            # CUSTOM_CHALLENGE (would raise KeyError in _set_tokens).
            tokens = await self._loop.run_in_executor(
                None,
                partial(self._authenticate_srp, cognito, password),
            )

            # Check for CUSTOM_CHALLENGE (not handled by pycognito)
            challenge_name = tokens.get("ChallengeName") if isinstance(tokens, dict) else None
            if challenge_name == "CUSTOM_CHALLENGE":
                self._mfa_cognito = cognito
                raise MFAChallengeRequired(
                    message="Custom authentication challenge required. "
                    "Call async_request_custom_challenge_code() to request "
                    "a verification code.",
                    mfa_type="CUSTOM_CHALLENGE",
                    mfa_tokens=tokens,
                )

            # Success path — tokens were set on cognito in _authenticate_srp
            await self._update_token(
                cognito.id_token,
                cognito.access_token,
                cognito.refresh_token,
            )

        except SoftwareTokenMFAChallengeException as err:
            self._mfa_cognito = cognito
            raise MFAChallengeRequired(
                message="Software Token MFA challenge required. Please provide the MFA code.",
                mfa_type="SOFTWARE_TOKEN_MFA",
                mfa_tokens=err.get_tokens(),
            ) from err

        except SMSMFAChallengeException as err:
            self._mfa_cognito = cognito
            raise MFAChallengeRequired(
                message="SMS MFA challenge required. Please provide the MFA code.",
                mfa_type="SMS_MFA",
                mfa_tokens=err.get_tokens(),
            ) from err

        except ClientError as err:
            raise _map_aws_exception(err) from err

        except BotoCoreError as err:
            raise UnknownError(str(err)) from err

    async def async_respond_to_mfa_challenge(
        self,
        mfa_code: str,
        mfa_type: str | None = None,
        mfa_tokens: dict[str, Any] | None = None,
    ) -> None:
        """Respond to an MFA challenge with the user's MFA code.

        For CUSTOM_CHALLENGE, the mfa_tokens must contain a "Session" key,
        typically from async_request_custom_challenge_code() response.

        :param mfa_code: The MFA code from the user's authenticator app, SMS,
            or email verification
        :type mfa_code: ``str``
        :param mfa_type: The type of MFA challenge
            ('SMS_MFA', 'SOFTWARE_TOKEN_MFA', or 'CUSTOM_CHALLENGE')
        :type mfa_type: ``Optional[str]``
        :param mfa_tokens: Optional tokens from the MFA exception or
            from async_request_custom_challenge_code()
        :type mfa_tokens: ``Optional[Dict[str, Any]]``
        :raises Unauthenticated: If no MFA challenge is in progress
        :raises UnknownError: If MFA type is unknown or verification fails
        """
        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        try:
            # Determine which cognito instance to use
            if mfa_tokens and self.username:
                cognito = await self._loop.run_in_executor(
                    None,
                    partial(self._create_cognito_client, username=self.username),
                )
            elif self._mfa_cognito is not None:
                cognito = self._mfa_cognito
            else:
                raise Unauthenticated("No MFA challenge in progress. Please login first.")

            # Respond to the appropriate MFA challenge type
            if mfa_type == "SOFTWARE_TOKEN_MFA" or (
                mfa_type is None and self._mfa_cognito is not None
            ):
                await self._loop.run_in_executor(
                    None,
                    partial(
                        cognito.respond_to_software_token_mfa_challenge,
                        mfa_code,
                        mfa_tokens,
                    ),
                )
            elif mfa_type == "SMS_MFA":
                await self._loop.run_in_executor(
                    None,
                    partial(
                        cognito.respond_to_sms_mfa_challenge,
                        mfa_code,
                        mfa_tokens,
                    ),
                )

            elif mfa_type == "CUSTOM_CHALLENGE":
                # Handle Kwikset's CUSTOM_CHALLENGE verification
                if mfa_tokens is None or "Session" not in mfa_tokens:
                    raise Unauthenticated(
                        "CUSTOM_CHALLENGE requires mfa_tokens with a Session key. "
                        "Call async_request_custom_challenge_code() first."
                    )

                # Use boto3 client directly to respond to CUSTOM_CHALLENGE
                response = await self._loop.run_in_executor(
                    None,
                    partial(
                        cognito.client.respond_to_auth_challenge,
                        ClientId=self.client_id,
                        ChallengeName="CUSTOM_CHALLENGE",
                        Session=mfa_tokens["Session"],
                        ChallengeResponses={
                            "ANSWER": mfa_code,
                            "USERNAME": self.username,
                        },
                    ),
                )

                # Extract tokens from the response
                auth_result = response.get("AuthenticationResult", {})
                if not auth_result:
                    raise UnknownError(
                        "CUSTOM_CHALLENGE verification did not return authentication tokens."
                    )

                # Manually set tokens since we bypassed pycognito's token handling
                id_token = auth_result.get("IdToken", "")
                access_token = auth_result.get("AccessToken", "")
                refresh_token = auth_result.get("RefreshToken", "")

                await self._update_token(id_token, access_token, refresh_token)
                self._mfa_cognito = None
                return

            else:
                raise UnknownError("Unknown MFA type")

            # Update tokens after successful MFA verification
            await self._update_token(
                cognito.id_token,
                cognito.access_token,
                cognito.refresh_token,
            )

            # Clean up stored cognito instance
            self._mfa_cognito = None

        except ClientError as err:
            raise _map_aws_exception(err) from err

        except BotoCoreError as err:
            raise UnknownError(str(err)) from err

    async def async_request_custom_challenge_code(
        self,
        code_type: str = "email",
        mfa_tokens: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Request a verification code for a CUSTOM_CHALLENGE.

        This is the first step after async_login() raises MFAChallengeRequired
        with mfa_type="CUSTOM_CHALLENGE". It triggers the delivery of a
        verification code via the specified medium (email or phone).

        After calling this method, call async_respond_to_mfa_challenge()
        with the received code to complete authentication.

        :param code_type: The delivery medium for the verification code.
            Must be "email" or "phone".
        :type code_type: ``str``
        :param mfa_tokens: The challenge tokens from the MFAChallengeRequired
            exception. Must contain a "Session" key.
        :type mfa_tokens: ``Optional[Dict[str, Any]]``
        :rtype: ``Dict[str, Any]``
        :returns: Updated challenge tokens (with new Session) for use
            with async_respond_to_mfa_challenge()
        :raises Unauthenticated: If no challenge is in progress
        :raises UnknownError: If the code request fails
        :raises ValueError: If code_type is not "email" or "phone"
        """
        if code_type not in ("email", "phone"):
            raise ValueError(f"code_type must be 'email' or 'phone', got {code_type!r}")

        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        if mfa_tokens is None or "Session" not in mfa_tokens:
            raise Unauthenticated("No CUSTOM_CHALLENGE in progress. Missing session token.")

        # Get the cognito client for boto3 access
        if self._mfa_cognito is not None:
            cognito = self._mfa_cognito
        elif self.username:
            cognito = await self._loop.run_in_executor(
                None,
                partial(self._create_cognito_client, username=self.username),
            )
        else:
            raise Unauthenticated("No challenge in progress. Please login first.")

        answer = CUSTOM_CHALLENGE_GENERATE_ANSWER % code_type

        try:
            response: dict[str, Any] = await self._loop.run_in_executor(
                None,
                partial(
                    cognito.client.respond_to_auth_challenge,
                    ClientId=self.client_id,
                    ChallengeName="CUSTOM_CHALLENGE",
                    Session=mfa_tokens["Session"],
                    ChallengeResponses={
                        "ANSWER": answer,
                        "USERNAME": self.username,
                    },
                ),
            )

            LOGGER.debug("Custom challenge code requested via %s", code_type)
            return response

        except ClientError as err:
            raise _map_aws_exception(err) from err

        except BotoCoreError as err:
            raise UnknownError(f"Failed to request custom challenge code: {err}") from err

    async def async_check_token(self) -> None:
        """Check token validity and renew if necessary.

        This method is called automatically before each API request
        to ensure the authentication token is still valid.

        :raises Unauthenticated: If token refresh fails
        """
        async with self._request_lock:
            cognito = await self._async_authenticated_cognito()
            if not cognito.check_token(renew=False):
                return

            try:
                await self.async_renew_access_token()
            except (Unauthenticated, UserNotFound) as err:
                LOGGER.error("Unable to refresh token: %s", err)
                raise

    async def async_renew_access_token(
        self,
        access_token: str | None = None,
        refresh_token: str | None = None,
    ) -> None:
        """Renew the access token using the refresh token.

        :param access_token: Optional new access token to set
        :type access_token: ``Optional[str]``
        :param refresh_token: Optional new refresh token to set
        :type refresh_token: ``Optional[str]``
        :raises Unauthenticated: If token renewal fails
        :raises UnknownError: If an unknown error occurs
        """
        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        if access_token is not None:
            self.access_token = access_token
        if refresh_token is not None:
            self.refresh_token = refresh_token

        cognito = await self._async_authenticated_cognito()

        try:
            await self._loop.run_in_executor(None, cognito.renew_access_token)
            # Pass refresh_token in case Cognito rotated it during refresh
            await self._update_token(
                cognito.id_token,
                cognito.access_token,
                cognito.refresh_token,
            )

        except ClientError as err:
            raise _map_aws_exception(err) from err

        except BotoCoreError as err:
            raise UnknownError(str(err)) from err

    async def _async_authenticated_cognito(self) -> pycognito.Cognito:
        """Return an authenticated Cognito instance.

        Creates a pycognito Cognito client with the current authentication
        tokens for operations like token refresh.

        :rtype: ``pycognito.Cognito``
        :raises Unauthenticated: If no authentication tokens are available
        """
        if self.access_token is None or self.refresh_token is None:
            raise Unauthenticated("No authentication found")

        if self._loop is None:
            self._loop = asyncio.get_event_loop()

        # Include id_token if available for complete token state
        kwargs: dict[str, Any] = {
            "access_token": self.access_token,
            "refresh_token": self.refresh_token,
        }
        if self.id_token is not None:
            kwargs["id_token"] = self.id_token

        return await self._loop.run_in_executor(
            None,
            partial(self._create_cognito_client, **kwargs),
        )

    def _create_cognito_client(self, **kwargs: Any) -> pycognito.Cognito:
        """Create a new Cognito client.

        NOTE: This method performs I/O operations.

        :param kwargs: Additional arguments passed to the Cognito client
        :rtype: ``pycognito.Cognito``
        """
        if self._boto_session is None:
            self._boto_session = boto3.session.Session()

        config = botocore.config.Config(
            signature_version=botocore.UNSIGNED,
            user_agent_extra=LOGIN_HEADER_USER_AGENT,
        )

        return pycognito.Cognito(
            user_pool_id=self.user_pool_id,
            client_id=self.client_id,
            user_pool_region=self.pool_region,
            botocore_config=config,
            session=self._boto_session,
            **kwargs,
        )

    def _authenticate_srp(
        self,
        cognito: pycognito.Cognito,
        password: str | None,
    ) -> dict[str, Any]:
        """Perform SRP authentication using pycognito's AWSSRP directly.

        This method bypasses Cognito.authenticate() to gain access to the raw
        response from AWS Cognito. This is necessary because Kwikset's Cognito
        setup uses CUSTOM_CHALLENGE, which pycognito doesn't natively handle.

        When pycognito encounters CUSTOM_CHALLENGE, its authenticate_user()
        returns the raw response dict (without raising an exception), but
        then _set_tokens() fails because there's no AuthenticationResult.
        By using AWSSRP directly, we can intercept the CUSTOM_CHALLENGE
        response before _set_tokens is called.

        NOTE: This method performs I/O operations (SRP auth with Cognito).
        Must be called within run_in_executor.

        :param cognito: The pycognito Cognito instance (provides the boto3 client)
        :type cognito: ``pycognito.Cognito``
        :param password: The user's password
        :type password: ``Optional[str]``
        :rtype: ``Dict[str, Any]``
        :raises SoftwareTokenMFAChallengeException: If software token MFA is required
        :raises SMSMFAChallengeException: If SMS MFA is required
        :raises ClientError: If an AWS error occurs
        """
        aws_srp = AWSSRP(
            username=cognito.username,
            password=password,
            pool_id=self.user_pool_id,
            client_id=self.client_id,
            client=cognito.client,
        )

        # authenticate_user() returns raw boto3 response:
        # - Success: {"AuthenticationResult": {"IdToken": ..., ...}}
        # - CUSTOM_CHALLENGE: {"ChallengeName": "CUSTOM_CHALLENGE", "Session": ..., ...}
        # - SMS_MFA: raises SMSMFAChallengeException
        # - SOFTWARE_TOKEN_MFA: raises SoftwareTokenMFAChallengeException
        # - NEW_PASSWORD_REQUIRED: raises ForceChangePasswordException
        tokens: dict[str, Any] = aws_srp.authenticate_user()

        # If auth succeeded (no challenge), set tokens on the cognito instance
        if "AuthenticationResult" in tokens:
            cognito._set_tokens(tokens)

        return tokens
