"""Unit tests for generate_jwt_token_with_impersonation."""
import logging
from typing import Any
from unittest.mock import MagicMock, call, patch

import httpx
import pytest

from owsclient.m2m.impersonation import (
    Auth0TokenException,
    ClientCredentials,
    OAuthToken,
    _generate_jwt_token_with_impersonation,
    generate_jwt_token_with_impersonation_with_retries,
)


def test__generate_jwt_token_with_impersonation(
    valid_client_credentials_data: dict[str, Any],
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    impersonated_identity_uuid: str,
) -> None:
    """Test successful token generation with 200 status code."""
    mock_client = MagicMock(spec=httpx.Client)
    mock_client.post.return_value = httpx.Response(
        200,
        json=valid_oauth_token_data,
    )
    result = _generate_jwt_token_with_impersonation(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    assert result.status_code == 200
    assert result.json() == valid_oauth_token_data
    mock_client.post.assert_called_once_with(
        "/oauth/token",
        json={
            **valid_client_credentials_data,
            "impersonate_identity_uuid": impersonated_identity_uuid,
        },
    )


def test__generate_jwt_token_with_impersonation_bad_payload(
    valid_client_credentials_data: dict[str, Any],
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    impersonated_identity_uuid: str,
) -> None:
    """Test 200 status code and invalid response payload."""
    invalid_oauth_token_data = {
        **valid_oauth_token_data,
        "access_token": "",
    }

    mock_client = MagicMock(spec=httpx.Client)
    mock_client.post.return_value = httpx.Response(
        200,
        json=invalid_oauth_token_data,
    )
    result = _generate_jwt_token_with_impersonation(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    assert result.status_code == 200
    assert result.json() == invalid_oauth_token_data
    mock_client.post.assert_called_once_with(
        "/oauth/token",
        json={
            **valid_client_credentials_data,
            "impersonate_identity_uuid": impersonated_identity_uuid,
        },
    )


@pytest.mark.parametrize(
    "expected_status_code, expected_text",
    [
        pytest.param(401, "401 Unauthorized"),
        pytest.param(429, "429 Too Many Requests"),
        pytest.param(502, "502 Bad Gateway"),
        pytest.param(503, "503 Service Unavailable"),
        pytest.param(504, "504 Gateway Timeout"),
    ],
)
def test__generate_jwt_token_with_impersonation_non_200(
    expected_status_code: int,
    expected_text: str,
    valid_client_credentials_data: dict[str, Any],
    valid_client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
) -> None:
    """Test token generation with non-200 status code."""
    mock_client = MagicMock(spec=httpx.Client)
    mock_client.post.return_value = httpx.Response(
        expected_status_code,
        text=expected_text,
    )

    result = _generate_jwt_token_with_impersonation(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    assert result.status_code == expected_status_code
    assert result.text == expected_text
    mock_client.post.assert_called_once_with(
        "/oauth/token",
        json={
            **valid_client_credentials_data,
            "impersonate_identity_uuid": impersonated_identity_uuid,
        },
    )


def test__generate_jwt_token_with_impersonation_connect_error(
    valid_client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
) -> None:
    """Test token generation with ConnectError raises."""
    mock_client = MagicMock(spec=httpx.Client)
    mock_client.post.side_effect = httpx.ConnectError("Connection failed")

    with pytest.raises(
        Auth0TokenException, match="Network error while requesting JWT Access Token"
    ):
        _generate_jwt_token_with_impersonation(
            client_credentials=valid_client_credentials,
            impersonated_identity_uuid=impersonated_identity_uuid,
            client=mock_client,
        )


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_impersonation_with_retries_success(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    valid_oauth_token: OAuthToken,
    impersonated_identity_uuid: str,
    auth0_url: str,
) -> None:
    """Test successful token generation with 200 status code on the first call to auth0."""
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client
    mock_generate_jwt.return_value = httpx.Response(
        200,
        json=valid_oauth_token_data,
    )
    result = generate_jwt_token_with_impersonation_with_retries(
        valid_client_credentials, impersonated_identity_uuid, auth0_url=auth0_url
    )

    mock_httpx.HTTPTransport.assert_called_once_with(retries=3)
    mock_httpx.Client.assert_called_once_with(
        base_url=auth0_url,
        transport=mock_transport,
        timeout=10,
    )

    assert isinstance(result, OAuthToken)
    assert result == valid_oauth_token
    mock_generate_jwt.assert_called_once_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    mock_get_backoff.assert_not_called()
    mock_sleep.assert_not_called()


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_impersonation_with_retries_raises_for_bad_response(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    impersonated_identity_uuid: str,
    auth0_url: str,
) -> None:
    """Test bad auth0 response body raises, even with 200 status code."""
    invalid_oauth_token_data = {
        **valid_oauth_token_data,
        "access_token": "",
    }
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client
    mock_generate_jwt.return_value = httpx.Response(
        200,
        json=invalid_oauth_token_data,
    )
    with pytest.raises(
        Auth0TokenException, match="JWT Access Token could not be parsed"
    ):
        generate_jwt_token_with_impersonation_with_retries(
            valid_client_credentials, impersonated_identity_uuid, auth0_url=auth0_url
        )

    mock_httpx.HTTPTransport.assert_called_once_with(retries=3)
    mock_httpx.Client.assert_called_once_with(
        base_url=auth0_url,
        transport=mock_transport,
        timeout=10,
    )

    mock_generate_jwt.assert_called_once_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    mock_get_backoff.assert_not_called()
    mock_sleep.assert_not_called()


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_impersonation_with_retries_failure(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
    auth0_url: str,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Test token generation failure with non-200/non-retryable status code."""
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client

    mock_generate_jwt.return_value = httpx.Response(401, text="Unauthorized")

    with pytest.raises(Auth0TokenException) as exc_info:
        generate_jwt_token_with_impersonation_with_retries(
            valid_client_credentials, impersonated_identity_uuid, auth0_url=auth0_url
        )

    mock_httpx.HTTPTransport.assert_called_once_with(retries=3)
    mock_httpx.Client.assert_called_once_with(
        base_url=auth0_url,
        transport=mock_transport,
        timeout=10,
    )

    assert "JWT Access Token could not be generated" in str(exc_info.value)
    mock_generate_jwt.assert_called_once_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    mock_get_backoff.assert_not_called()
    mock_sleep.assert_not_called()
    assert (
        caplog.records[0].msg == "ImpersonationM2MTokenManager unable to generate JWT"
    )


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_impersonation_with_retries_network_error(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
    auth0_url: str,
) -> None:
    """Test token generation fails on network error."""
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client

    mock_generate_jwt.side_effect = Auth0TokenException(
        "Network error while requesting JWT Access Token"
    )

    with pytest.raises(
        Auth0TokenException, match="Network error while requesting JWT Access Token"
    ):
        generate_jwt_token_with_impersonation_with_retries(
            valid_client_credentials, impersonated_identity_uuid, auth0_url=auth0_url
        )
    mock_generate_jwt.assert_called_once_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    mock_get_backoff.assert_not_called()
    mock_sleep.assert_not_called()


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_retries_eventually_succeeds(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    valid_oauth_token: OAuthToken,
    impersonated_identity_uuid: str,
    auth0_url: str,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Test token generation with retryable status codes."""
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client

    mock_generate_jwt.side_effect = [
        httpx.Response(502, text="Bad Gateway"),
        httpx.Response(503, text="Service Unavailable"),
        httpx.Response(504, text="Gateway Timeout"),
        httpx.Response(
            200,
            json=valid_oauth_token_data,
        ),
    ]
    mock_get_backoff.side_effect = [4, 7, 2]

    with caplog.at_level(logging.INFO):
        result = generate_jwt_token_with_impersonation_with_retries(
            valid_client_credentials,
            impersonated_identity_uuid,
            auth0_url=auth0_url,
            connect_retries=1,
            timeout=1,
        )

    assert isinstance(result, OAuthToken)
    assert result == valid_oauth_token

    mock_httpx.HTTPTransport.assert_called_once_with(retries=1)
    mock_httpx.Client.assert_called_once_with(
        base_url=auth0_url,
        transport=mock_transport,
        timeout=1,
    )

    assert mock_generate_jwt.call_count == 4
    mock_generate_jwt.assert_called_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    assert mock_get_backoff.call_count == 3
    mock_get_backoff.assert_has_calls(
        [
            call(1, 3, 1),
            call(1, 3, 2),
            call(1, 3, 3),
        ]
    )
    assert mock_sleep.call_count == 3
    mock_sleep.assert_has_calls(
        [
            call(4),
            call(7),
            call(2),
        ]
    )

    for record in caplog.records:
        assert record.levelname in ("INFO"), f"Unexpected log {record.msg}"
        assert (
            record.msg
            == "ImpersonationM2MTokenManager is retrying request to generate a JWT"
        )


@patch("owsclient.m2m.impersonation.sleep")
@patch("owsclient.m2m.impersonation.get_backoff_with_full_jitter")
@patch("owsclient.m2m.impersonation._generate_jwt_token_with_impersonation")
@patch("owsclient.m2m.impersonation.httpx")
def test_generate_jwt_token_with_retries_eventually_raises(
    mock_httpx: MagicMock,
    mock_generate_jwt: MagicMock,
    mock_get_backoff: MagicMock,
    mock_sleep: MagicMock,
    valid_client_credentials: ClientCredentials,
    valid_oauth_token_data: dict[str, Any],
    impersonated_identity_uuid: str,
    auth0_url: str,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Test token generation with retryable status codes but not enough retries."""
    mock_transport = MagicMock(spec=httpx.HTTPTransport)
    mock_client = MagicMock(spec=httpx.Client)
    mock_httpx.HTTPTransport.return_value = mock_transport
    mock_httpx.Client.return_value.__enter__.return_value = mock_client

    mock_generate_jwt.side_effect = [
        httpx.Response(429, text="Too Many Requests"),
        httpx.Response(503, text="Service Unavailable"),
        httpx.Response(504, text="Gateway Timeout"),
        httpx.Response(
            200,
            json=valid_oauth_token_data,
        ),
    ]
    mock_get_backoff.side_effect = [7, 8]

    with pytest.raises(
        Auth0TokenException,
        match="'JWT Access Token could not be generated after 2 attempts', 'Gateway Timeout'",
    ):
        with caplog.at_level(logging.INFO):
            generate_jwt_token_with_impersonation_with_retries(
                client_credentials=valid_client_credentials,
                impersonated_identity_uuid=impersonated_identity_uuid,
                auth0_url=auth0_url,
                connect_retries=0,
                timeout=35,
                server_error_retries=2,
            )

    mock_httpx.HTTPTransport.assert_called_once_with(retries=0)
    mock_httpx.Client.assert_called_once_with(
        base_url=auth0_url,
        transport=mock_transport,
        timeout=35,
    )

    assert mock_generate_jwt.call_count == 3
    mock_generate_jwt.assert_called_with(
        client_credentials=valid_client_credentials,
        impersonated_identity_uuid=impersonated_identity_uuid,
        client=mock_client,
    )
    assert mock_get_backoff.call_count == 2
    mock_get_backoff.assert_has_calls(
        [
            call(1, 3, 1),
            call(1, 3, 2),
        ]
    )
    assert mock_sleep.call_count == 2
    mock_sleep.assert_has_calls(
        [
            call(7),
            call(8),
        ]
    )

    for record in caplog.records:
        assert record.levelname in ("INFO", "WARNING"), f"Unexpected log {record.msg}"
        if record.levelname == "INFO":
            assert (
                record.msg
                == "ImpersonationM2MTokenManager is retrying request to generate a JWT"
            )
        if record.levelname == "WARNING":
            assert record.msg == "ImpersonationM2MTokenManager unable to generate JWT"