"""Tests for load manager — rate limiting, concurrency, and SLO tracking."""

from __future__ import annotations

import threading
import time

from orxaq.orchestration.load_manager import (
    ConcurrencyLimiter,
    LoadConfig,
    LoadManager,
    LoadStats,
    RateLimiter,
    RequestMetrics,
)

# ── LoadConfig ───────────────────────────────────────────


class TestLoadConfig:
    def test_defaults(self):
        c = LoadConfig()
        assert c.max_qps == 50.0
        assert c.max_concurrent == 10
        assert c.slo_p99_ms == 5000.0

    def test_to_dict(self):
        c = LoadConfig(max_qps=100)
        d = c.to_dict()
        assert d["max_qps"] == 100


# ── RequestMetrics ───────────────────────────────────────


class TestRequestMetrics:
    def test_creation(self):
        m = RequestMetrics(
            provider="vllm",
            latency_ms=150.0,
            queued_ms=10.0,
            success=True,
        )
        assert m.provider == "vllm"
        assert m.latency_ms == 150.0
        assert m.success is True


# ── LoadStats ────────────────────────────────────────────


class TestLoadStats:
    def test_defaults(self):
        s = LoadStats()
        assert s.total_requests == 0
        assert s.latency_p99_ms == 0.0

    def test_to_dict(self):
        s = LoadStats(total_requests=100, latency_p99_ms=250.5)
        d = s.to_dict()
        assert d["total_requests"] == 100
        assert d["latency_p99_ms"] == 250.5


# ── RateLimiter ──────────────────────────────────────────


class TestRateLimiter:
    def test_acquire_within_rate(self):
        rl = RateLimiter(max_qps=1000.0)
        assert rl.acquire(timeout_s=0.1) is True

    def test_available_tokens(self):
        rl = RateLimiter(max_qps=100.0)
        tokens = rl.available_tokens
        assert tokens > 0

    def test_rate_limiting_works(self):
        rl = RateLimiter(max_qps=2.0)
        # Drain all tokens
        rl.acquire(timeout_s=0.1)
        rl.acquire(timeout_s=0.1)
        # Next should fail quickly (low timeout)
        assert rl.acquire(timeout_s=0.02) is False


# ── ConcurrencyLimiter ──────────────────────────────────


class TestConcurrencyLimiter:
    def test_acquire_release(self):
        cl = ConcurrencyLimiter(max_concurrent=2)
        assert cl.acquire(timeout_s=0.1) is True
        assert cl.active == 1
        assert cl.available == 1
        cl.release()
        assert cl.active == 0
        assert cl.available == 2

    def test_blocks_at_limit(self):
        cl = ConcurrencyLimiter(max_concurrent=1)
        assert cl.acquire(timeout_s=0.1) is True
        # Already at limit
        assert cl.acquire(timeout_s=0.02) is False
        cl.release()
        assert cl.acquire(timeout_s=0.1) is True
        cl.release()


# ── LoadManager ──────────────────────────────────────────


class TestLoadManagerAcquire:
    def test_acquire_success(self):
        mgr = LoadManager(LoadConfig(max_qps=1000, max_concurrent=10))
        assert mgr.acquire(timeout_s=0.5) is True
        mgr.release()

    def test_acquire_tracks_count(self):
        mgr = LoadManager(LoadConfig(max_qps=1000, max_concurrent=10))
        mgr.acquire(timeout_s=0.5)
        mgr.release()
        s = mgr.stats
        assert s.total_requests == 1


class TestLoadManagerMetrics:
    def test_record_success(self):
        mgr = LoadManager()
        mgr.record_success("vllm", latency_ms=100.0)
        mgr.record_success("vllm", latency_ms=200.0)
        s = mgr.stats
        assert s.total_completed == 2

    def test_record_error(self):
        mgr = LoadManager()
        mgr.record_error("openai", latency_ms=500.0)
        s = mgr.stats
        assert s.total_errors == 1
        assert s.total_completed == 1

    def test_percentiles(self):
        mgr = LoadManager()
        for i in range(100):
            mgr.record_success("vllm", latency_ms=float(i))
        s = mgr.stats
        assert s.latency_p50_ms > 0
        assert s.latency_p95_ms > s.latency_p50_ms
        assert s.latency_p99_ms >= s.latency_p95_ms

    def test_slo_violations(self):
        mgr = LoadManager(LoadConfig(slo_p99_ms=100.0))
        mgr.record_success("vllm", latency_ms=50.0)
        mgr.record_success("vllm", latency_ms=200.0)  # violation
        mgr.record_success("vllm", latency_ms=150.0)  # violation
        s = mgr.stats
        assert s.slo_violations == 2

    def test_rps_tracking(self):
        mgr = LoadManager()
        mgr.record_success("vllm", latency_ms=10.0)
        s = mgr.stats
        assert s.requests_per_second > 0


class TestLoadManagerProviderStats:
    def test_per_provider(self):
        mgr = LoadManager()
        mgr.record_success("vllm", latency_ms=50.0)
        mgr.record_success("vllm", latency_ms=100.0)
        mgr.record_success("openai", latency_ms=200.0)
        mgr.record_error("openai")

        ps = mgr.provider_stats()
        assert ps["vllm"]["total"] == 2
        assert ps["vllm"]["errors"] == 0
        assert ps["openai"]["total"] == 2
        assert ps["openai"]["errors"] == 1


class TestLoadManagerSummary:
    def test_summary(self):
        mgr = LoadManager()
        mgr.record_success("vllm", latency_ms=100.0)
        s = mgr.summary()
        assert "config" in s
        assert "stats" in s
        assert "providers" in s
        assert "rate_limiter_tokens" in s
        assert "concurrency_available" in s


class TestLoadManagerReset:
    def test_reset(self):
        mgr = LoadManager()
        mgr.record_success("vllm", latency_ms=100.0)
        mgr.record_error("openai")
        mgr.reset()
        s = mgr.stats
        assert s.total_completed == 0
        assert s.total_errors == 0
        assert s.total_requests == 0


class TestLoadManagerThreaded:
    def test_concurrent_acquire_release(self):
        mgr = LoadManager(
            LoadConfig(
                max_qps=1000,
                max_concurrent=5,
            )
        )
        errors: list[str] = []

        def worker(i: int):
            try:
                if mgr.acquire(timeout_s=2.0):
                    mgr.record_success(f"p{i % 3}", latency_ms=float(i))
                    time.sleep(0.01)
                    mgr.release()
            except Exception as e:
                errors.append(str(e))

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(20)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0
        s = mgr.stats
        assert s.total_completed > 0


class TestLoadManagerMaxMetrics:
    def test_metrics_cap(self):
        mgr = LoadManager()
        mgr._max_metrics = 50
        for i in range(100):
            mgr.record_success("vllm", latency_ms=float(i))
        assert len(mgr._metrics) == 50


# ── Targeted coverage (lines 213-215, 218-220, 262) ──────


class TestLoadManagerRateLimitReject:
    """Cover lines 213-215: rate limiter rejects request."""

    def test_exhausted_rate_limiter_rejects(self):
        cfg = LoadConfig(max_qps=1.0, max_concurrent=100, queue_timeout_s=0.01)
        mgr = LoadManager(config=cfg)
        # Drain rate limiter
        mgr._rate_limiter._tokens = 0.0
        result = mgr.acquire(timeout_s=0.001)
        assert result is False
        assert mgr._total_rejected >= 1


class TestLoadManagerConcurrencyTimeout:
    """Cover lines 218-220: concurrency limiter timeout."""

    def test_concurrency_exhausted_times_out(self):
        cfg = LoadConfig(max_qps=1000.0, max_concurrent=1, queue_timeout_s=0.01)
        mgr = LoadManager(config=cfg)
        # Acquire the only slot
        assert mgr.acquire(timeout_s=1.0) is True
        # Second acquire should timeout
        result = mgr.acquire(timeout_s=0.001)
        assert result is False
        assert mgr._total_timeouts >= 1
        mgr.release()


class TestLoadManagerErrorMetricsOverflow:
    """Cover line 262: metrics list overflow trimming."""

    def test_error_metrics_trimmed(self):
        mgr = LoadManager()
        mgr._max_metrics = 5
        for i in range(10):
            mgr.record_error("vllm", latency_ms=float(i))
        assert len(mgr._metrics) == 5
