import pytest
import time
from aiva.email.throttle_manager import ThrottleManager
from aiva.email.throttle_config import ThrottleConfig


@pytest.fixture
def throttle_manager():
    config = {
        "example.com": ThrottleConfig(rate=2, period=1, burst_size=4),
        "gmail.com": ThrottleConfig(rate=1, period=0.5, burst_size=2)
    }
    return ThrottleManager(config)


def test_throttle_manager_allow_request_within_rate(throttle_manager):
    assert throttle_manager.allow_request("test@example.com")
    assert throttle_manager.allow_request("another@example.com")
    assert not throttle_manager.allow_request("third@example.com")


def test_throttle_manager_allow_request_burst_size(throttle_manager):
    for _ in range(4):
        assert throttle_manager.allow_request("test@example.com")
    assert not throttle_manager.allow_request("fifth@example.com")


def test_throttle_manager_allow_request_period_reset(throttle_manager):
    assert throttle_manager.allow_request("test@example.com")
    assert not throttle_manager.allow_request("test@example.com")
    time.sleep(1.1)  # Wait for the period to expire
    assert throttle_manager.allow_request("test@example.com")


def test_throttle_manager_allow_request_different_domains(throttle_manager):
    assert throttle_manager.allow_request("test@example.com")
    assert not throttle_manager.allow_request("test2@example.com")
    assert throttle_manager.allow_request("test@gmail.com")
    assert not throttle_manager.allow_request("test2@gmail.com")


def test_throttle_manager_allow_request_no_config_domain(throttle_manager):
    assert throttle_manager.allow_request("test@unknown.com")  # Should allow since no config exists


def test_throttle_manager_config_update():
    config = {
        "example.com": ThrottleConfig(rate=2, period=1, burst_size=4),
    }
    throttle_manager = ThrottleManager(config)
    assert throttle_manager.allow_request("test@example.com")
    config["example.com"] = ThrottleConfig(rate=5, period=1, burst_size=10)
    throttle_manager.update_config(config)
    for _ in range(5):
        assert throttle_manager.allow_request("test@example.com")
    assert not throttle_manager.allow_request("test@example.com")
