"""Tests for Scheduler Bridge (Task 2.6/2.7)"""

from sagellm_kv_cache.load_aware_scheduler import PlacementStrategy
from sagellm_kv_cache.models import EvictionPolicy, RequestType, SchedulerRequest
from sagellm_kv_cache.scheduler_bridge import SchedulerBridge


class TestSchedulerBridge:
    """Test suite for SchedulerBridge."""

    def test_initialization(self):
        """Test basic initialization."""
        bridge = SchedulerBridge(max_tokens=1024, block_size=16)

        assert bridge.max_tokens == 1024
        assert bridge.block_size == 16
        assert bridge.kv_pool is not None
        assert bridge.prefix_cache is not None

    def test_allocate_request(self):
        """Test allocating KV cache for a request."""
        bridge = SchedulerBridge(max_tokens=1024)

        handle = bridge.allocate_request("req1", num_tokens=100)

        assert handle.num_tokens == 100
        assert "req1" in bridge._request_handles

        # Check pool stats
        stats = bridge.get_stats()
        assert stats["pool"]["allocated_tokens"] == 100
        assert stats["num_requests"] == 1

    def test_free_request(self):
        """Test freeing KV cache for a request."""
        bridge = SchedulerBridge(max_tokens=1024)

        bridge.allocate_request("req1", num_tokens=100)
        bridge.free_request("req1")

        assert "req1" not in bridge._request_handles
        stats = bridge.get_stats()
        assert stats["pool"]["allocated_tokens"] == 0

    def test_plan_single_request(self):
        """Test creating plan for single request."""
        bridge = SchedulerBridge(max_tokens=1024)

        req = SchedulerRequest(
            request_id="req1",
            trace_id="trace1",
            request_type=RequestType.PREFILL,
            prompt_len=100,
            max_new_tokens=50,
            kv_tokens_required=150,
        )

        plan = bridge.plan([req])

        assert plan.batch_size == 1
        assert plan.total_kv_tokens >= 150
        assert "req1" in plan.placement

    def test_plan_multiple_requests(self):
        """Test creating plan for multiple requests."""
        bridge = SchedulerBridge(max_tokens=2048)

        reqs = [
            SchedulerRequest(
                request_id=f"req{i}",
                trace_id=f"trace{i}",
                request_type=RequestType.PREFILL,
                prompt_len=50,
                max_new_tokens=50,
                kv_tokens_required=100,
            )
            for i in range(5)
        ]

        plan = bridge.plan(reqs)

        assert plan.batch_size == 5
        assert plan.total_kv_tokens >= 500

    def test_plan_exceeds_budget(self):
        """Test plan when budget is exceeded."""
        bridge = SchedulerBridge(max_tokens=200)

        # First request uses most of budget
        req1 = SchedulerRequest(
            request_id="req1",
            trace_id="trace1",
            request_type=RequestType.PREFILL,
            prompt_len=150,
            max_new_tokens=50,
            kv_tokens_required=180,
        )

        # Second request won't fit
        req2 = SchedulerRequest(
            request_id="req2",
            trace_id="trace2",
            request_type=RequestType.PREFILL,
            prompt_len=50,
            max_new_tokens=50,
            kv_tokens_required=100,
        )

        bridge.plan([req1, req2])

        # At least one should fail to allocate due to budget
        assert len(bridge._request_handles) < 2

    def test_eviction_triggered(self):
        """Test that eviction is triggered when needed."""
        bridge = SchedulerBridge(max_tokens=300, eviction_policy=EvictionPolicy.LRU)

        # Allocate three requests
        bridge.allocate_request("req1", 100)
        h2 = bridge.allocate_request("req2", 100)
        h3 = bridge.allocate_request("req3", 100)

        # Touch h2 and h3 to make them more recent
        h2.touch()
        h3.touch()

        # Now try to allocate req4 (needs eviction)
        bridge.allocate_request("req4", 150)

        # req1 should be evicted (oldest)
        assert "req1" not in bridge._request_handles
        assert "req4" in bridge._request_handles

    def test_eviction_respects_pinned(self):
        """Test that pinned handles are not evicted."""
        bridge = SchedulerBridge(max_tokens=300)

        h1 = bridge.allocate_request("req1", 100)
        bridge.allocate_request("req2", 100)
        bridge.allocate_request("req3", 100)

        # Pin h1
        bridge.kv_pool.pin(h1)

        # Try to allocate req4 (needs eviction)
        bridge.allocate_request("req4", 150)

        # h1 should NOT be evicted (pinned)
        assert "req1" in bridge._request_handles
        # h2 or h3 should be evicted
        assert "req4" in bridge._request_handles

    def test_reset(self):
        """Test resetting bridge state."""
        bridge = SchedulerBridge(max_tokens=1024)

        bridge.allocate_request("req1", 100)
        bridge.allocate_request("req2", 100)

        assert len(bridge._request_handles) == 2

        bridge.reset()

        assert len(bridge._request_handles) == 0
        stats = bridge.get_stats()
        assert stats["pool"]["allocated_tokens"] == 0

    def test_prefix_cache_disabled(self):
        """Test bridge with prefix cache disabled."""
        bridge = SchedulerBridge(max_tokens=1024, enable_prefix_cache=False)

        assert bridge.prefix_cache is None
        stats = bridge.get_stats()
        assert "prefix_cache" not in stats

    def test_device_affinity(self):
        """Test device affinity in plan."""
        bridge = SchedulerBridge(max_tokens=1024)

        req = SchedulerRequest(
            request_id="req1",
            trace_id="trace1",
            request_type=RequestType.PREFILL,
            prompt_len=100,
            max_new_tokens=50,
            kv_tokens_required=150,
            device_affinity="cuda:1",
        )

        plan = bridge.plan([req])

        assert plan.placement["req1"] == "cuda:1"

    def test_repr(self):
        """Test string representation."""
        bridge = SchedulerBridge(max_tokens=1024, block_size=16)
        repr_str = repr(bridge)

        assert "SchedulerBridge" in repr_str
        assert "max_tokens=1024" in repr_str
        assert "block_size=16" in repr_str

    def test_load_aware_placement_lowest_kv(self):
        """Test load-aware placement picks lower KV occupancy device."""
        bridge = SchedulerBridge(
            max_tokens=2000,
            enable_load_awareness=True,
            placement_strategy=PlacementStrategy.LOWEST_KV,
            devices=["cuda:0", "cuda:1"],
        )

        bridge.allocate_request("seed0", num_tokens=900, device="cuda:0")
        bridge.allocate_request("seed1", num_tokens=100, device="cuda:1")

        req = SchedulerRequest(
            request_id="req-lowest-kv",
            trace_id="trace-lowest-kv",
            request_type=RequestType.PREFILL,
            prompt_len=32,
            max_new_tokens=16,
            kv_tokens_required=48,
        )
        plan = bridge.plan([req])

        assert plan.placement["req-lowest-kv"] == "cuda:1"

    def test_load_metrics_exposed_in_stats(self):
        """Test per-device metrics are visible in bridge stats."""
        bridge = SchedulerBridge(
            max_tokens=1024,
            enable_load_awareness=True,
            devices=["cuda:0", "cuda:1"],
        )
        bridge.allocate_request("req-metrics", num_tokens=128, device="cuda:1")

        stats = bridge.get_stats()

        assert "load_metrics" in stats
        assert "placement" in stats
        assert "cuda:1" in stats["load_metrics"]
        assert stats["load_metrics"]["cuda:1"]["active_requests"] >= 1
