"""Tests for the rate limiting module."""

from __future__ import annotations

import threading
import time

import pytest

from oclawma.ratelimit import (
    RateLimitConfig,
    RateLimiter,
    RateLimitError,
    RateLimitInfo,
    TokenBucket,
)


class TestTokenBucket:
    """Test TokenBucket implementation."""

    def test_default_creation(self):
        """Test creating a token bucket with defaults."""
        bucket = TokenBucket(rate=10.0, burst=5)

        assert bucket.rate == 10.0
        assert bucket.burst == 5
        assert bucket.tokens == 5  # Starts full
        assert bucket.total_requests == 0
        assert bucket.blocked_requests == 0

    def test_custom_creation(self):
        """Test creating a token bucket with custom values."""
        bucket = TokenBucket(rate=100.0, burst=20, initial_tokens=10)

        assert bucket.rate == 100.0
        assert bucket.burst == 20
        assert bucket.tokens == 10

    def test_consume_tokens(self):
        """Test consuming tokens from the bucket."""
        bucket = TokenBucket(rate=10.0, burst=5)

        # Consume 1 token
        assert bucket.consume(1) is True
        assert int(bucket.tokens) == 4

        # Consume 2 more tokens
        assert bucket.consume(2) is True
        assert int(bucket.tokens) == 2

    def test_consume_exceeds_available(self):
        """Test consuming more tokens than available."""
        bucket = TokenBucket(rate=10.0, burst=5)

        # Try to consume more than available
        assert bucket.consume(10) is False
        assert bucket.tokens == 5  # Tokens unchanged

    def test_consume_zero_tokens(self):
        """Test consuming zero tokens always succeeds."""
        bucket = TokenBucket(rate=10.0, burst=5)
        bucket.tokens = 0

        assert bucket.consume(0) is True
        # Tokens may have refilled slightly due to time elapsed
        assert bucket.tokens < 0.1  # Should be close to zero

    def test_token_refill(self):
        """Test that tokens refill over time."""
        bucket = TokenBucket(rate=10.0, burst=5)
        bucket.tokens = 0  # Empty the bucket

        # Wait for refill (0.1s at 10 tokens/sec = 1 token)
        time.sleep(0.1)
        bucket._refill()

        # Should have approximately 1 token (with some tolerance)
        assert 0.9 <= bucket.tokens <= 1.1

    def test_token_refill_capped_at_burst(self):
        """Test that tokens don't exceed burst limit."""
        bucket = TokenBucket(rate=10.0, burst=5)

        # Wait and refill multiple times
        time.sleep(0.5)
        bucket._refill()

        # Should be capped at burst
        assert bucket.tokens == 5

    def test_get_stats(self):
        """Test getting bucket statistics."""
        bucket = TokenBucket(rate=10.0, burst=5)
        bucket.consume(2)

        stats = bucket.get_stats()

        assert stats["tokens"] == 3
        assert stats["burst"] == 5
        assert stats["rate"] == 10.0
        assert stats["total_requests"] == 1
        assert stats["blocked_requests"] == 0

    def test_reset(self):
        """Test resetting the bucket."""
        bucket = TokenBucket(rate=10.0, burst=5)
        bucket.consume(3)
        bucket.total_requests = 100
        bucket.blocked_requests = 50

        bucket.reset()

        assert bucket.tokens == 5  # Reset to burst
        assert bucket.total_requests == 0
        assert bucket.blocked_requests == 0


class TestRateLimitConfig:
    """Test RateLimitConfig dataclass."""

    def test_default_creation(self):
        """Test creating config with defaults."""
        config = RateLimitConfig()

        assert config.rate == 10.0
        assert config.burst == 5
        assert config.per_job_type == {}
        assert config.enabled is True
        assert config.default_behavior == "block"

    def test_custom_creation(self):
        """Test creating config with custom values."""
        config = RateLimitConfig(
            rate=100.0,
            burst=20,
            per_job_type={"email": RateLimitConfig(rate=5.0, burst=2)},
            enabled=True,
            default_behavior="drop",
        )

        assert config.rate == 100.0
        assert config.burst == 20
        assert "email" in config.per_job_type
        assert config.per_job_type["email"].rate == 5.0

    def test_to_dict(self):
        """Test converting config to dictionary."""
        config = RateLimitConfig(rate=10.0, burst=5, default_behavior="defer")

        data = config.to_dict()

        assert data["rate"] == 10.0
        assert data["burst"] == 5
        assert data["enabled"] is True
        assert data["default_behavior"] == "defer"

    def test_from_dict(self):
        """Test creating config from dictionary."""
        data = {
            "rate": 50.0,
            "burst": 10,
            "enabled": False,
            "default_behavior": "drop",
            "per_job_type": {
                "sms": {"rate": 5.0, "burst": 2},
            },
        }

        config = RateLimitConfig.from_dict(data)

        assert config.rate == 50.0
        assert config.burst == 10
        assert config.enabled is False
        assert config.default_behavior == "drop"
        assert config.per_job_type["sms"].rate == 5.0


class TestRateLimitInfo:
    """Test RateLimitInfo dataclass."""

    def test_default_creation(self):
        """Test creating info with defaults."""
        info = RateLimitInfo()

        assert info.allowed is True
        assert info.tokens_remaining == 0
        assert info.tokens_total == 0
        assert info.reset_after == 0.0
        assert info.retry_after == 0.0
        assert info.job_type == "default"

    def test_custom_creation(self):
        """Test creating info with custom values."""
        info = RateLimitInfo(
            allowed=False,
            tokens_remaining=2,
            tokens_total=10,
            reset_after=0.5,
            retry_after=0.1,
            job_type="email",
        )

        assert info.allowed is False
        assert info.tokens_remaining == 2
        assert info.tokens_total == 10
        assert info.reset_after == 0.5
        assert info.retry_after == 0.1
        assert info.job_type == "email"

    def test_to_headers(self):
        """Test converting info to HTTP-style headers."""
        info = RateLimitInfo(
            allowed=True,
            tokens_remaining=8,
            tokens_total=10,
            reset_after=0.5,
            retry_after=0.0,
            job_type="email",
        )

        headers = info.to_headers()

        assert headers["X-RateLimit-Limit"] == "10"
        assert headers["X-RateLimit-Remaining"] == "8"
        assert headers["X-RateLimit-Reset"] == "0.500"
        assert headers["X-RateLimit-Retry-After"] == "0.000"
        assert headers["X-RateLimit-JobType"] == "email"


class TestRateLimiterBasic:
    """Test basic RateLimiter functionality."""

    @pytest.fixture
    def limiter(self):
        """Create a rate limiter for testing."""
        config = RateLimitConfig(rate=10.0, burst=5)
        return RateLimiter(config)

    def test_creation(self):
        """Test creating a rate limiter."""
        config = RateLimitConfig(rate=10.0, burst=5)
        limiter = RateLimiter(config)

        assert limiter.config.rate == 10.0
        assert limiter.config.burst == 5
        assert "default" in limiter._buckets

    def test_get_bucket_for_job_type(self, limiter):
        """Test getting bucket for job type."""
        # Unknown job types get their own bucket (same config as default)
        bucket = limiter._get_bucket("unknown")
        assert bucket is not None
        assert bucket.rate == limiter.config.rate
        assert bucket.burst == limiter.config.burst
        # Unknown job types are stored separately
        assert "unknown" in limiter._buckets

    def test_get_bucket_per_job_type(self):
        """Test getting bucket with per-job-type config."""
        config = RateLimitConfig(
            rate=10.0,
            burst=5,
            per_job_type={"email": RateLimitConfig(rate=5.0, burst=2)},
        )
        limiter = RateLimiter(config)

        # Should have separate bucket for email
        email_bucket = limiter._get_bucket("email")
        default_bucket = limiter._get_bucket("default")

        assert email_bucket.rate == 5.0
        assert email_bucket.burst == 2
        assert default_bucket.rate == 10.0
        assert default_bucket.burst == 5

    def test_try_acquire_success(self, limiter):
        """Test successful token acquisition."""
        info = limiter.try_acquire("default")

        assert info.allowed is True
        assert info.tokens_remaining == 4  # 5 - 1
        assert info.tokens_total == 5

    def test_try_acquire_failure(self, limiter):
        """Test failed token acquisition."""
        # Consume all tokens
        for _ in range(5):
            limiter.try_acquire("default")

        # Next acquisition should fail
        info = limiter.try_acquire("default")

        assert info.allowed is False
        assert info.tokens_remaining == 0
        assert info.retry_after > 0

    def test_try_acquire_disabled(self):
        """Test acquisition when rate limiting is disabled."""
        config = RateLimitConfig(enabled=False)
        limiter = RateLimiter(config)

        info = limiter.try_acquire("default")

        assert info.allowed is True
        assert info.tokens_remaining == -1  # Unlimited


class TestRateLimiterAcquire:
    """Test RateLimiter.acquire() blocking method."""

    def test_acquire_immediate(self):
        """Test acquire when tokens available."""
        config = RateLimitConfig(rate=10.0, burst=5)
        limiter = RateLimiter(config)

        info = limiter.acquire("default")

        assert info.allowed is True
        assert info.tokens_remaining == 4

    def test_acquire_blocks(self):
        """Test acquire blocks until token available."""
        config = RateLimitConfig(rate=100.0, burst=1)  # 1 token, fast refill
        limiter = RateLimiter(config)

        # Consume the only token
        limiter.try_acquire("default")

        # Acquire should block briefly then succeed
        start = time.time()
        info = limiter.acquire("default", timeout=2.0)
        elapsed = time.time() - start

        assert info.allowed is True
        assert elapsed >= 0.01  # Should have waited at least a bit

    def test_acquire_timeout(self):
        """Test acquire times out when no tokens available."""
        config = RateLimitConfig(rate=0.1, burst=1)  # Very slow refill
        limiter = RateLimiter(config)

        # Consume the only token
        limiter.try_acquire("default")

        # Acquire should timeout
        info = limiter.acquire("default", timeout=0.1)

        assert info.allowed is False

    def test_acquire_non_blocking(self):
        """Test acquire with no blocking."""
        config = RateLimitConfig(rate=10.0, burst=1)
        limiter = RateLimiter(config)

        # Consume the only token
        limiter.try_acquire("default")

        # Non-blocking acquire should fail immediately
        info = limiter.acquire("default", blocking=False)

        assert info.allowed is False


class TestRateLimiterGetQuota:
    """Test RateLimiter.get_quota() method."""

    def test_get_quota_available(self):
        """Test getting quota when tokens available."""
        config = RateLimitConfig(rate=10.0, burst=5)
        limiter = RateLimiter(config)

        info = limiter.get_quota("default")

        assert info.allowed is True
        assert info.tokens_remaining == 5
        assert info.tokens_total == 5

    def test_get_quota_does_not_consume(self):
        """Test get_quota doesn't consume tokens."""
        config = RateLimitConfig(rate=10.0, burst=5)
        limiter = RateLimiter(config)

        # Get quota multiple times
        limiter.get_quota("default")
        limiter.get_quota("default")
        info = limiter.get_quota("default")

        # Should still have full bucket
        assert info.tokens_remaining == 5


class TestRateLimiterConcurrency:
    """Test thread safety of rate limiter."""

    def test_concurrent_acquire(self):
        """Test concurrent token acquisition."""
        config = RateLimitConfig(rate=1000.0, burst=100)
        limiter = RateLimiter(config)

        results = []
        errors = []

        def acquire_task():
            try:
                info = limiter.try_acquire("default")
                results.append(info.allowed)
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=acquire_task) for _ in range(50)]

        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0
        # All 50 should succeed (burst is 100)
        assert sum(results) == 50

    def test_concurrent_acquire_and_refill(self):
        """Test concurrent acquisition with refill."""
        config = RateLimitConfig(rate=100.0, burst=10)
        limiter = RateLimiter(config)

        results = []
        errors = []

        def acquire_task():
            try:
                # Try to acquire with short timeout
                info = limiter.acquire("default", timeout=0.5)
                results.append(info.allowed)
            except Exception as e:
                errors.append(e)

        # Start more threads than burst
        threads = [threading.Thread(target=acquire_task) for _ in range(20)]

        for t in threads:
            t.start()
            time.sleep(0.01)  # Stagger starts slightly

        for t in threads:
            t.join()

        assert len(errors) == 0
        # Some should succeed due to refill during execution
        assert sum(results) >= 10


class TestRateLimiterReset:
    """Test RateLimiter reset functionality."""

    def test_reset_all_buckets(self):
        """Test resetting all buckets."""
        config = RateLimitConfig(
            rate=10.0,
            burst=5,
            per_job_type={"email": RateLimitConfig(rate=5.0, burst=2)},
        )
        limiter = RateLimiter(config)

        # Consume some tokens
        limiter.try_acquire("default")
        limiter.try_acquire("email")

        # Reset all
        limiter.reset()

        # Check both buckets are reset
        default_info = limiter.get_quota("default")
        email_info = limiter.get_quota("email")

        assert default_info.tokens_remaining == 5
        assert email_info.tokens_remaining == 2

    def test_reset_specific_bucket(self):
        """Test resetting a specific bucket."""
        config = RateLimitConfig(
            rate=10.0,
            burst=5,
            per_job_type={"email": RateLimitConfig(rate=5.0, burst=2)},
        )
        limiter = RateLimiter(config)

        # Consume tokens from both
        limiter.try_acquire("default")
        limiter.try_acquire("email")

        # Reset only email bucket
        limiter.reset("email")

        default_info = limiter.get_quota("default")
        email_info = limiter.get_quota("email")

        assert default_info.tokens_remaining == 4  # Still consumed
        assert email_info.tokens_remaining == 2  # Reset


class TestRateLimiterStats:
    """Test RateLimiter statistics."""

    def test_get_stats(self):
        """Test getting overall statistics."""
        config = RateLimitConfig(
            rate=10.0,
            burst=5,
            per_job_type={"email": RateLimitConfig(rate=5.0, burst=2)},
        )
        limiter = RateLimiter(config)

        # Make some requests
        limiter.try_acquire("default")
        limiter.try_acquire("default")
        limiter.try_acquire("email")

        stats = limiter.get_stats()

        assert "default" in stats
        assert "email" in stats
        assert stats["default"]["total_requests"] == 2
        assert stats["email"]["total_requests"] == 1


class TestRateLimitError:
    """Test RateLimitError exception."""

    def test_exception_creation(self):
        """Test creating the exception."""
        info = RateLimitInfo(allowed=False, retry_after=1.5, job_type="email")
        exc = RateLimitError("Rate limit exceeded", info)

        assert str(exc) == "Rate limit exceeded"
        assert exc.info is info
        assert exc.info.retry_after == 1.5


class TestRateLimiterEdgeCases:
    """Test edge cases and special scenarios."""

    def test_zero_rate(self):
        """Test behavior with zero rate."""
        config = RateLimitConfig(rate=0.0, burst=5)
        limiter = RateLimiter(config)

        # Should be able to consume initial tokens
        info = limiter.try_acquire("default")
        assert info.allowed is True

    def test_large_burst(self):
        """Test with large burst value."""
        config = RateLimitConfig(rate=1.0, burst=10000)
        limiter = RateLimiter(config)

        info = limiter.get_quota("default")
        assert info.tokens_total == 10000

    def test_fractional_rate(self):
        """Test with fractional rate (slow refill)."""
        config = RateLimitConfig(rate=0.5, burst=2)  # 1 token every 2 seconds
        limiter = RateLimiter(config)

        info = limiter.get_quota("default")
        assert info.tokens_total == 2

    def test_negative_tokens_not_allowed(self):
        """Test that tokens don't go negative."""
        config = RateLimitConfig(rate=10.0, burst=5)
        limiter = RateLimiter(config)

        # Try to consume when empty
        for _ in range(10):
            limiter.try_acquire("default")

        bucket = limiter._get_bucket("default")
        assert bucket.tokens >= 0

    def test_job_type_isolation(self):
        """Test that different job types are isolated."""
        config = RateLimitConfig(
            rate=10.0,
            burst=5,
            per_job_type={
                "email": RateLimitConfig(rate=5.0, burst=2),
                "sms": RateLimitConfig(rate=3.0, burst=1),
            },
        )
        limiter = RateLimiter(config)

        # Consume all email tokens
        limiter.try_acquire("email")
        limiter.try_acquire("email")

        # SMS should still have its token
        info = limiter.try_acquire("sms")
        assert info.allowed is True

        # Default should still have its tokens
        info = limiter.try_acquire("default")
        assert info.allowed is True
