from unittest.mock import AsyncMock, patch

import pytest
from server.auth.token_manager import TokenManager, create_encryption_utility

from openhands.integrations.service_types import ProviderType


@pytest.fixture
def token_manager():
    with patch('server.config.get_config') as mock_get_config:
        mock_config = mock_get_config.return_value
        mock_config.jwt_secret.get_secret_value.return_value = 'test_secret'
        return TokenManager(external=False)


def test_create_encryption_utility():
    """Test the encryption utility creation and functionality."""
    secret_key = b'test_secret_key_that_is_32_bytes_lng'
    encrypt_payload, decrypt_payload, encrypt_text, decrypt_text = (
        create_encryption_utility(secret_key)
    )

    # Test text encryption/decryption
    original_text = 'This is a test message'
    encrypted = encrypt_text(original_text)
    decrypted = decrypt_text(encrypted)
    assert decrypted == original_text
    assert encrypted != original_text

    # Test payload encryption/decryption
    original_payload = {'key1': 'value1', 'key2': 123, 'nested': {'inner': 'value'}}
    encrypted = encrypt_payload(original_payload)
    decrypted = decrypt_payload(encrypted)
    assert decrypted == original_payload
    assert encrypted != original_payload


@pytest.mark.asyncio
async def test_get_keycloak_tokens_success(token_manager):
    """Test successful retrieval of Keycloak tokens."""
    mock_token_response = {
        'access_token': 'test_access_token',
        'refresh_token': 'test_refresh_token',
    }

    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_token = AsyncMock(return_value=mock_token_response)

        access_token, refresh_token = await token_manager.get_keycloak_tokens(
            'test_code', 'http://test.com/callback'
        )

        assert access_token == 'test_access_token'
        assert refresh_token == 'test_refresh_token'
        mock_keycloak.return_value.a_token.assert_called_once_with(
            grant_type='authorization_code',
            code='test_code',
            redirect_uri='http://test.com/callback',
        )


@pytest.mark.asyncio
async def test_get_keycloak_tokens_missing_tokens(token_manager):
    """Test handling of missing tokens in Keycloak response."""
    mock_token_response = {
        'access_token': 'test_access_token',
        # Missing refresh_token
    }

    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_token = AsyncMock(return_value=mock_token_response)

        access_token, refresh_token = await token_manager.get_keycloak_tokens(
            'test_code', 'http://test.com/callback'
        )

        assert access_token is None
        assert refresh_token is None


@pytest.mark.asyncio
async def test_get_keycloak_tokens_exception(token_manager):
    """Test handling of exceptions during token retrieval."""
    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_token = AsyncMock(
            side_effect=Exception('Test error')
        )

        access_token, refresh_token = await token_manager.get_keycloak_tokens(
            'test_code', 'http://test.com/callback'
        )

        assert access_token is None
        assert refresh_token is None


@pytest.mark.asyncio
async def test_verify_keycloak_token_valid(token_manager):
    """Test verification of a valid Keycloak token."""
    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_userinfo = AsyncMock(
            return_value={'sub': 'test_user_id'}
        )

        access_token, refresh_token = await token_manager.verify_keycloak_token(
            'test_access_token', 'test_refresh_token'
        )

        assert access_token == 'test_access_token'
        assert refresh_token == 'test_refresh_token'
        mock_keycloak.return_value.a_userinfo.assert_called_once_with(
            'test_access_token'
        )


@pytest.mark.asyncio
async def test_verify_keycloak_token_refresh(token_manager):
    """Test refreshing an invalid Keycloak token."""
    from keycloak.exceptions import KeycloakAuthenticationError

    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_userinfo = AsyncMock(
            side_effect=KeycloakAuthenticationError('Invalid token')
        )
        mock_keycloak.return_value.a_refresh_token = AsyncMock(
            return_value={
                'access_token': 'new_access_token',
                'refresh_token': 'new_refresh_token',
            }
        )

        access_token, refresh_token = await token_manager.verify_keycloak_token(
            'test_access_token', 'test_refresh_token'
        )

        assert access_token == 'new_access_token'
        assert refresh_token == 'new_refresh_token'
        mock_keycloak.return_value.a_userinfo.assert_called_once_with(
            'test_access_token'
        )
        mock_keycloak.return_value.a_refresh_token.assert_called_once_with(
            'test_refresh_token'
        )


@pytest.mark.asyncio
async def test_get_user_info(token_manager):
    """Test getting user info from a Keycloak token."""
    mock_user_info = {
        'sub': 'test_user_id',
        'name': 'Test User',
        'email': 'test@example.com',
    }

    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_userinfo = AsyncMock(return_value=mock_user_info)

        user_info = await token_manager.get_user_info('test_access_token')

        assert user_info == mock_user_info
        mock_keycloak.return_value.a_userinfo.assert_called_once_with(
            'test_access_token'
        )


@pytest.mark.asyncio
async def test_get_user_info_empty_token(token_manager):
    """Test handling of empty token when getting user info."""
    user_info = await token_manager.get_user_info('')

    assert user_info == {}


@pytest.mark.asyncio
async def test_store_idp_tokens(token_manager):
    """Test storing identity provider tokens."""
    mock_idp_tokens = {
        'access_token': 'github_access_token',
        'refresh_token': 'github_refresh_token',
        'access_token_expires_at': 1000,
        'refresh_token_expires_at': 2000,
    }

    with (
        patch.object(
            token_manager, 'get_idp_tokens_from_keycloak', return_value=mock_idp_tokens
        ),
        patch.object(token_manager, '_store_idp_tokens') as mock_store,
    ):
        await token_manager.store_idp_tokens(
            ProviderType.GITHUB, 'test_user_id', 'test_access_token'
        )

        mock_store.assert_called_once_with(
            'test_user_id',
            ProviderType.GITHUB,
            'github_access_token',
            'github_refresh_token',
            1000,
            2000,
        )


@pytest.mark.asyncio
async def test_get_idp_token(token_manager):
    """Test getting an identity provider token."""
    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info',
            AsyncMock(return_value={'sub': 'test_user_id'}),
        ),
        patch('server.auth.token_manager.AuthTokenStore') as mock_token_store_cls,
    ):
        mock_token_store = AsyncMock()
        mock_token_store.return_value.load_tokens.return_value = {
            'access_token': token_manager.encrypt_text('github_access_token'),
        }
        mock_token_store_cls.get_instance = mock_token_store

        token = await token_manager.get_idp_token(
            'test_access_token', ProviderType.GITHUB
        )

        assert token == 'github_access_token'
        mock_token_store_cls.get_instance.assert_called_once_with(
            keycloak_user_id='test_user_id', idp=ProviderType.GITHUB
        )
        mock_token_store.return_value.load_tokens.assert_called_once()


@pytest.mark.asyncio
async def test_refresh(token_manager):
    """Test refreshing a token."""
    mock_tokens = {
        'access_token': 'new_access_token',
        'refresh_token': 'new_refresh_token',
    }

    with patch('server.auth.token_manager.get_keycloak_openid') as mock_keycloak:
        mock_keycloak.return_value.a_refresh_token = AsyncMock(return_value=mock_tokens)

        result = await token_manager.refresh('test_refresh_token')

        assert result == mock_tokens
        mock_keycloak.return_value.a_refresh_token.assert_called_once_with(
            'test_refresh_token'
        )
