import time
from typing import List

import pytest
from django.conf import settings

import promptgame.gameui.redis_utils as ru
from promptgame.gameui.rate_limiting import rate_limit_check
from promptgame.settings.custom_types import RateLimitBucketConfig
from promptgame.tests.helpers import sync_clear_redis_test_db


def setup_function() -> None:
    sync_clear_redis_test_db()


def teardown_function() -> None:
    sync_clear_redis_test_db()


@pytest.mark.asyncio
async def test_basic_rate_limiting():
    account_id = 1
    bucket_config: List[RateLimitBucketConfig] = [
        dict(bucket_name="api", refill_rate=1.0, max_capacity=5)
    ]
    for _ in range(5):
        result = await rate_limit_check(account_id, bucket_config)
        assert result.request_allowed

    result = await rate_limit_check(account_id, bucket_config)
    assert not result.request_allowed
    assert 0 < result.wait_time_s <= 1


@pytest.mark.asyncio
async def test_multiple_buckets():
    account_id = 2
    bucket_configs: List[RateLimitBucketConfig] = [
        dict(bucket_name="api", refill_rate=1.0, max_capacity=5),
        dict(bucket_name="upload", refill_rate=1.0, max_capacity=3),
    ]

    for _ in range(3):
        result = await rate_limit_check(account_id, bucket_configs)
        assert result.request_allowed

    result = await rate_limit_check(account_id, bucket_configs)
    assert not result.request_allowed
    assert 0 < result.wait_time_s <= 1


@pytest.mark.asyncio
async def test_refill_over_time():
    account_id = 3
    bucket_config: List[RateLimitBucketConfig] = [
        dict(bucket_name="api", refill_rate=1.0, max_capacity=5)
    ]
    for _ in range(5):
        result = await rate_limit_check(account_id, bucket_config)
        assert result.request_allowed

    time.sleep(1)  # Let some tokens refill

    result = await rate_limit_check(account_id, bucket_config)
    assert result.request_allowed


@pytest.mark.asyncio
async def test_diff_refill_capacities_and_rates():
    account_id = 5
    bucket_configs: List[RateLimitBucketConfig] = [
        # High-capacity, low-rate (fills in 10s)
        dict(bucket_name="low_rate", refill_rate=1.0, max_capacity=10),
        # Low-capacity, high-rate (fills in 300ms)
        dict(bucket_name="high_rate", refill_rate=10.0, max_capacity=3),
    ]

    # Exhaust the high-rate bucket with a quick burst
    for _ in range(3):
        result = await rate_limit_check(account_id, bucket_configs)
        assert result.request_allowed

    # Check if a new request is denied due to the high-rate bucket being
    # exhausted
    result = await rate_limit_check(account_id, bucket_configs)
    assert not result.request_allowed

    # wait for high-rate bucket to refill (high-capacity bucket will only have 7
    # tokens left)
    time.sleep(0.30001)

    # Exhaust the low-rate bucket while keeping the high-rate bucket
    # sufficiently refilled (note that the first 7 loops and the 300ms sleep
    # above are enough to add an extra token to the low-rate bucket as well)
    for _ in range(8):
        result = await rate_limit_check(account_id, bucket_configs)
        assert result.request_allowed
        time.sleep(1 / 10 + 0.001)

    # At this point, the low-rate bucket should be exhausted, but the high-rate
    # bucket should still have tokens
    result = await rate_limit_check(account_id, bucket_configs)
    assert not result.request_allowed


@pytest.mark.asyncio
async def test_redis_key_expiration():
    client = await ru.get_async_client()

    account_id = 6
    bucket_config: List[RateLimitBucketConfig] = [
        dict(bucket_name="api", refill_rate=1.0, max_capacity=5)
    ]
    for _ in range(4):  # Use 4 tokens, leaving 1 token
        await rate_limit_check(account_id, bucket_config)

    # Check the TTL of the tokens_key
    tokens_key = f"{settings.REDIS_KEY_PREFIX}:ratelim:{account_id}:api:tokens"
    initial_ttl = await client.ttl(tokens_key)

    # make sure that it's around 4s (time required to get back to 5 tokens)
    assert abs(initial_ttl - 4) < 0.5

    # Now, let's consume one more token
    await rate_limit_check(account_id, bucket_config)

    # The TTL should be approximately 5 seconds (give or take some time for
    # execution)
    after_request_ttl = await client.ttl(tokens_key)
    assert abs(after_request_ttl - 5) < 0.5
