"""Tests for load-aware scheduler and load metrics (Task 2.7)."""

from __future__ import annotations

from sagellm_kv_cache.load_aware_scheduler import LoadAwareScheduler, PlacementStrategy
from sagellm_kv_cache.load_metrics import DeviceLoadMetrics, LoadMetricsCollector
from sagellm_kv_cache.models import RequestType, SchedulerRequest


class TestLoadMetrics:
    """Test suite for load metrics."""

    def test_device_load_metrics_basic(self) -> None:
        metrics = DeviceLoadMetrics.create(
            "cuda:0", kv_tokens_capacity=1000, memory_capacity_mb=100
        )
        metrics.update_kv_usage(300)
        metrics.update_request_count(active=2, queued=1)
        metrics.update_memory_usage(20)

        assert metrics.kv_utilization() == 0.3
        assert metrics.memory_utilization() == 0.2
        assert metrics.total_requests() == 3
        assert metrics.available_kv_tokens() == 700

    def test_load_metrics_collector(self) -> None:
        collector = LoadMetricsCollector()
        collector.register_device("cuda:0", kv_tokens_capacity=1000)
        collector.register_device("cuda:1", kv_tokens_capacity=1000)
        collector.update_kv_usage("cuda:0", 500)

        assert set(collector.list_devices()) == {"cuda:0", "cuda:1"}
        assert collector.get_metrics("cuda:0").kv_tokens_used == 500


class TestLoadAwareScheduler:
    """Test suite for load-aware placement strategies."""

    def _collector(self) -> LoadMetricsCollector:
        collector = LoadMetricsCollector()
        collector.register_device("cuda:0", kv_tokens_capacity=1000)
        collector.register_device("cuda:1", kv_tokens_capacity=1000)
        return collector

    def test_lowest_kv_strategy(self) -> None:
        collector = self._collector()
        collector.update_kv_usage("cuda:0", 700)
        collector.update_kv_usage("cuda:1", 200)

        scheduler = LoadAwareScheduler(collector, strategy=PlacementStrategy.LOWEST_KV)
        selected = scheduler.select_device(["cuda:0", "cuda:1"], kv_tokens_required=100)

        assert selected == "cuda:1"

    def test_shortest_queue_strategy(self) -> None:
        collector = self._collector()
        collector.update_request_count("cuda:0", active=4, queued=2)
        collector.update_request_count("cuda:1", active=1, queued=1)

        scheduler = LoadAwareScheduler(collector, strategy=PlacementStrategy.SHORTEST_QUEUE)
        selected = scheduler.select_device(["cuda:0", "cuda:1"], kv_tokens_required=100)

        assert selected == "cuda:1"

    def test_disable_load_awareness_fallback_round_robin(self) -> None:
        collector = self._collector()
        scheduler = LoadAwareScheduler(
            collector,
            strategy=PlacementStrategy.LOWEST_KV,
            enable_load_awareness=False,
        )

        first = scheduler.select_device(["cuda:0", "cuda:1"])
        second = scheduler.select_device(["cuda:0", "cuda:1"])

        assert (first, second) == ("cuda:0", "cuda:1")

    def test_request_device_affinity_is_honored(self) -> None:
        collector = self._collector()
        scheduler = LoadAwareScheduler(collector, strategy=PlacementStrategy.LOWEST_KV)

        request = SchedulerRequest(
            request_id="req-affinity",
            trace_id="trace-affinity",
            request_type=RequestType.PREFILL,
            prompt_len=16,
            max_new_tokens=8,
            kv_tokens_required=24,
            device_affinity="cuda:1",
        )

        selected = scheduler.select_device_for_request(request)
        assert selected == "cuda:1"
