"""Tests for placement and load balancing functionality."""

from __future__ import annotations

import pytest

from sagellm_control import (
    EngineInfo,
    EngineState,
    LoadBalancer,
    RequestMetadata,
    RequestType,
)


@pytest.fixture
def sample_engines() -> list[EngineInfo]:
    """Create sample engines with different load characteristics."""
    return [
        EngineInfo(
            engine_id="engine-0",
            model_id="Qwen2-7B",
            host="localhost",
            port=8000,
            state=EngineState.READY,
            active_requests=10,
            queue_length=5,
            kv_cache_usage_bytes=5_000_000_000,
            kv_cache_capacity_bytes=10_000_000_000,
            gpu_utilization=0.5,
        ),
        EngineInfo(
            engine_id="engine-1",
            model_id="Qwen2-7B",
            host="localhost",
            port=8001,
            state=EngineState.READY,
            active_requests=20,
            queue_length=10,
            kv_cache_usage_bytes=8_000_000_000,
            kv_cache_capacity_bytes=10_000_000_000,
            gpu_utilization=0.8,
        ),
        EngineInfo(
            engine_id="engine-2",
            model_id="Qwen2-7B",
            host="localhost",
            port=8002,
            state=EngineState.READY,
            active_requests=5,
            queue_length=2,
            kv_cache_usage_bytes=3_000_000_000,
            kv_cache_capacity_bytes=10_000_000_000,
            gpu_utilization=0.3,
        ),
    ]


@pytest.fixture
def sample_request() -> RequestMetadata:
    """Create a sample request."""
    return RequestMetadata(
        request_id="test-req-1",
        model_id="Qwen2-7B",
        prompt="Test prompt",
        request_type=RequestType.LLM_CHAT,
    )


def test_engine_info_load_metrics(sample_engines: list[EngineInfo]) -> None:
    """Test EngineInfo load metric calculations."""
    engine = sample_engines[0]

    # Test KV cache utilization
    assert engine.kv_cache_utilization == 0.5

    # Test total load calculation
    assert 0.0 <= engine.total_load <= 1.0

    # Engine with higher load should have higher total_load
    assert sample_engines[1].total_load > sample_engines[0].total_load
    assert sample_engines[2].total_load < sample_engines[0].total_load


def test_engine_info_to_dict(sample_engines: list[EngineInfo]) -> None:
    """Test EngineInfo serialization includes new fields."""
    engine = sample_engines[0]
    data = engine.to_dict()

    assert "queue_length" in data
    assert "kv_cache_usage_bytes" in data
    assert "kv_cache_capacity_bytes" in data
    assert "kv_cache_utilization" in data
    assert "gpu_utilization" in data
    assert "total_load" in data


def test_load_aware_router(
    sample_engines: list[EngineInfo], sample_request: RequestMetadata
) -> None:
    """Test LoadAwareRouter selects engine with lowest total load."""
    balancer = LoadBalancer(strategy="load_aware")
    engine = balancer.route(sample_request, sample_engines)

    assert engine is not None
    # Should select engine-2 (lowest load)
    assert engine.engine_id == "engine-2"


def test_kv_cache_aware_router(
    sample_engines: list[EngineInfo],
    sample_request: RequestMetadata,
) -> None:
    """Test KVCacheAwareRouter selects engine with lowest KV cache usage."""
    balancer = LoadBalancer(strategy="kv_cache_aware")
    engine = balancer.route(sample_request, sample_engines)

    assert engine is not None
    # Should select engine-2 (lowest KV cache utilization: 0.3)
    assert engine.engine_id == "engine-2"


def test_least_loaded_router(
    sample_engines: list[EngineInfo],
    sample_request: RequestMetadata,
) -> None:
    """Test LeastLoadedRouter selects engine with fewest active requests."""
    balancer = LoadBalancer(strategy="least_loaded")
    engine = balancer.route(sample_request, sample_engines)

    assert engine is not None
    # Should select engine-2 (5 active requests)
    assert engine.engine_id == "engine-2"


def test_route_with_decision(
    sample_engines: list[EngineInfo],
    sample_request: RequestMetadata,
) -> None:
    """Test route_with_decision returns placement decision."""
    balancer = LoadBalancer(strategy="load_aware")
    engine, decision = balancer.route_with_decision(sample_request, sample_engines)

    assert engine is not None
    assert decision is not None

    # Check decision attributes
    assert decision.engine_id == engine.engine_id
    assert decision.strategy == "load_aware"
    assert decision.decision_time_ms >= 0.0
    assert decision.reason != ""
    assert decision.candidate_count == len(sample_engines)
    assert decision.selected_engine_load >= 0.0
    assert decision.load_variance >= 0.0

    # Check metadata
    assert "active_requests" in decision.metadata
    assert "queue_length" in decision.metadata
    assert "kv_cache_utilization" in decision.metadata
    assert "gpu_utilization" in decision.metadata


def test_placement_decision_to_dict(
    sample_engines: list[EngineInfo],
    sample_request: RequestMetadata,
) -> None:
    """Test PlacementDecision serialization."""
    balancer = LoadBalancer(strategy="load_aware")
    _, decision = balancer.route_with_decision(sample_request, sample_engines)

    assert decision is not None
    data = decision.to_dict()

    assert "engine_id" in data
    assert "strategy" in data
    assert "decision_time_ms" in data
    assert "reason" in data
    assert "candidate_count" in data
    assert "selected_engine_load" in data
    assert "load_variance" in data
    assert "metadata" in data


def test_load_balancer_strategies() -> None:
    """Test that all strategies are registered."""
    expected_strategies = ["round_robin", "least_loaded", "load_aware", "kv_cache_aware"]

    for strategy in expected_strategies:
        balancer = LoadBalancer(strategy=strategy)
        assert balancer.strategy == strategy


def test_load_balancer_invalid_strategy() -> None:
    """Test that invalid strategy raises error."""
    with pytest.raises(ValueError, match="Unknown strategy"):
        LoadBalancer(strategy="invalid_strategy")


def test_route_with_no_healthy_engines(sample_request: RequestMetadata) -> None:
    """Test routing when no engines are healthy."""
    unhealthy_engines = [
        EngineInfo(
            engine_id="engine-0",
            model_id="Qwen2-7B",
            host="localhost",
            port=8000,
            state=EngineState.ERROR,  # Not healthy
        )
    ]

    balancer = LoadBalancer(strategy="load_aware")
    engine = balancer.route(sample_request, unhealthy_engines)

    assert engine is None


def test_route_with_empty_candidates(sample_request: RequestMetadata) -> None:
    """Test routing with empty candidate list."""
    balancer = LoadBalancer(strategy="load_aware")
    engine = balancer.route(sample_request, [])

    assert engine is None


def test_load_variance_calculation(sample_engines: list[EngineInfo]) -> None:
    """Test that load variance is calculated correctly."""
    # Calculate expected variance
    loads = [e.total_load for e in sample_engines]
    mean_load = sum(loads) / len(loads)
    expected_variance = sum((load - mean_load) ** 2 for load in loads) / len(loads)

    # Get variance from placement decision
    balancer = LoadBalancer(strategy="load_aware")
    request = RequestMetadata(
        request_id="test-req",
        model_id="Qwen2-7B",
        prompt="Test",
    )
    _, decision = balancer.route_with_decision(request, sample_engines)

    assert decision is not None
    assert abs(decision.load_variance - expected_variance) < 1e-6


def test_decision_reason_contains_metrics(
    sample_engines: list[EngineInfo],
    sample_request: RequestMetadata,
) -> None:
    """Test that decision reason contains relevant metrics."""
    balancer = LoadBalancer(strategy="load_aware")
    _, decision = balancer.route_with_decision(sample_request, sample_engines)

    assert decision is not None
    reason = decision.reason.lower()

    # Should mention key metrics
    assert "load" in reason or "kv" in reason or "queue" in reason


def test_multiple_requests_load_distribution(sample_engines: list[EngineInfo]) -> None:
    """Test that load-aware strategy distributes requests better over time."""
    balancer = LoadBalancer(strategy="load_aware")

    # Track engine selection
    selections = {e.engine_id: 0 for e in sample_engines}

    # Route multiple requests
    for i in range(20):
        request = RequestMetadata(
            request_id=f"req-{i}",
            model_id="Qwen2-7B",
            prompt=f"Test {i}",
        )

        engine = balancer.route(request, sample_engines)
        if engine:
            selections[engine.engine_id] += 1
            # Simulate load increase
            engine.active_requests += 1

    # Engine-2 (initially least loaded) should get more requests
    # but not all due to load balancing
    assert selections["engine-2"] > 0

    # Check distribution is reasonably balanced
    # Filter out engines that got selected
    selected_counts = [count for count in selections.values() if count > 0]
    if len(selected_counts) > 1:
        # Distribution shouldn't be extremely skewed
        assert max(selected_counts) / min(selected_counts) < 10
