"""
Test key rotation manager functionality
"""

import os
import sys
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock

import pytest

sys.path.insert(0, os.path.abspath("../../../.."))

from litellm.proxy._types import (
    GenerateKeyResponse,
    LiteLLM_VerificationToken,
)
from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager


class TestKeyRotationManager:
    """Test the KeyRotationManager class functionality."""

    @pytest.mark.asyncio
    async def test_should_rotate_key_logic(self):
        """
        Test the core logic for determining when a key should be rotated.

        This tests:
        - Keys with null key_rotation_at should rotate immediately
        - Keys with future key_rotation_at should not rotate
        - Keys with past key_rotation_at should rotate
        """
        # Setup
        mock_prisma_client = AsyncMock()
        manager = KeyRotationManager(mock_prisma_client)

        now = datetime.now(timezone.utc)

        # Test Case 1: No rotation time set (key_rotation_at = None) - should rotate
        key_no_rotation_time = LiteLLM_VerificationToken(
            token="test-token-1",
            auto_rotate=True,
            rotation_interval="30s",
            key_rotation_at=None,
            rotation_count=0,
        )

        assert manager._should_rotate_key(key_no_rotation_time, now) is True

        # Test Case 2: Future rotation time - should NOT rotate
        key_future_rotation = LiteLLM_VerificationToken(
            token="test-token-2",
            auto_rotate=True,
            rotation_interval="30s",
            key_rotation_at=now + timedelta(seconds=10),
            rotation_count=1,
        )

        assert manager._should_rotate_key(key_future_rotation, now) is False

        # Test Case 3: Past rotation time - should rotate
        key_past_rotation = LiteLLM_VerificationToken(
            token="test-token-3",
            auto_rotate=True,
            rotation_interval="30s",
            key_rotation_at=now - timedelta(seconds=10),
            rotation_count=2,
        )

        assert manager._should_rotate_key(key_past_rotation, now) is True

        # Test Case 4: Exact rotation time - should rotate
        key_exact_rotation = LiteLLM_VerificationToken(
            token="test-token-4",
            auto_rotate=True,
            rotation_interval="30s",
            key_rotation_at=now,
            rotation_count=1,
        )

        assert manager._should_rotate_key(key_exact_rotation, now) is True

        # Test Case 5: No rotation interval - should NOT rotate
        key_no_interval = LiteLLM_VerificationToken(
            token="test-token-5",
            auto_rotate=True,
            rotation_interval=None,
            key_rotation_at=None,
            rotation_count=0,
        )

        assert manager._should_rotate_key(key_no_interval, now) is False

    @pytest.mark.asyncio
    async def test_find_keys_needing_rotation(self):
        """
        Test finding keys that need rotation from database.

        This tests:
        - Only keys with auto_rotate=True are considered
        - Database query filters by key_rotation_at properly
        - Keys are returned based on key_rotation_at being null or <= now
        """
        # Setup
        mock_prisma_client = AsyncMock()
        manager = KeyRotationManager(mock_prisma_client)

        # Use a fixed timestamp to avoid timing issues in tests
        now = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)

        # Mock database response - these are the keys the database query would return
        mock_keys = [
            LiteLLM_VerificationToken(
                token="token-1",
                auto_rotate=True,
                rotation_interval="30s",
                key_rotation_at=None,  # Should rotate (null key_rotation_at)
                rotation_count=0,
            ),
            LiteLLM_VerificationToken(
                token="token-2",
                auto_rotate=True,
                rotation_interval="60s",
                key_rotation_at=now
                - timedelta(seconds=10),  # Should rotate (past time)
                rotation_count=1,
            ),
        ]

        mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = (
            mock_keys
        )

        # Mock datetime.now to return our fixed timestamp
        from unittest.mock import patch

        with patch(
            "litellm.proxy.common_utils.key_rotation_manager.datetime"
        ) as mock_datetime:
            mock_datetime.now.return_value = now
            mock_datetime.side_effect = lambda *args, **kwargs: datetime(
                *args, **kwargs
            )

            # Execute
            keys_needing_rotation = await manager._find_keys_needing_rotation()

        # Verify database query - should use OR condition for key_rotation_at
        mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
            where={
                "auto_rotate": True,
                "OR": [{"key_rotation_at": None}, {"key_rotation_at": {"lte": now}}],
            }
        )

        # Verify all keys returned by database query are included (no additional filtering)
        assert len(keys_needing_rotation) == 2

        tokens_needing_rotation = [key.token for key in keys_needing_rotation]
        assert "token-1" in tokens_needing_rotation  # Null key_rotation_at
        assert "token-2" in tokens_needing_rotation  # Past key_rotation_at

    @pytest.mark.asyncio
    async def test_rotate_key_updates_database(self):
        """
        Test that key rotation properly updates the database with new rotation info.

        This tests:
        - Rotation count is incremented
        - last_rotation_at is set to current time
        - key_rotation_at is set to next rotation time
        - New key token is updated (not old one)
        """
        # Setup
        mock_prisma_client = AsyncMock()
        manager = KeyRotationManager(mock_prisma_client)

        # Mock key to rotate
        key_to_rotate = LiteLLM_VerificationToken(
            token="old-token",
            auto_rotate=True,
            rotation_interval="30s",
            last_rotation_at=None,
            key_rotation_at=None,
            rotation_count=0,
        )

        # Mock regenerate_key_fn response
        mock_response = GenerateKeyResponse(
            key="new-api-key", token_id="new-token-id", user_id="test-user"
        )

        # Mock the regenerate function
        from unittest.mock import patch

        with patch(
            "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
            return_value=mock_response,
        ):
            with patch(
                "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook"
            ):
                # Execute
                await manager._rotate_key(key_to_rotate)

        # Verify database update was called with correct data
        mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()

        call_args = mock_prisma_client.db.litellm_verificationtoken.update.call_args

        # Check the WHERE clause targets the new token
        assert call_args[1]["where"]["token"] == "new-token-id"

        # Check the data being updated
        update_data = call_args[1]["data"]
        assert update_data["rotation_count"] == 1  # Incremented from 0
        assert "last_rotation_at" in update_data
        assert isinstance(update_data["last_rotation_at"], datetime)
        assert "key_rotation_at" in update_data
        assert isinstance(update_data["key_rotation_at"], datetime)

        # Verify key_rotation_at is set to future time (30s from now)
        now = datetime.now(timezone.utc)
        next_rotation = update_data["key_rotation_at"]
        time_diff = (next_rotation - now).total_seconds()
        assert (
            25 <= time_diff <= 35
        )  # Should be around 30 seconds, allow some tolerance

    @pytest.mark.asyncio
    async def test_cleanup_expired_deprecated_keys(self):
        """
        Test that _cleanup_expired_deprecated_keys deletes expired deprecated keys.
        """
        mock_prisma_client = AsyncMock()
        mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.return_value = (
            3
        )
        manager = KeyRotationManager(mock_prisma_client)

        await manager._cleanup_expired_deprecated_keys()

        mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once()
        call_args = (
            mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.call_args
        )
        assert "revoke_at" in call_args[1]["where"]
        assert call_args[1]["where"]["revoke_at"]["lt"] is not None

    @pytest.mark.asyncio
    async def test_rotate_key_passes_grace_period(self):
        """
        Test that _rotate_key passes grace_period in RegenerateKeyRequest.
        """
        mock_prisma_client = AsyncMock()
        manager = KeyRotationManager(mock_prisma_client)

        key_to_rotate = LiteLLM_VerificationToken(
            token="old-token",
            auto_rotate=True,
            rotation_interval="30s",
            key_rotation_at=None,
            rotation_count=0,
        )

        mock_response = GenerateKeyResponse(
            key="new-api-key",
            token_id="new-token-id",
            user_id="test-user",
        )

        from unittest.mock import patch

        with patch(
            "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
            new_callable=AsyncMock,
        ) as mock_regenerate:
            mock_regenerate.return_value = mock_response
            with patch(
                "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
                new_callable=AsyncMock,
            ):
                with patch(
                    "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD",
                    "48h",
                ):
                    await manager._rotate_key(key_to_rotate)

            mock_regenerate.assert_called_once()
            call_args = mock_regenerate.call_args
            regenerate_request = call_args[1]["data"]
            assert regenerate_request.grace_period == "48h"
