"""Tests for rbsp.runtime.throttle."""

from __future__ import annotations

import threading
import time

import pytest

from rbsp.runtime.throttle import (
    BackpressureController,
    SlidingWindowCounter,
    Throttle,
    TokenBucket,
)


# ---------------------------------------------------------------------------
# TokenBucket
# ---------------------------------------------------------------------------

class TestTokenBucket:
    def test_initial_capacity(self):
        bucket = TokenBucket(rate=10, capacity=5)
        assert bucket.available() == 5

    def test_try_acquire_success(self):
        bucket = TokenBucket(rate=10, capacity=5)
        assert bucket.try_acquire(3) is True
        assert bucket.available() == 2

    def test_try_acquire_failure(self):
        bucket = TokenBucket(rate=10, capacity=5)
        assert bucket.try_acquire(6) is False
        assert bucket.available() == 5

    def test_try_acquire_drains_bucket(self):
        bucket = TokenBucket(rate=10, capacity=3)
        assert bucket.try_acquire() is True
        assert bucket.try_acquire() is True
        assert bucket.try_acquire() is True
        assert bucket.try_acquire() is False

    def test_acquire_with_timeout_success(self):
        bucket = TokenBucket(rate=1000, capacity=5)
        assert bucket.acquire(tokens=1, timeout=1.0) is True

    def test_acquire_with_timeout_failure(self):
        bucket = TokenBucket(rate=0.1, capacity=1)
        bucket.try_acquire()  # drain
        assert bucket.acquire(tokens=1, timeout=0.05) is False

    def test_refill_over_time(self):
        bucket = TokenBucket(rate=100, capacity=10)
        bucket.try_acquire(10)  # drain all
        assert bucket.available() == 0
        time.sleep(0.06)
        assert bucket.available() >= 1

    def test_capacity_is_ceiling(self):
        bucket = TokenBucket(rate=100, capacity=5)
        time.sleep(0.1)  # would add 10 tokens, but cap is 5
        assert bucket.available() <= 5

    def test_acquire_blocks_then_succeeds(self):
        bucket = TokenBucket(rate=200, capacity=1)
        bucket.try_acquire()  # drain
        # rate=200 => 1 token every 5ms, so should succeed well within 0.5s
        assert bucket.acquire(tokens=1, timeout=0.5) is True

    def test_thread_safety(self):
        # Use a very low refill rate so tokens don't replenish during the test
        bucket = TokenBucket(rate=0.001, capacity=100)
        results = []

        def worker():
            count = 0
            for _ in range(50):
                if bucket.try_acquire():
                    count += 1
            results.append(count)

        threads = [threading.Thread(target=worker) for _ in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        # Total acquired should not exceed capacity (100)
        assert sum(results) <= 100


# ---------------------------------------------------------------------------
# SlidingWindowCounter
# ---------------------------------------------------------------------------

class TestSlidingWindowCounter:
    def test_record_and_count(self):
        counter = SlidingWindowCounter(window_seconds=10.0, max_count=100)
        for _ in range(5):
            assert counter.record() is True
        assert counter.count() >= 5

    def test_max_count_enforced(self):
        counter = SlidingWindowCounter(window_seconds=60.0, max_count=3)
        assert counter.record() is True
        assert counter.record() is True
        assert counter.record() is True
        assert counter.record() is False

    def test_is_exceeded(self):
        counter = SlidingWindowCounter(window_seconds=60.0, max_count=2)
        assert counter.is_exceeded() is False
        counter.record()
        counter.record()
        assert counter.is_exceeded() is True

    def test_initial_count_is_zero(self):
        counter = SlidingWindowCounter(window_seconds=10.0, max_count=100)
        assert counter.count() == 0

    def test_count_does_not_go_negative(self):
        counter = SlidingWindowCounter(window_seconds=0.01, max_count=100)
        counter.record()
        time.sleep(0.05)
        assert counter.count() >= 0

    def test_thread_safety(self):
        counter = SlidingWindowCounter(window_seconds=60.0, max_count=1000)
        barrier = threading.Barrier(4)

        def worker():
            barrier.wait()
            for _ in range(50):
                counter.record()

        threads = [threading.Thread(target=worker) for _ in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert counter.count() <= 200  # 4 * 50


# ---------------------------------------------------------------------------
# BackpressureController
# ---------------------------------------------------------------------------

class TestBackpressureController:
    def test_acquire_and_release(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=2)
        assert bp.acquire() is True
        assert bp.acquire() is True
        assert bp.acquire() is False  # at capacity
        bp.release(50.0)
        assert bp.acquire() is True

    def test_current_concurrency(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=5)
        assert bp.current_concurrency() == 0
        bp.acquire()
        bp.acquire()
        assert bp.current_concurrency() == 2
        bp.release(10.0)
        assert bp.current_concurrency() == 1

    def test_should_shed_load_when_full(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=1)
        bp.acquire()
        assert bp.should_shed_load() is True

    def test_should_not_shed_load_when_idle(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=5)
        assert bp.should_shed_load() is False

    def test_should_shed_load_high_latency(self):
        bp = BackpressureController(target_latency_ms=10.0, max_concurrent=10)
        # Report high latencies to trigger shedding
        bp.acquire()
        bp.release(100.0)
        bp.acquire()
        bp.release(100.0)
        bp.acquire()
        bp.release(100.0)
        assert bp.should_shed_load() is True

    def test_concurrency_does_not_go_negative(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=2)
        bp.release(10.0)  # release without acquire
        assert bp.current_concurrency() == 0

    def test_adaptive_decrease(self):
        bp = BackpressureController(target_latency_ms=10.0, max_concurrent=8)
        # Force the adjust interval to pass
        bp._last_adjust = time.monotonic() - 2.0
        # Report latencies way above target (>1.5x)
        for _ in range(5):
            bp.acquire()
            bp.release(100.0)
        # max_concurrent should have decreased
        assert bp._max_concurrent < 8

    def test_adaptive_increase(self):
        bp = BackpressureController(target_latency_ms=100.0, max_concurrent=2)
        bp._ceiling = 10
        # Force the adjust interval to pass
        bp._last_adjust = time.monotonic() - 2.0
        # Report latencies well below target (<0.5x)
        for _ in range(5):
            bp.acquire()
            bp.release(1.0)
        # max_concurrent should have increased
        assert bp._max_concurrent > 2


# ---------------------------------------------------------------------------
# Throttle (combined)
# ---------------------------------------------------------------------------

class TestThrottle:
    def test_allow_query(self):
        throttle = Throttle(queries_per_second=100, max_concurrent_queries=5)
        assert throttle.allow_query() is True

    def test_allow_index(self):
        throttle = Throttle(max_concurrent_indexing=2)
        assert throttle.allow_index() is True

    def test_query_concurrency_limit(self):
        throttle = Throttle(
            queries_per_second=1000,
            max_concurrent_queries=2,
            max_concurrent_indexing=2,
        )
        assert throttle.allow_query() is True
        assert throttle.allow_query() is True
        assert throttle.allow_query() is False  # backpressure limit
        throttle.release_query(10.0)
        assert throttle.allow_query() is True

    def test_index_concurrency_limit(self):
        throttle = Throttle(max_concurrent_indexing=1)
        assert throttle.allow_index() is True
        assert throttle.allow_index() is False
        throttle.release_index(10.0)
        assert throttle.allow_index() is True

    def test_stats_keys(self):
        throttle = Throttle()
        stats = throttle.stats()
        expected_keys = {
            "query_tokens_available",
            "query_concurrency",
            "query_window_count",
            "query_overloaded",
            "index_concurrency",
            "index_window_count",
            "index_overloaded",
        }
        assert set(stats.keys()) == expected_keys

    def test_stats_initial_values(self):
        throttle = Throttle(queries_per_second=50, max_concurrent_queries=5)
        stats = throttle.stats()
        assert stats["query_concurrency"] == 0
        assert stats["index_concurrency"] == 0
        assert stats["query_overloaded"] is False
        assert stats["index_overloaded"] is False
        assert stats["query_window_count"] == 0
        assert stats["index_window_count"] == 0

    def test_stats_after_operations(self):
        throttle = Throttle(
            queries_per_second=1000,
            max_concurrent_queries=10,
            max_concurrent_indexing=5,
        )
        throttle.allow_query()
        throttle.allow_index()
        stats = throttle.stats()
        assert stats["query_concurrency"] == 1
        assert stats["index_concurrency"] == 1
        assert stats["query_window_count"] >= 1
        assert stats["index_window_count"] >= 1

    def test_release_query(self):
        throttle = Throttle(
            queries_per_second=1000, max_concurrent_queries=1
        )
        throttle.allow_query()
        assert throttle.stats()["query_concurrency"] == 1
        throttle.release_query(5.0)
        assert throttle.stats()["query_concurrency"] == 0

    def test_release_index(self):
        throttle = Throttle(max_concurrent_indexing=1)
        throttle.allow_index()
        assert throttle.stats()["index_concurrency"] == 1
        throttle.release_index(5.0)
        assert throttle.stats()["index_concurrency"] == 0

    def test_rate_limit_blocks_queries(self):
        throttle = Throttle(
            queries_per_second=2,
            max_concurrent_queries=100,
        )
        # Drain the token bucket (capacity = 2)
        assert throttle.allow_query() is True
        throttle.release_query(1.0)
        assert throttle.allow_query() is True
        throttle.release_query(1.0)
        # Third should be rate-limited
        assert throttle.allow_query() is False
