"""Tests for KV-aware scheduling policy."""

from __future__ import annotations

import pytest

from sagellm_control.policies.kv_aware import KVAwareSchedulingPolicy
from sagellm_control.types import (
    EngineInfo,
    EngineState,
    RequestMetadata,
    RequestPriority,
    RequestType,
)


@pytest.fixture
def engines() -> list[EngineInfo]:
    """Create test engines."""
    return [
        EngineInfo(
            engine_id="engine-001",
            model_id="Qwen2-7B",
            host="localhost",
            port=8001,
            state=EngineState.READY,
            active_requests=0,
        ),
        EngineInfo(
            engine_id="engine-002",
            model_id="Qwen2-7B",
            host="localhost",
            port=8002,
            state=EngineState.READY,
            active_requests=2,
        ),
        EngineInfo(
            engine_id="engine-003",
            model_id="Qwen2-7B",
            host="localhost",
            port=8003,
            state=EngineState.READY,
            active_requests=1,
        ),
    ]


@pytest.fixture
def test_request() -> RequestMetadata:
    """Create test request."""
    return RequestMetadata(
        request_id="req-001",
        trace_id="trace-001",
        model_id="Qwen2-7B",
        prompt="Hello, how are you?",
        prompt_tokens=100,
        max_tokens=50,
        priority=RequestPriority.NORMAL,
        request_type=RequestType.LLM_GENERATE,
    )


def test_kv_aware_policy_initialization() -> None:
    """Test KV-aware policy can be initialized."""
    policy = KVAwareSchedulingPolicy(
        scheduler_bridge=None,
        max_kv_tokens_per_engine=4096,
        kv_budget_threshold=0.8,
    )
    assert policy.name == "kv_aware"
    assert policy._max_kv_tokens_per_engine == 4096
    assert policy._kv_budget_threshold == 0.8


def test_kv_tokens_estimation(test_request: RequestMetadata) -> None:
    """Test KV token estimation."""
    policy = KVAwareSchedulingPolicy()

    # Test with prompt_tokens set
    estimated = policy._estimate_kv_tokens(test_request)
    assert estimated == 100 + 50  # prompt + max_tokens

    # Test without prompt_tokens (estimates from prompt length)
    test_request.prompt_tokens = 0
    estimated = policy._estimate_kv_tokens(test_request)
    assert estimated > 50  # At least max_tokens


def test_can_admit_request(engines: list[EngineInfo]) -> None:
    """Test request admission check."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    # Should admit small request
    assert policy._can_admit_request("engine-001", 100) is True

    # Should reject request exceeding budget
    assert policy._can_admit_request("engine-001", 1500) is False

    # Test threshold warning
    # Allocate up to 70% of budget
    policy._engine_kv_usage["engine-001"] = 700
    # This should trigger warning but still admit
    assert policy._can_admit_request("engine-001", 200) is True


def test_select_best_engine_basic(engines: list[EngineInfo]) -> None:
    """Test engine selection based on KV budget."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    # All engines have no KV usage, should select first
    selected = policy._select_best_engine(engines, kv_tokens_needed=100)
    assert selected is not None
    assert selected.engine_id == "engine-001"


def test_select_best_engine_with_usage(engines: list[EngineInfo]) -> None:
    """Test engine selection considers KV utilization."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    # Set different KV usage levels
    policy._engine_kv_usage["engine-001"] = 800  # 80% used
    policy._engine_kv_usage["engine-002"] = 200  # 20% used
    policy._engine_kv_usage["engine-003"] = 500  # 50% used

    # Should select engine-002 (lowest utilization)
    selected = policy._select_best_engine(engines, kv_tokens_needed=100)
    assert selected is not None
    assert selected.engine_id == "engine-002"


def test_select_best_engine_insufficient_budget(engines: list[EngineInfo]) -> None:
    """Test engine selection when no engine has sufficient budget."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    # Fill all engines to capacity
    for engine in engines:
        policy._engine_kv_usage[engine.engine_id] = 950

    # Request 200 tokens - no engine can admit
    selected = policy._select_best_engine(engines, kv_tokens_needed=200)
    assert selected is None


def test_schedule_request_success(
    test_request: RequestMetadata,
    engines: list[EngineInfo],
) -> None:
    """Test successful request scheduling."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    decision = policy.schedule(test_request, engines, [])

    assert decision.is_scheduled is True
    assert decision.engine_id == "engine-001"
    assert "KV-aware" in decision.reason


def test_schedule_request_no_capacity(
    test_request: RequestMetadata,
    engines: list[EngineInfo],
) -> None:
    """Test request scheduling when no capacity available."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=100)  # Very small budget

    decision = policy.schedule(test_request, engines, [])

    assert decision.is_scheduled is False
    assert decision.engine_id is None
    assert "No engine with sufficient KV budget" in decision.reason


def test_release_request() -> None:
    """Test releasing KV cache resources."""
    policy = KVAwareSchedulingPolicy()

    # Allocate some tokens
    policy._engine_kv_usage["engine-001"] = 500

    # Release tokens
    policy.release_request("req-001", "engine-001", 200)

    assert policy._engine_kv_usage["engine-001"] == 300

    # Release more than allocated (should not go negative)
    policy.release_request("req-002", "engine-001", 1000)
    assert policy._engine_kv_usage["engine-001"] == 0


def test_get_kv_utilization_metrics() -> None:
    """Test KV utilization metrics."""
    policy = KVAwareSchedulingPolicy(max_kv_tokens_per_engine=1000)

    policy._engine_kv_usage["engine-001"] = 500
    policy._engine_kv_usage["engine-002"] = 750

    metrics = policy.get_kv_utilization_metrics()

    assert metrics["engine-001"] == 0.5
    assert metrics["engine-002"] == 0.75


def test_reorder_queue() -> None:
    """Test queue reordering by KV token requirements."""
    policy = KVAwareSchedulingPolicy()

    # Create requests with different sizes
    requests = [
        RequestMetadata(
            request_id="req-001",
            prompt_tokens=500,
            max_tokens=100,
        ),
        RequestMetadata(
            request_id="req-002",
            prompt_tokens=100,
            max_tokens=50,
        ),
        RequestMetadata(
            request_id="req-003",
            prompt_tokens=300,
            max_tokens=75,
        ),
    ]

    reordered = policy.reorder_queue(requests)

    # Should be sorted by total KV tokens (ascending)
    assert reordered[0].request_id == "req-002"  # Smallest: 150
    assert reordered[1].request_id == "req-003"  # Medium: 375
    assert reordered[2].request_id == "req-001"  # Largest: 600


def test_kv_aware_policy_with_unhealthy_engines(
    test_request: RequestMetadata,
    engines: list[EngineInfo],
) -> None:
    """Test scheduling skips unhealthy engines."""
    policy = KVAwareSchedulingPolicy()

    # Mark first engine as unhealthy
    engines[0].state = EngineState.ERROR

    decision = policy.schedule(test_request, engines, [])

    # Should select engine-002 or engine-003, not engine-001
    assert decision.is_scheduled is True
    assert decision.engine_id != "engine-001"
