"""
tests/test_provider_rate_limiter.py — Tests for cortex.providers.rate_limiter

Coverage:
  Token bucket / burst control:
    1.  Allows requests under RPM limit (no throttle)
    2.  Throttles requests when RPM burst is exhausted (time advances in sleep)
    3.  Token-bucket throttles when TPM burst is exhausted
    4.  Refills tokens over time
  Per-provider isolation:
    5.  Two providers with different configs don't interfere
    6.  Unconfigured provider gets auto-defaults
  Backoff on 429:
    7.  record_429() sets backoff — next acquire waits
    8.  record_429() uses Retry-After when provided
    9.  Consecutive 429s produce increasing backoff (exponential)
    10. record_success() resets the consecutive-429 counter
  Metrics:
    11. requests_throttled increments on throttle
    12. backoffs_triggered increments on 429
    13. all_metrics() returns per-provider snapshot
  Config:
    14. from_config() reads provider_rate_limits dict
    15. configure_from_dict() works
    16. configure() with RateLimitConfig works
  Status / reset:
    17. status() includes backoff_remaining_s
    18. reset() removes provider state
    19. reset_all() clears everything
    20. RateLimitConfig defaults auto-fill burst values
"""
from __future__ import annotations

import asyncio
import time
from typing import List
from unittest.mock import patch

import pytest

from cortex.providers.rate_limiter import (
    ProviderRateLimiter,
    RateLimitConfig,
    RateLimitMetrics,
    _AsyncTokenBucket,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def run(coro):
    """Run a coroutine synchronously."""
    return asyncio.run(coro)


class FakeClock:
    """A fake monotonic clock that can be advanced manually."""

    def __init__(self, start: float = 0.0) -> None:
        self._t = start

    def __call__(self) -> float:
        return self._t

    def advance(self, seconds: float) -> None:
        self._t += seconds


def make_fake_sleep(clock: FakeClock, recorded: List[float]):
    """Return an async fake_sleep that advances the clock and records calls."""

    async def fake_sleep(seconds: float) -> None:
        clock.advance(seconds)
        recorded.append(seconds)

    return fake_sleep


# ---------------------------------------------------------------------------
# 1. Allows requests under RPM limit (no throttle)
# ---------------------------------------------------------------------------


def test_allows_under_rpm():
    limiter = ProviderRateLimiter()
    limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

    async def _run():
        for _ in range(5):
            async with limiter.acquire("openai"):
                pass

    run(_run())
    m = limiter.metrics("openai")
    assert m.requests_allowed == 5
    assert m.requests_throttled == 0


# ---------------------------------------------------------------------------
# 2. Throttles when RPM burst is exhausted
# ---------------------------------------------------------------------------


def test_throttles_when_burst_exhausted():
    """With burst=1, second request must throttle (sleep & wait)."""
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(requests_per_minute=60, tokens_per_minute=0, burst_requests=1),
        )

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            # First request drains burst=1
            async with limiter.acquire("openai"):
                pass
            # Second request: bucket empty → must sleep before consuming
            async with limiter.acquire("openai"):
                pass

    run(_run())
    assert len(sleep_calls) >= 1
    assert all(s >= 0 for s in sleep_calls)


# ---------------------------------------------------------------------------
# 3. Token-bucket throttles when TPM burst is exhausted
# ---------------------------------------------------------------------------


def test_throttles_when_tpm_exhausted():
    """Token bucket throttles when a call requests more than remaining tokens."""
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(
                requests_per_minute=0,       # no RPM limit
                tokens_per_minute=6000,      # 100 tok/sec refill
                burst_requests=1000,
                burst_tokens=100,
            ),
        )

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            # Drain burst_tokens=100 first
            async with limiter.acquire("openai", tokens=100):
                pass

            # Now bucket is empty; next call for 1 token must throttle
            async with limiter.acquire("openai", tokens=1):
                pass

    run(_run())
    assert len(sleep_calls) >= 1


# ---------------------------------------------------------------------------
# 4. Tokens refill over time
# ---------------------------------------------------------------------------


def test_tokens_refill_over_time():
    """Advancing the clock by ≥1s should allow a second call without extra sleep."""
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(
                requests_per_minute=60,  # 1 req/sec refill
                tokens_per_minute=0,
                burst_requests=1,
            ),
        )

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            # Drain burst
            async with limiter.acquire("openai"):
                pass

            # Advance clock by 1 full second
            clock.advance(1.0)
            pre = len(sleep_calls)

            # Should pass without extra sleeping
            async with limiter.acquire("openai"):
                pass

        assert len(sleep_calls) == pre  # no new sleeps after clock advance

    run(_run())


# ---------------------------------------------------------------------------
# 5. Per-provider isolation
# ---------------------------------------------------------------------------


def test_per_provider_isolation():
    """Draining provider A must not affect provider B."""
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(requests_per_minute=60, tokens_per_minute=0, burst_requests=1),
        )
        limiter.configure(
            "anthropic",
            RateLimitConfig(requests_per_minute=60, tokens_per_minute=0, burst_requests=5),
        )

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            # Drain openai (burst=1)
            async with limiter.acquire("openai"):
                pass
            async with limiter.acquire("openai"):
                pass  # triggers throttle

            openai_throttled = limiter.metrics("openai").requests_throttled

            # anthropic burst=5 — should not throttle at all
            for _ in range(5):
                async with limiter.acquire("anthropic"):
                    pass
            anthropic_throttled = limiter.metrics("anthropic").requests_throttled

        assert openai_throttled >= 1
        assert anthropic_throttled == 0

    run(_run())


# ---------------------------------------------------------------------------
# 6. Unconfigured provider gets auto-defaults
# ---------------------------------------------------------------------------


def test_unconfigured_provider_auto_defaults():
    async def _run():
        limiter = ProviderRateLimiter()
        # Don't configure "venice" — should auto-create from defaults
        async with limiter.acquire("venice"):
            pass
        m = limiter.metrics("venice")
        assert m.requests_allowed == 1

    run(_run())


# ---------------------------------------------------------------------------
# 7. record_429 triggers backoff sleep
# ---------------------------------------------------------------------------


def test_record_429_triggers_backoff():
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            # Record a 429 with explicit retry_after
            await limiter.record_429("openai", retry_after=5.0)
            # backoff_until is now set to clock() + 5.0 = 5.0
            # (clock is at 0.0 so backoff_until=5.0)

            # Next acquire should sleep for the backoff period
            async with limiter.acquire("openai"):
                pass

        assert len(sleep_calls) >= 1
        assert any(s > 0 for s in sleep_calls)

    run(_run())


# ---------------------------------------------------------------------------
# 8. record_429 uses Retry-After header value when provided
# ---------------------------------------------------------------------------


def test_record_429_uses_retry_after():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        clock = FakeClock(start=1000.0)
        with patch("time.monotonic", clock):
            await limiter.record_429("openai", retry_after=30.0)

        state = limiter._providers["openai"]
        # backoff_until should be approximately clock_at_call + 30
        assert abs(state.backoff_until - 1030.0) < 1.0
        assert state.consecutive_429s == 1

    run(_run())


# ---------------------------------------------------------------------------
# 9. Consecutive 429s produce increasing backoff
# ---------------------------------------------------------------------------


def test_consecutive_429s_exponential_backoff():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        backoffs: List[float] = []
        clock = FakeClock(start=0.0)

        with patch("time.monotonic", clock):
            for _ in range(5):
                await limiter.record_429("openai")
                state = limiter._providers["openai"]
                # Compute how long from now the backoff lasts
                backoffs.append(state.backoff_until - clock())
                # Advance clock past current backoff so next record_429 starts fresh
                clock.advance(state.backoff_until - clock() + 0.1)

        # Backoffs should grow (exponential); last >> first (with jitter tolerance)
        assert backoffs[-1] > backoffs[0], f"Expected exponential growth: {backoffs}"

    run(_run())


# ---------------------------------------------------------------------------
# 10. record_success resets consecutive-429 counter
# ---------------------------------------------------------------------------


def test_record_success_resets_counter():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        await limiter.record_429("openai")
        await limiter.record_429("openai")
        assert limiter._providers["openai"].consecutive_429s == 2

        limiter.record_success("openai")
        assert limiter._providers["openai"].consecutive_429s == 0

    run(_run())


# ---------------------------------------------------------------------------
# 11. requests_throttled increments on throttle
# ---------------------------------------------------------------------------


def test_throttled_metric_increments():
    clock = FakeClock()
    sleep_calls: List[float] = []

    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(requests_per_minute=60, tokens_per_minute=0, burst_requests=1),
        )

        with (
            patch("time.monotonic", clock),
            patch("asyncio.sleep", make_fake_sleep(clock, sleep_calls)),
        ):
            async with limiter.acquire("openai"):
                pass
            async with limiter.acquire("openai"):  # throttled
                pass

        m = limiter.metrics("openai")
        assert m.requests_throttled >= 1

    run(_run())


# ---------------------------------------------------------------------------
# 12. backoffs_triggered increments on 429
# ---------------------------------------------------------------------------


def test_backoffs_triggered_increments():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        await limiter.record_429("openai")
        await limiter.record_429("openai")
        m = limiter.metrics("openai")
        assert m.backoffs_triggered == 2

    run(_run())


# ---------------------------------------------------------------------------
# 13. all_metrics() returns per-provider snapshot
# ---------------------------------------------------------------------------


def test_all_metrics_returns_all_providers():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))
        limiter.configure("anthropic", RateLimitConfig(requests_per_minute=50, tokens_per_minute=0))

        async with limiter.acquire("openai"):
            pass
        async with limiter.acquire("anthropic"):
            pass

        all_m = limiter.all_metrics()
        assert "openai" in all_m
        assert "anthropic" in all_m
        assert all_m["openai"].requests_allowed == 1
        assert all_m["anthropic"].requests_allowed == 1

    run(_run())


# ---------------------------------------------------------------------------
# 14. from_config() reads provider_rate_limits dict
# ---------------------------------------------------------------------------


def test_from_config_reads_provider_rate_limits():
    class FakeConfig:
        provider_rate_limits = {
            "openai": {"requests_per_minute": 30, "tokens_per_minute": 45_000},
            "anthropic": {"requests_per_minute": 20, "tokens_per_minute": 60_000},
        }

    limiter = ProviderRateLimiter.from_config(FakeConfig())
    assert "openai" in limiter._providers
    assert "anthropic" in limiter._providers
    assert limiter._providers["openai"].config.requests_per_minute == 30
    assert limiter._providers["anthropic"].config.tokens_per_minute == 60_000


# ---------------------------------------------------------------------------
# 15. configure_from_dict works
# ---------------------------------------------------------------------------


def test_configure_from_dict():
    limiter = ProviderRateLimiter()
    limiter.configure_from_dict("venice", {"requests_per_minute": 10, "tokens_per_minute": 5000})
    assert limiter._providers["venice"].config.requests_per_minute == 10
    assert limiter._providers["venice"].config.tokens_per_minute == 5000


# ---------------------------------------------------------------------------
# 16. configure() with RateLimitConfig
# ---------------------------------------------------------------------------


def test_configure_with_rate_limit_config():
    limiter = ProviderRateLimiter()
    cfg = RateLimitConfig(requests_per_minute=25, tokens_per_minute=50_000)
    limiter.configure("bailian", cfg)
    assert limiter._providers["bailian"].config.requests_per_minute == 25


# ---------------------------------------------------------------------------
# 17. status() includes backoff_remaining_s
# ---------------------------------------------------------------------------


def test_status_includes_backoff_remaining():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        clock = FakeClock(start=100.0)
        with patch("time.monotonic", clock):
            await limiter.record_429("openai", retry_after=15.0)
            s = limiter.status()

        assert "openai" in s
        assert s["openai"]["backoff_remaining_s"] > 0

    run(_run())


# ---------------------------------------------------------------------------
# 18. reset() removes provider state
# ---------------------------------------------------------------------------


def test_reset_removes_provider():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        async with limiter.acquire("openai"):
            pass
        assert "openai" in limiter._providers

        limiter.reset("openai")
        assert "openai" not in limiter._providers

        # After reset, next acquire auto-creates from defaults
        async with limiter.acquire("openai"):
            pass
        assert limiter.metrics("openai").requests_allowed == 1

    run(_run())


# ---------------------------------------------------------------------------
# 19. reset_all() clears everything
# ---------------------------------------------------------------------------


def test_reset_all_clears_all():
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))
        limiter.configure("anthropic", RateLimitConfig(requests_per_minute=50, tokens_per_minute=0))

        async with limiter.acquire("openai"):
            pass
        async with limiter.acquire("anthropic"):
            pass

        limiter.reset_all()
        assert len(limiter._providers) == 0

    run(_run())


# ---------------------------------------------------------------------------
# 20. RateLimitConfig defaults auto-fill burst values
# ---------------------------------------------------------------------------


def test_rate_limit_config_auto_burst():
    cfg = RateLimitConfig(requests_per_minute=30, tokens_per_minute=45_000)
    assert cfg.burst_requests == 30
    assert cfg.burst_tokens == 45_000

    cfg2 = RateLimitConfig(requests_per_minute=30, tokens_per_minute=45_000,
                           burst_requests=5, burst_tokens=1000)
    assert cfg2.burst_requests == 5
    assert cfg2.burst_tokens == 1000


# ---------------------------------------------------------------------------
# B-023: Negative token amounts should not inflate bucket
# ---------------------------------------------------------------------------


def test_negative_token_amount_clamped_to_zero():
    """B-023: Passing a negative amount to consume() should not add tokens back."""
    clock = FakeClock(start=1000.0)
    with patch("time.monotonic", clock):
        bucket = _AsyncTokenBucket(capacity=10.0, refill_rate=1.0)
        # Consume 5 tokens
        ok, _ = bucket.try_consume(5.0)
        assert ok
        before = bucket.available()
        assert abs(before - 5.0) < 0.01

        # Negative consume should NOT increase tokens (clock frozen)
        ok2, _ = bucket.try_consume(-10.0)
        assert ok2  # 0.0 is always consumable
        after = bucket.available()
        assert after <= before + 0.01, f"Negative consume inflated bucket: {before} → {after}"


def test_negative_token_amount_in_acquire():
    """B-023: acquire() with negative tokens should not inflate bucket."""
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(requests_per_minute=0, tokens_per_minute=6000,
                            burst_requests=1000, burst_tokens=10),
        )
        # Drain token bucket
        async with limiter.acquire("openai", tokens=10):
            pass

        # Negative tokens should not refill bucket
        async with limiter.acquire("openai", tokens=-100):
            pass

        state = limiter._providers["openai"]
        # Token bucket should not have been inflated
        assert state.tok_bucket.available() <= 10.0  # at most capacity after refill

    run(_run())


# ---------------------------------------------------------------------------
# B-024: record_success() should clear active backoff
# ---------------------------------------------------------------------------


def test_record_success_clears_backoff():
    """B-024: After a 429 backoff, record_success() should clear backoff_until."""
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        clock = FakeClock(start=100.0)
        with patch("time.monotonic", clock):
            await limiter.record_429("openai", retry_after=30.0)
            state = limiter._providers["openai"]
            assert state.backoff_until > 0

            # Record success should clear backoff
            limiter.record_success("openai")
            assert state.backoff_until == 0.0
            assert state.consecutive_429s == 0

    run(_run())


def test_record_success_prevents_future_backoff_sleep():
    """B-024: After record_success(), next acquire should NOT sleep for old backoff."""
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure("openai", RateLimitConfig(requests_per_minute=60, tokens_per_minute=0))

        await limiter.record_429("openai", retry_after=30.0)
        state = limiter._providers["openai"]
        assert state.backoff_until > 0, "backoff should be set after 429"

        limiter.record_success("openai")
        assert state.backoff_until == 0.0, "backoff should be cleared after success"

        # Verify next acquire doesn't see backoff (backoff_until=0 < any time.monotonic())
        # This is the key invariant: backoff_until=0 means no wait
        now = time.monotonic()
        assert state.backoff_until <= now, "backoff_until should be in the past"

    run(_run())


# ---------------------------------------------------------------------------
# B-025: burst=1 when RPM=0 causes subtle throttling
# ---------------------------------------------------------------------------


def test_rpm_zero_does_not_throttle():
    """B-025: When RPM=0 (unlimited), burst should be large enough that
    multiple rapid requests don't throttle."""
    async def _run():
        limiter = ProviderRateLimiter()
        limiter.configure(
            "openai",
            RateLimitConfig(requests_per_minute=0, tokens_per_minute=0),
        )

        # With the fix, burst should be very large (not 1)
        state = limiter._providers["openai"]
        assert state.config.burst_requests >= 1000, \
            f"burst_requests should be large when RPM=0, got {state.config.burst_requests}"

        # 100 rapid requests should not throttle at all
        for _ in range(100):
            async with limiter.acquire("openai"):
                pass
        m = limiter.metrics("openai")
        assert m.requests_allowed == 100
        assert m.requests_throttled == 0

    run(_run())


def test_rpm_zero_burst_config():
    """B-025: RateLimitConfig with RPM=0 should set burst to a large value."""
    cfg = RateLimitConfig(requests_per_minute=0, tokens_per_minute=0)
    assert cfg.burst_requests >= 1000, f"Expected large burst, got {cfg.burst_requests}"
    assert cfg.burst_tokens >= 1000, f"Expected large burst, got {cfg.burst_tokens}"
