"""Tests for the API class.

This module contains comprehensive tests for the main API class,
including initialization, authentication, token management,
context manager behavior, and HTTP request handling.
"""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from aiohttp import ClientSession
from aiohttp.client_exceptions import (
    ClientConnectorError,
    ClientPayloadError,
    ClientResponseError,
    ServerConnectionError,
)
from botocore.exceptions import BotoCoreError
from pycognito.exceptions import (
    SMSMFAChallengeException,
    SoftwareTokenMFAChallengeException,
)

from aiokwikset import API
from aiokwikset.errors import (
    ConnectionError,
    MFAChallengeRequired,
    PasswordChangeRequired,
    RequestError,
    Unauthenticated,
    UnknownError,
    UserExists,
    UserNotConfirmed,
    UserNotFound,
)

from .conftest import (
    MOCK_CUSTOM_CHALLENGE_AUTH_RESULT,
    MOCK_CUSTOM_CHALLENGE_CODE_RESPONSE,
    MOCK_CUSTOM_CHALLENGE_RESPONSE,
    MockCognito,
    create_client_error,
    create_mock_cognito_with_client,
)

# ============================================================================
# API Initialization Tests
# ============================================================================


class TestAPIInitialization:
    """Tests for API class initialization."""

    def test_init_default(self) -> None:
        """Test API initialization with default parameters."""
        api = API()

        assert api.websession is None
        assert api.username is None
        assert api.id_token is None
        assert api.access_token is None
        assert api.refresh_token is None
        assert api.user_pool_region == "us-east-1"
        assert api.device is None
        assert api.user is None

    def test_init_with_websession(self) -> None:
        """Test API initialization with external websession."""
        mock_session = MagicMock(spec=ClientSession)
        api = API(websession=mock_session)

        assert api.websession is mock_session

    def test_init_with_custom_region(self) -> None:
        """Test API initialization with custom pool region."""
        api = API(user_pool_region="us-west-2")

        assert api.user_pool_region == "us-west-2"

    def test_init_with_username(self) -> None:
        """Test API initialization with username."""
        api = API(username="test@example.com")

        assert api.username == "test@example.com"

    def test_init_with_tokens(self) -> None:
        """Test API initialization with pre-existing tokens."""
        api = API(
            id_token="id_token_value",
            access_token="access_token_value",
            refresh_token="refresh_token_value",
        )

        assert api.id_token == "id_token_value"
        assert api.access_token == "access_token_value"
        assert api.refresh_token == "refresh_token_value"

    def test_init_with_token_callback(self) -> None:
        """Test API initialization with token update callback."""

        async def callback(id_t: str, access_t: str, refresh_t: str) -> None:
            pass

        api = API(token_update_callback=callback)

        assert api.token_update_callback is callback


# ============================================================================
# API Properties Tests
# ============================================================================


class TestAPIProperties:
    """Tests for API class properties."""

    def test_is_authenticated_true(self) -> None:
        """Test is_authenticated returns True when all tokens present."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        assert api.is_authenticated is True

    def test_is_authenticated_false_no_tokens(self) -> None:
        """Test is_authenticated returns False with no tokens."""
        api = API()

        assert api.is_authenticated is False

    def test_is_authenticated_false_partial_tokens(self) -> None:
        """Test is_authenticated returns False with partial tokens."""
        api = API(id_token="id_token")

        assert api.is_authenticated is False

        api2 = API(id_token="id_token", access_token="access_token")

        assert api2.is_authenticated is False

    def test_repr_hides_tokens(self) -> None:
        """Test __repr__ does not expose sensitive token data."""
        api = API(
            username="test@example.com",
            id_token="secret_id_token",
            access_token="secret_access_token",
            refresh_token="secret_refresh_token",
        )

        repr_str = repr(api)

        assert "secret" not in repr_str
        assert "token" not in repr_str.lower() or "status=" in repr_str
        assert "test@example.com" in repr_str
        assert "authenticated" in repr_str

    def test_repr_unauthenticated(self) -> None:
        """Test __repr__ shows unauthenticated status."""
        api = API(username="test@example.com")

        repr_str = repr(api)

        assert "unauthenticated" in repr_str


# ============================================================================
# Context Manager Tests
# ============================================================================


class TestAPIContextManager:
    """Tests for API async context manager behavior."""

    @pytest.mark.asyncio
    async def test_aenter_returns_api(self) -> None:
        """Test __aenter__ returns the API instance."""
        api = API()

        async with api as context_api:
            assert context_api is api

    @pytest.mark.asyncio
    async def test_aexit_cleans_up_resources(self) -> None:
        """Test __aexit__ cleans up tokens and state."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )
        api._mfa_cognito = MagicMock()

        async with api:
            pass

        assert api.id_token is None
        assert api.access_token is None
        assert api.refresh_token is None
        assert api._mfa_cognito is None

    @pytest.mark.asyncio
    async def test_context_manager_with_exception(self) -> None:
        """Test context manager cleans up even on exception."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        with pytest.raises(ValueError):
            async with api:
                raise ValueError("Test exception")

        assert api.id_token is None
        assert api.access_token is None


# ============================================================================
# async_close Tests
# ============================================================================


class TestAsyncClose:
    """Tests for async_close method."""

    @pytest.mark.asyncio
    async def test_async_close_clears_tokens(self) -> None:
        """Test async_close clears all authentication tokens."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        await api.async_close()

        assert api.id_token is None
        assert api.access_token is None
        assert api.refresh_token is None

    @pytest.mark.asyncio
    async def test_async_close_clears_internal_state(self) -> None:
        """Test async_close clears internal state."""
        api = API()
        api._mfa_cognito = MagicMock()
        api._boto_session = MagicMock()
        api._loop = asyncio.get_event_loop()

        await api.async_close()

        assert api._mfa_cognito is None
        assert api._boto_session is None
        assert api._loop is None

    @pytest.mark.asyncio
    async def test_async_close_clears_endpoint_handlers(self) -> None:
        """Test async_close clears endpoint handlers."""
        api = API()
        api.device = MagicMock()
        api.user = MagicMock()

        await api.async_close()

        assert api.device is None
        assert api.user is None


# ============================================================================
# Authentication Tests
# ============================================================================


class TestAsyncLogin:
    """Tests for async_login method."""

    @pytest.mark.asyncio
    async def test_login_success(self) -> None:
        """Test successful login without MFA."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MockCognito()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = {
                    "AuthenticationResult": {
                        "IdToken": "mock_id_token",
                        "AccessToken": "mock_access_token",
                        "RefreshToken": "mock_refresh_token",
                    }
                }

                await api.async_login("test@example.com", "password123")

        assert api.is_authenticated
        assert api.id_token == "mock_id_token"
        assert api.access_token == "mock_access_token"
        assert api.refresh_token == "mock_refresh_token"
        assert api.username == "test@example.com"

    @pytest.mark.asyncio
    async def test_login_initializes_endpoint_handlers(self) -> None:
        """Test successful login initializes device and user handlers."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MockCognito()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = {
                    "AuthenticationResult": {
                        "IdToken": "mock_id_token",
                        "AccessToken": "mock_access_token",
                        "RefreshToken": "mock_refresh_token",
                    }
                }

                await api.async_login("test@example.com", "password123")

        assert api.device is not None
        assert api.user is not None

    @pytest.mark.asyncio
    async def test_login_software_token_mfa(self) -> None:
        """Test login raises MFAChallengeRequired for software token MFA."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                exc = SoftwareTokenMFAChallengeException.__new__(SoftwareTokenMFAChallengeException)
                exc._tokens = {"session": "mfa_session_token"}
                mock_srp.side_effect = exc

                with pytest.raises(MFAChallengeRequired) as exc_info:
                    await api.async_login("test@example.com", "password123")

        assert exc_info.value.mfa_type == "SOFTWARE_TOKEN_MFA"

    @pytest.mark.asyncio
    async def test_login_sms_mfa(self) -> None:
        """Test login raises MFAChallengeRequired for SMS MFA."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                exc = SMSMFAChallengeException.__new__(SMSMFAChallengeException)
                exc._tokens = {"session": "sms_session_token"}
                mock_srp.side_effect = exc

                with pytest.raises(MFAChallengeRequired) as exc_info:
                    await api.async_login("test@example.com", "password123")

        assert exc_info.value.mfa_type == "SMS_MFA"

    @pytest.mark.asyncio
    async def test_login_user_not_found(self) -> None:
        """Test login raises UserNotFound for unknown user."""
        api = API()

        error = create_client_error("UserNotFoundException", "User not found")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=error),
                pytest.raises(UserNotFound),
            ):
                await api.async_login("unknown@example.com", "password123")

    @pytest.mark.asyncio
    async def test_login_user_not_confirmed(self) -> None:
        """Test login raises UserNotConfirmed for unconfirmed user."""
        api = API()

        error = create_client_error("UserNotConfirmedException", "User not confirmed")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=error),
                pytest.raises(UserNotConfirmed),
            ):
                await api.async_login("unconfirmed@example.com", "password123")

    @pytest.mark.asyncio
    async def test_login_invalid_credentials(self) -> None:
        """Test login raises Unauthenticated for invalid credentials."""
        api = API()

        error = create_client_error("NotAuthorizedException", "Incorrect password")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=error),
                pytest.raises(Unauthenticated),
            ):
                await api.async_login("test@example.com", "wrong_password")

    @pytest.mark.asyncio
    async def test_login_password_change_required(self) -> None:
        """Test login raises PasswordChangeRequired when needed."""
        api = API()

        error = create_client_error("PasswordResetRequiredException", "Password reset required")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=error),
                pytest.raises(PasswordChangeRequired),
            ):
                await api.async_login("test@example.com", "password123")

    @pytest.mark.asyncio
    async def test_login_user_exists(self) -> None:
        """Test login raises UserExists for existing user."""
        api = API()

        error = create_client_error("UsernameExistsException", "User already exists")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=error),
                pytest.raises(UserExists),
            ):
                await api.async_login("existing@example.com", "password123")

    @pytest.mark.asyncio
    async def test_login_unknown_botocore_error(self) -> None:
        """Test login raises UnknownError for generic BotoCoreError."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with (
                patch.object(api, "_authenticate_srp", side_effect=BotoCoreError()),
                pytest.raises(UnknownError),
            ):
                await api.async_login("test@example.com", "password123")


# ============================================================================
# MFA Response Tests
# ============================================================================


class TestAsyncRespondToMFAChallenge:
    """Tests for async_respond_to_mfa_challenge method."""

    @pytest.mark.asyncio
    async def test_mfa_response_software_token(self) -> None:
        """Test successful software token MFA response."""
        api = API(username="test@example.com")
        api._mfa_cognito = MockCognito()

        await api.async_respond_to_mfa_challenge(
            mfa_code="123456",
            mfa_type="SOFTWARE_TOKEN_MFA",
        )

        assert api.is_authenticated

    @pytest.mark.asyncio
    async def test_mfa_response_sms(self) -> None:
        """Test successful SMS MFA response."""
        api = API(username="test@example.com")
        api._mfa_cognito = MockCognito()

        await api.async_respond_to_mfa_challenge(
            mfa_code="123456",
            mfa_type="SMS_MFA",
        )

        assert api.is_authenticated

    @pytest.mark.asyncio
    async def test_mfa_response_with_tokens(self) -> None:
        """Test MFA response with mfa_tokens from exception."""
        api = API(username="test@example.com")

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MockCognito()
            mock_create.return_value = mock_cognito

            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="SOFTWARE_TOKEN_MFA",
                mfa_tokens={"session": "mfa_session"},
            )

        assert api.is_authenticated

    @pytest.mark.asyncio
    async def test_mfa_response_no_challenge(self) -> None:
        """Test MFA response raises Unauthenticated without active challenge."""
        api = API()

        with pytest.raises(Unauthenticated):
            await api.async_respond_to_mfa_challenge(mfa_code="123456")

    @pytest.mark.asyncio
    async def test_mfa_response_unknown_type(self) -> None:
        """Test MFA response raises UnknownError for invalid MFA type."""
        api = API(username="test@example.com")
        api._mfa_cognito = MockCognito()

        with pytest.raises(UnknownError):
            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="UNKNOWN_MFA",
            )

    @pytest.mark.asyncio
    async def test_mfa_response_clears_stored_cognito(self) -> None:
        """Test MFA response clears stored cognito instance after success."""
        api = API(username="test@example.com")
        api._mfa_cognito = MockCognito()

        await api.async_respond_to_mfa_challenge(
            mfa_code="123456",
            mfa_type="SOFTWARE_TOKEN_MFA",
        )

        assert api._mfa_cognito is None


# ============================================================================
# Token Management Tests
# ============================================================================


class TestTokenManagement:
    """Tests for token management methods."""

    @pytest.mark.asyncio
    async def test_renew_access_token_success(self) -> None:
        """Test successful token renewal."""
        api = API(
            id_token="old_id_token",
            access_token="old_access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "_async_authenticated_cognito") as mock_auth:
            mock_cognito = MockCognito()
            mock_auth.return_value = mock_cognito

            await api.async_renew_access_token()

        assert api.id_token == "renewed_id_token"
        assert api.access_token == "renewed_access_token"

    @pytest.mark.asyncio
    async def test_renew_access_token_with_new_tokens(self) -> None:
        """Test token renewal with explicit new tokens."""
        api = API(
            id_token="old_id_token",
            access_token="old_access_token",
            refresh_token="old_refresh_token",
        )

        async def mock_auth():
            return MockCognito()

        api._async_authenticated_cognito = mock_auth

        await api.async_renew_access_token(
            access_token="new_access_token",
            refresh_token="new_refresh_token",
        )

        # Tokens should be updated after renewal
        assert api.id_token == "renewed_id_token"
        assert api.access_token == "renewed_access_token"

    @pytest.mark.asyncio
    async def test_renew_access_token_invalid_refresh(self) -> None:
        """Test token renewal fails with invalid refresh token."""
        api = API(
            access_token="access_token",
            refresh_token="invalid_refresh_token",
        )

        error = create_client_error("NotAuthorizedException", "Invalid Refresh Token")

        with patch.object(api, "_async_authenticated_cognito") as mock_auth:
            mock_cognito = MockCognito(raise_on_renew=error)
            mock_auth.return_value = mock_cognito

            with pytest.raises(Unauthenticated):
                await api.async_renew_access_token()

    @pytest.mark.asyncio
    async def test_renew_access_token_network_error(self) -> None:
        """Test token renewal handles network errors."""
        api = API(
            access_token="access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "_async_authenticated_cognito") as mock_auth:
            mock_cognito = MockCognito(raise_on_renew=BotoCoreError())
            mock_auth.return_value = mock_cognito

            with pytest.raises(UnknownError):
                await api.async_renew_access_token()

    @pytest.mark.asyncio
    async def test_check_token_no_renewal_needed(self) -> None:
        """Test async_check_token when token is still valid."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "_async_authenticated_cognito") as mock_auth:
            mock_cognito = MockCognito(token_expired=False)
            mock_auth.return_value = mock_cognito

            # Should not raise and should not renew
            await api.async_check_token()

    @pytest.mark.asyncio
    async def test_check_token_renewal_needed(self) -> None:
        """Test async_check_token when token needs renewal."""
        api = API(
            id_token="old_id_token",
            access_token="old_access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "_async_authenticated_cognito") as mock_auth:
            mock_cognito = MockCognito(token_expired=True)
            mock_auth.return_value = mock_cognito

            with patch.object(api, "async_renew_access_token") as mock_renew:
                mock_renew.return_value = None

                await api.async_check_token()

                mock_renew.assert_called_once()

    @pytest.mark.asyncio
    async def test_token_update_callback_called(self) -> None:
        """Test token update callback is invoked after token update."""
        callback_called = False
        received_tokens: dict[str, str] = {}

        async def callback(id_t: str, access_t: str, refresh_t: str) -> None:
            nonlocal callback_called, received_tokens
            callback_called = True
            received_tokens = {
                "id_token": id_t,
                "access_token": access_t,
                "refresh_token": refresh_t,
            }

        api = API(token_update_callback=callback)

        await api._update_token("id_token", "access_token", "refresh_token")

        assert callback_called
        assert received_tokens["id_token"] == "id_token"
        assert received_tokens["access_token"] == "access_token"
        assert received_tokens["refresh_token"] == "refresh_token"

    @pytest.mark.asyncio
    async def test_token_update_callback_failure_logged(self) -> None:
        """Test token update callback failure is handled gracefully."""

        async def failing_callback(
            _id_t: str,
            _access_t: str,
            _refresh_t: str,
        ) -> None:
            raise Exception("Callback failed")

        api = API(token_update_callback=failing_callback)

        # Should not raise, just log warning
        await api._update_token("id_token", "access_token", "refresh_token")

        assert api.is_authenticated


# ============================================================================
# HTTP Request Tests
# ============================================================================


class TestRequest:
    """Tests for _request method."""

    @pytest.mark.asyncio
    async def test_request_adds_auth_header(self) -> None:
        """Test request adds Authorization header."""
        mock_response = AsyncMock()
        mock_response.json = AsyncMock(return_value={"data": []})
        mock_response.raise_for_status = MagicMock()

        mock_cm = AsyncMock()
        mock_cm.__aenter__.return_value = mock_response
        mock_cm.__aexit__.return_value = None

        mock_session = MagicMock()
        mock_session.request.return_value = mock_cm

        api = API(
            websession=mock_session,
            id_token="test_id_token",
            access_token="test_access_token",
            refresh_token="test_refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        await api._request("get", "https://api.example.com/test")

        # Verify request was called with correct headers
        call_args = mock_session.request.call_args
        headers = call_args.kwargs.get("headers", {})
        assert headers.get("Authorization") == "Bearer test_id_token"

    @pytest.mark.asyncio
    async def test_request_reuses_provided_websession(self) -> None:
        """Test request reuses provided external websession."""
        mock_session = AsyncMock(spec=ClientSession)
        mock_response = AsyncMock()
        mock_response.json = AsyncMock(return_value={"data": []})
        mock_response.raise_for_status = MagicMock()
        mock_session.request.return_value.__aenter__.return_value = mock_response

        api = API(
            websession=mock_session,
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "async_check_token"):
            await api._request("get", "https://api.example.com/test")

        mock_session.request.assert_called_once()
        # Session should not be closed when provided externally
        mock_session.close.assert_not_called()

    @pytest.mark.asyncio
    async def test_request_creates_and_closes_session(self) -> None:
        """Test request creates and closes session when not provided."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        with patch("aiokwikset.api.ClientSession") as mock_session_class:
            mock_session = MagicMock()
            mock_response = AsyncMock()
            mock_response.json = AsyncMock(return_value={"data": []})
            mock_response.raise_for_status = MagicMock()

            mock_cm = AsyncMock()
            mock_cm.__aenter__.return_value = mock_response
            mock_cm.__aexit__.return_value = None
            mock_session.request.return_value = mock_cm
            mock_session.close = AsyncMock()
            mock_session_class.return_value = mock_session

            await api._request("get", "https://api.example.com/test")

            mock_session.close.assert_called_once()

    @pytest.mark.asyncio
    async def test_request_handles_client_response_error(self) -> None:
        """Test request handles ClientResponseError."""
        mock_response = AsyncMock()
        mock_response.json = AsyncMock(return_value={})
        mock_response.raise_for_status = MagicMock(
            side_effect=ClientResponseError(
                request_info=MagicMock(),
                history=(),
                status=401,
            )
        )

        mock_cm = AsyncMock()
        mock_cm.__aenter__.return_value = mock_response
        mock_cm.__aexit__.return_value = None

        mock_session = MagicMock()
        mock_session.request.return_value = mock_cm

        api = API(
            websession=mock_session,
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        # 401 responses now raise Unauthenticated for better HA reauth flow support
        with pytest.raises(Unauthenticated) as exc_info:
            await api._request("get", "https://api.example.com/test")

        assert "Authentication failed" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_request_handles_client_connector_error(self) -> None:
        """Test request handles ClientConnectorError with retry logic."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        with patch("aiohttp.ClientSession") as mock_session_class:
            mock_session = MagicMock()
            mock_session.request.side_effect = ClientConnectorError(
                connection_key=MagicMock(),
                os_error=OSError("Connection refused"),
            )
            mock_session.close = AsyncMock()
            mock_session_class.return_value = mock_session

            # Now raises ConnectionError after retries for better HA handling
            with pytest.raises(ConnectionError) as exc_info:
                await api._request("get", "https://api.example.com/test")

            assert "Unable to connect" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_request_handles_server_connection_error(self) -> None:
        """Test request handles ServerConnectionError with retry logic."""
        mock_session = MagicMock()
        mock_session.request.side_effect = ServerConnectionError("Server disconnected")

        api = API(
            websession=mock_session,
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        # Now raises ConnectionError after retries for better HA handling
        with pytest.raises(ConnectionError) as exc_info:
            await api._request("get", "https://api.example.com/test")

        assert "Server connection" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_request_handles_payload_error(self) -> None:
        """Test request handles ClientPayloadError."""
        mock_session = MagicMock()
        mock_session.request.side_effect = ClientPayloadError("Invalid payload")

        api = API(
            websession=mock_session,
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        async def mock_check_token():
            pass

        api.async_check_token = mock_check_token  # type: ignore

        with pytest.raises(RequestError) as exc_info:
            await api._request("get", "https://api.example.com/test")

        assert "Payload error" in str(exc_info.value)


# ============================================================================
# Authenticated Cognito Tests
# ============================================================================


class TestAuthenticatedCognito:
    """Tests for _async_authenticated_cognito method."""

    @pytest.mark.asyncio
    async def test_authenticated_cognito_raises_without_tokens(self) -> None:
        """Test _async_authenticated_cognito raises when not authenticated."""
        api = API()

        with pytest.raises(Unauthenticated):
            await api._async_authenticated_cognito()

    @pytest.mark.asyncio
    async def test_authenticated_cognito_raises_partial_tokens(self) -> None:
        """Test _async_authenticated_cognito raises with partial tokens."""
        api = API(access_token="access_token")  # Missing refresh_token

        with pytest.raises(Unauthenticated):
            await api._async_authenticated_cognito()

    @pytest.mark.asyncio
    async def test_authenticated_cognito_returns_client(self) -> None:
        """Test _async_authenticated_cognito returns Cognito client."""
        api = API(
            id_token="id_token",
            access_token="access_token",
            refresh_token="refresh_token",
        )

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MockCognito()
            mock_create.return_value = mock_cognito

            result = await api._async_authenticated_cognito()

        assert result is mock_cognito


# ============================================================================
# CUSTOM_CHALLENGE Login Tests
# ============================================================================


class TestCustomChallengeLogin:
    """Tests for CUSTOM_CHALLENGE detection during login."""

    @pytest.mark.asyncio
    async def test_login_custom_challenge_detected(self) -> None:
        """Test login raises MFAChallengeRequired for CUSTOM_CHALLENGE."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_cognito.username = "test@example.com"
            mock_cognito.client = MagicMock()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = MOCK_CUSTOM_CHALLENGE_RESPONSE

                with pytest.raises(MFAChallengeRequired) as exc_info:
                    await api.async_login("test@example.com", "password123")

        assert exc_info.value.mfa_type == "CUSTOM_CHALLENGE"
        assert exc_info.value.mfa_tokens["Session"] == "mock_custom_challenge_session"
        assert api._mfa_cognito is mock_cognito

    @pytest.mark.asyncio
    async def test_login_custom_challenge_stores_mfa_tokens(self) -> None:
        """Test CUSTOM_CHALLENGE stores complete mfa_tokens including ChallengeParameters."""
        api = API()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = MOCK_CUSTOM_CHALLENGE_RESPONSE

                with pytest.raises(MFAChallengeRequired) as exc_info:
                    await api.async_login("test@example.com", "password123")

        assert "ChallengeParameters" in exc_info.value.mfa_tokens
        assert exc_info.value.mfa_tokens["ChallengeName"] == "CUSTOM_CHALLENGE"

    @pytest.mark.asyncio
    async def test_login_success_after_srp_refactor(self) -> None:
        """Test login still works for non-challenge responses after SRP refactor."""
        api = API()

        mock_cognito = MockCognito()

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = {
                    "AuthenticationResult": {
                        "IdToken": "mock_id_token",
                        "AccessToken": "mock_access_token",
                        "RefreshToken": "mock_refresh_token",
                    }
                }

                await api.async_login("test@example.com", "password123")

        assert api.is_authenticated


# ============================================================================
# Custom Challenge Code Request Tests
# ============================================================================


class TestAsyncRequestCustomChallengeCode:
    """Tests for async_request_custom_challenge_code method."""

    @pytest.mark.asyncio
    async def test_request_code_email(self) -> None:
        """Test requesting verification code via email."""
        api = API(username="test@example.com")

        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[MOCK_CUSTOM_CHALLENGE_CODE_RESPONSE],
        )
        api._mfa_cognito = mock_cognito

        result = await api.async_request_custom_challenge_code(
            code_type="email",
            mfa_tokens={"Session": "initial_session"},
        )

        assert result["Session"] == "mock_custom_challenge_code_session"

        # Verify the correct ANSWER was sent
        call_kwargs = mock_cognito.client.respond_to_auth_challenge.call_args
        assert "answerType:generateCode,medium:email,codeType:login" in str(call_kwargs)

    @pytest.mark.asyncio
    async def test_request_code_phone(self) -> None:
        """Test requesting verification code via phone."""
        api = API(username="test@example.com")

        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[MOCK_CUSTOM_CHALLENGE_CODE_RESPONSE],
        )
        api._mfa_cognito = mock_cognito

        result = await api.async_request_custom_challenge_code(
            code_type="phone",
            mfa_tokens={"Session": "initial_session"},
        )

        assert result["Session"] == "mock_custom_challenge_code_session"

        call_kwargs = mock_cognito.client.respond_to_auth_challenge.call_args
        assert "answerType:generateCode,medium:phone,codeType:login" in str(call_kwargs)

    @pytest.mark.asyncio
    async def test_request_code_invalid_type(self) -> None:
        """Test requesting code with invalid code_type raises ValueError."""
        api = API(username="test@example.com")
        api._mfa_cognito = MagicMock()

        with pytest.raises(ValueError, match="code_type must be"):
            await api.async_request_custom_challenge_code(
                code_type="fax",
                mfa_tokens={"Session": "session"},
            )

    @pytest.mark.asyncio
    async def test_request_code_no_session(self) -> None:
        """Test requesting code without session raises Unauthenticated."""
        api = API(username="test@example.com")
        api._mfa_cognito = MagicMock()

        with pytest.raises(Unauthenticated, match="Missing session token"):
            await api.async_request_custom_challenge_code(
                code_type="email",
                mfa_tokens={},
            )

    @pytest.mark.asyncio
    async def test_request_code_no_mfa_tokens(self) -> None:
        """Test requesting code without mfa_tokens raises Unauthenticated."""
        api = API(username="test@example.com")
        api._mfa_cognito = MagicMock()

        with pytest.raises(Unauthenticated, match="Missing session token"):
            await api.async_request_custom_challenge_code(
                code_type="email",
                mfa_tokens=None,
            )

    @pytest.mark.asyncio
    async def test_request_code_no_cognito(self) -> None:
        """Test requesting code without login raises Unauthenticated."""
        api = API()

        with pytest.raises(Unauthenticated, match="Please login first"):
            await api.async_request_custom_challenge_code(
                code_type="email",
                mfa_tokens={"Session": "session"},
            )

    @pytest.mark.asyncio
    async def test_request_code_aws_error(self) -> None:
        """Test requesting code with AWS error raises mapped exception."""
        api = API(username="test@example.com")

        error = create_client_error("NotAuthorizedException", "Invalid session")
        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[error],
        )
        api._mfa_cognito = mock_cognito

        with pytest.raises(Unauthenticated):
            await api.async_request_custom_challenge_code(
                code_type="email",
                mfa_tokens={"Session": "expired_session"},
            )

    @pytest.mark.asyncio
    async def test_request_code_creates_cognito_without_stored(self) -> None:
        """Test requesting code creates cognito client when _mfa_cognito is None."""
        api = API(username="test@example.com")
        api._mfa_cognito = None

        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = create_mock_cognito_with_client(
                respond_side_effect=[MOCK_CUSTOM_CHALLENGE_CODE_RESPONSE],
            )
            mock_create.return_value = mock_cognito

            result = await api.async_request_custom_challenge_code(
                code_type="email",
                mfa_tokens={"Session": "initial_session"},
            )

        assert result["Session"] == "mock_custom_challenge_code_session"
        mock_create.assert_called_once()


# ============================================================================
# Custom Challenge MFA Response Tests
# ============================================================================


class TestCustomChallengeMFAResponse:
    """Tests for CUSTOM_CHALLENGE in async_respond_to_mfa_challenge."""

    @pytest.mark.asyncio
    async def test_custom_challenge_success(self) -> None:
        """Test successful CUSTOM_CHALLENGE verification."""
        api = API(username="test@example.com")

        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[MOCK_CUSTOM_CHALLENGE_AUTH_RESULT],
        )

        with patch.object(api, "_create_cognito_client", return_value=mock_cognito):
            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens={"Session": "code_session"},
            )

        assert api.is_authenticated
        assert api.id_token == "custom_challenge_id_token"
        assert api.access_token == "custom_challenge_access_token"
        assert api.refresh_token == "custom_challenge_refresh_token"
        assert api._mfa_cognito is None

    @pytest.mark.asyncio
    async def test_custom_challenge_no_session(self) -> None:
        """Test CUSTOM_CHALLENGE without session raises Unauthenticated."""
        api = API(username="test@example.com")
        api._mfa_cognito = MagicMock()

        with pytest.raises(Unauthenticated, match="CUSTOM_CHALLENGE requires mfa_tokens"):
            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens={},
            )

    @pytest.mark.asyncio
    async def test_custom_challenge_no_mfa_tokens(self) -> None:
        """Test CUSTOM_CHALLENGE without mfa_tokens raises Unauthenticated."""
        api = API(username="test@example.com")
        api._mfa_cognito = MagicMock()

        with pytest.raises(Unauthenticated, match="CUSTOM_CHALLENGE requires mfa_tokens"):
            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens=None,
            )

    @pytest.mark.asyncio
    async def test_custom_challenge_no_auth_result(self) -> None:
        """Test CUSTOM_CHALLENGE without AuthenticationResult raises UnknownError."""
        api = API(username="test@example.com")

        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[
                {
                    "ChallengeName": "CUSTOM_CHALLENGE",
                    "Session": "another_session",
                }
            ],
        )

        with (
            patch.object(api, "_create_cognito_client", return_value=mock_cognito),
            pytest.raises(UnknownError, match="did not return authentication tokens"),
        ):
            await api.async_respond_to_mfa_challenge(
                mfa_code="wrong_code",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens={"Session": "code_session"},
            )

    @pytest.mark.asyncio
    async def test_custom_challenge_aws_error(self) -> None:
        """Test CUSTOM_CHALLENGE with AWS error raises mapped exception."""
        api = API(username="test@example.com")

        error = create_client_error("NotAuthorizedException", "Invalid code")
        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[error],
        )

        with (
            patch.object(api, "_create_cognito_client", return_value=mock_cognito),
            pytest.raises(Unauthenticated),
        ):
            await api.async_respond_to_mfa_challenge(
                mfa_code="000000",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens={"Session": "code_session"},
            )

    @pytest.mark.asyncio
    async def test_custom_challenge_correct_boto3_call(self) -> None:
        """Test CUSTOM_CHALLENGE makes correct boto3 API call."""
        api = API(username="test@example.com")

        mock_cognito = create_mock_cognito_with_client(
            respond_side_effect=[MOCK_CUSTOM_CHALLENGE_AUTH_RESULT],
        )

        with patch.object(api, "_create_cognito_client", return_value=mock_cognito):
            await api.async_respond_to_mfa_challenge(
                mfa_code="654321",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens={"Session": "test_session_123"},
            )

        mock_cognito.client.respond_to_auth_challenge.assert_called_once_with(
            ClientId=api.client_id,
            ChallengeName="CUSTOM_CHALLENGE",
            Session="test_session_123",
            ChallengeResponses={
                "ANSWER": "654321",
                "USERNAME": "test@example.com",
            },
        )


# ============================================================================
# _authenticate_srp Unit Tests
# ============================================================================


class TestAuthenticateSrp:
    """Tests for _authenticate_srp method."""

    def test_srp_success_sets_tokens(self) -> None:
        """Test _authenticate_srp sets tokens on cognito for success response."""
        api = API()

        mock_cognito = MagicMock()
        mock_cognito.username = "test@example.com"
        mock_cognito.client = MagicMock()

        success_response = {
            "AuthenticationResult": {
                "IdToken": "id",
                "AccessToken": "access",
                "RefreshToken": "refresh",
            }
        }

        with patch("aiokwikset.api.AWSSRP") as mock_awssrp:
            mock_srp_instance = MagicMock()
            mock_srp_instance.authenticate_user.return_value = success_response
            mock_awssrp.return_value = mock_srp_instance

            result = api._authenticate_srp(mock_cognito, "password")

        assert result == success_response
        mock_cognito._set_tokens.assert_called_once_with(success_response)

    def test_srp_custom_challenge_no_set_tokens(self) -> None:
        """Test _authenticate_srp does NOT call _set_tokens for CUSTOM_CHALLENGE."""
        api = API()

        mock_cognito = MagicMock()
        mock_cognito.username = "test@example.com"
        mock_cognito.client = MagicMock()

        challenge_response = {
            "ChallengeName": "CUSTOM_CHALLENGE",
            "Session": "session_123",
        }

        with patch("aiokwikset.api.AWSSRP") as mock_awssrp:
            mock_srp_instance = MagicMock()
            mock_srp_instance.authenticate_user.return_value = challenge_response
            mock_awssrp.return_value = mock_srp_instance

            result = api._authenticate_srp(mock_cognito, "password")

        assert result == challenge_response
        mock_cognito._set_tokens.assert_not_called()

    def test_srp_mfa_exception_propagates(self) -> None:
        """Test _authenticate_srp propagates pycognito MFA exceptions."""
        api = API()

        mock_cognito = MagicMock()
        mock_cognito.username = "test@example.com"
        mock_cognito.client = MagicMock()

        with patch("aiokwikset.api.AWSSRP") as mock_awssrp:
            mock_srp_instance = MagicMock()
            exc = SMSMFAChallengeException("Do SMS MFA", {"Session": "s"})
            mock_srp_instance.authenticate_user.side_effect = exc
            mock_awssrp.return_value = mock_srp_instance

            with pytest.raises(SMSMFAChallengeException):
                api._authenticate_srp(mock_cognito, "password")

    def test_srp_creates_awssrp_with_correct_params(self) -> None:
        """Test _authenticate_srp creates AWSSRP with correct parameters."""
        api = API()

        mock_cognito = MagicMock()
        mock_cognito.username = "test@example.com"
        mock_cognito.client = MagicMock()

        with patch("aiokwikset.api.AWSSRP") as mock_awssrp:
            mock_srp_instance = MagicMock()
            mock_srp_instance.authenticate_user.return_value = {
                "AuthenticationResult": {"IdToken": "t", "AccessToken": "t", "RefreshToken": "t"}
            }
            mock_awssrp.return_value = mock_srp_instance

            api._authenticate_srp(mock_cognito, "test_password")

        mock_awssrp.assert_called_once_with(
            username="test@example.com",
            password="test_password",
            pool_id=api.user_pool_id,
            client_id=api.client_id,
            client=mock_cognito.client,
        )


# ============================================================================
# Full CUSTOM_CHALLENGE Flow Integration Test
# ============================================================================


class TestCustomChallengeFullFlow:
    """Integration test for the full CUSTOM_CHALLENGE flow."""

    @pytest.mark.asyncio
    async def test_full_custom_challenge_flow(self) -> None:
        """Test complete flow: login → request code → verify code."""
        api = API()

        # Step 1: Login triggers CUSTOM_CHALLENGE
        with patch.object(api, "_create_cognito_client") as mock_create:
            mock_cognito = MagicMock()
            mock_cognito.username = "test@example.com"
            mock_cognito.client = MagicMock()
            mock_create.return_value = mock_cognito

            with patch.object(api, "_authenticate_srp") as mock_srp:
                mock_srp.return_value = MOCK_CUSTOM_CHALLENGE_RESPONSE

                with pytest.raises(MFAChallengeRequired) as exc_info:
                    await api.async_login("test@example.com", "password123")

        challenge_tokens = exc_info.value.mfa_tokens
        assert exc_info.value.mfa_type == "CUSTOM_CHALLENGE"

        # Step 2: Request code via email
        mock_cognito.client.respond_to_auth_challenge.return_value = (
            MOCK_CUSTOM_CHALLENGE_CODE_RESPONSE
        )

        code_tokens = await api.async_request_custom_challenge_code(
            code_type="email",
            mfa_tokens=challenge_tokens,
        )

        assert code_tokens["Session"] == "mock_custom_challenge_code_session"

        # Step 3: Submit verification code
        mock_cognito.client.respond_to_auth_challenge.return_value = (
            MOCK_CUSTOM_CHALLENGE_AUTH_RESULT
        )

        with patch.object(api, "_create_cognito_client", return_value=mock_cognito):
            await api.async_respond_to_mfa_challenge(
                mfa_code="123456",
                mfa_type="CUSTOM_CHALLENGE",
                mfa_tokens=code_tokens,
            )

        assert api.is_authenticated
        assert api.id_token == "custom_challenge_id_token"
        assert api._mfa_cognito is None
