"""Tests for KVComputeCoOptimizationStrategy (issue #52)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import pytest

from sagellm_core.runtime_optimizer import KVComputeCoOptimizationStrategy, RuntimeOptimizer


@dataclass
class _Batch:
    """Mock batch for testing."""

    batch_id: int
    requests: list[Any]
    batch_type: str
    block_tables: list[list[int]] | None = None
    metadata: dict[str, Any] | None = None

    @property
    def num_requests(self) -> int:
        return len(self.requests)

    @property
    def is_prefill(self) -> bool:
        return self.batch_type == "prefill"

    @property
    def is_decode(self) -> bool:
        return self.batch_type == "decode"


class _MockKVCacheManager:
    """Mock KV cache manager for testing."""

    def __init__(self, utilization: float = 0.5):
        self._utilization = utilization

    def get_stats(self) -> dict[str, Any]:
        return {
            "pool": {
                "utilization": self._utilization,
                "allocated_tokens": 1000,
            },
            "max_tokens": 2000,
        }


class _MockBackend:
    """Mock backend with KV cache manager."""

    def __init__(self, kv_cache_manager: _MockKVCacheManager):
        self.kv_cache_manager = kv_cache_manager


def _prefill_batch(batch_id: int, num_requests: int) -> _Batch:
    return _Batch(
        batch_id=batch_id,
        requests=list(range(num_requests)),
        batch_type="prefill",
        block_tables=[[0, 1] for _ in range(num_requests)],
    )


def _decode_batch(batch_id: int, num_requests: int) -> _Batch:
    return _Batch(
        batch_id=batch_id,
        requests=list(range(num_requests)),
        batch_type="decode",
        block_tables=[[1, 2] for _ in range(num_requests)],
    )


def test_kv_compute_coopt_strategy_initialization() -> None:
    """Test strategy initialization with default parameters."""
    strategy = KVComputeCoOptimizationStrategy()

    assert strategy.enabled is True
    assert strategy.enable_kv_aware_scheduling is True
    assert strategy.enable_speculative_decoding is False
    assert strategy.enable_early_exit is False
    assert strategy.enable_kv_transfer is True
    assert strategy.recompute_threshold == 0.3
    assert strategy.speculative_lookahead == 4


def test_kv_compute_coopt_strategy_custom_parameters() -> None:
    """Test strategy initialization with custom parameters."""
    strategy = KVComputeCoOptimizationStrategy(
        enable_kv_aware_scheduling=True,
        enable_speculative_decoding=True,
        enable_early_exit=True,
        enable_kv_transfer=True,
        recompute_threshold=0.5,
        speculative_lookahead=8,
        early_exit_layers=[4, 8, 12],
        kv_transfer_threshold=2048 * 1024,
    )

    assert strategy.enable_speculative_decoding is True
    assert strategy.enable_early_exit is True
    assert strategy.recompute_threshold == 0.5
    assert strategy.speculative_lookahead == 8
    assert strategy.early_exit_layers == [4, 8, 12]
    assert strategy.kv_transfer_threshold == 2048 * 1024


def test_kv_compute_coopt_strategy_invalid_parameters() -> None:
    """Test strategy initialization with invalid parameters."""
    with pytest.raises(ValueError, match="recompute_threshold must be in"):
        KVComputeCoOptimizationStrategy(recompute_threshold=1.5)

    with pytest.raises(ValueError, match="speculative_lookahead must be"):
        KVComputeCoOptimizationStrategy(speculative_lookahead=0)

    with pytest.raises(ValueError, match="kv_transfer_threshold must be"):
        KVComputeCoOptimizationStrategy(kv_transfer_threshold=-1)


def test_kv_aware_scheduling_low_utilization() -> None:
    """Test KV-aware scheduling with low cache utilization."""
    kv_manager = _MockKVCacheManager(utilization=0.2)
    backend = _MockBackend(kv_manager)
    strategy = KVComputeCoOptimizationStrategy(enable_kv_aware_scheduling=True)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    # With low utilization, order should remain the same
    prefill = _prefill_batch(1, 2)
    decode = _decode_batch(2, 2)

    optimized = optimizer.optimize_step([prefill, decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 2
    assert optimized[0].batch_type == "prefill"
    assert optimized[1].batch_type == "decode"


def test_kv_aware_scheduling_high_utilization() -> None:
    """Test KV-aware scheduling with high cache utilization."""
    kv_manager = _MockKVCacheManager(utilization=0.8)
    backend = _MockBackend(kv_manager)
    strategy = KVComputeCoOptimizationStrategy(
        enable_kv_aware_scheduling=True, recompute_threshold=0.3
    )
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    # With high utilization, decode should be prioritized
    prefill = _prefill_batch(1, 2)
    decode = _decode_batch(2, 2)

    optimized = optimizer.optimize_step([prefill, decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 2
    assert optimized[0].batch_type == "decode"
    assert optimized[1].batch_type == "prefill"


def test_speculative_decoding_optimization() -> None:
    """Test speculative decoding metadata injection."""
    strategy = KVComputeCoOptimizationStrategy(
        enable_speculative_decoding=True, speculative_lookahead=6
    )
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    decode = _decode_batch(1, 3)
    optimized = optimizer.optimize_step([decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 1
    assert hasattr(optimized[0], "metadata")
    assert optimized[0].metadata["speculative_lookahead"] == 6
    assert optimized[0].metadata["enable_speculation"] is True


def test_early_exit_optimization() -> None:
    """Test early exit strategy metadata injection."""
    strategy = KVComputeCoOptimizationStrategy(
        enable_early_exit=True, early_exit_layers=[6, 12, 18]
    )
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    decode = _decode_batch(1, 2)
    optimized = optimizer.optimize_step([decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 1
    assert hasattr(optimized[0], "metadata")
    assert optimized[0].metadata["early_exit_layers"] == [6, 12, 18]
    assert optimized[0].metadata["enable_early_exit"] is True


def test_kv_transfer_optimization() -> None:
    """Test KV transfer optimization for large batches."""
    kv_manager = _MockKVCacheManager(utilization=0.5)
    backend = _MockBackend(kv_manager)
    strategy = KVComputeCoOptimizationStrategy(enable_kv_transfer=True, kv_transfer_threshold=1024)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    # Large batch should trigger KV transfer optimization
    decode = _decode_batch(1, 10)
    optimized = optimizer.optimize_step([decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 1
    assert hasattr(optimized[0], "metadata")
    assert optimized[0].metadata["enable_kv_compression"] is True
    assert "kv_transfer_batch_size" in optimized[0].metadata


def test_kv_transfer_small_batch_no_optimization() -> None:
    """Test KV transfer skipped for small batches."""
    kv_manager = _MockKVCacheManager(utilization=0.5)
    backend = _MockBackend(kv_manager)
    strategy = KVComputeCoOptimizationStrategy(
        enable_kv_transfer=True,
        kv_transfer_threshold=1024 * 1024 * 10,  # 10MB
    )
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    # Small batch should not trigger KV transfer optimization
    decode = _decode_batch(1, 2)
    optimized = optimizer.optimize_step([decode])

    assert isinstance(optimized, list)
    assert len(optimized) == 1


def test_combined_optimizations() -> None:
    """Test all optimizations enabled together."""
    kv_manager = _MockKVCacheManager(utilization=0.8)
    backend = _MockBackend(kv_manager)
    strategy = KVComputeCoOptimizationStrategy(
        enable_kv_aware_scheduling=True,
        enable_speculative_decoding=True,
        enable_early_exit=True,
        enable_kv_transfer=True,
        recompute_threshold=0.3,
    )
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    prefill = _prefill_batch(1, 2)
    decode = _decode_batch(2, 5)

    optimized = optimizer.optimize_step([prefill, decode])

    # Should reorder to prioritize decode
    assert len(optimized) == 2
    assert optimized[0].batch_type == "decode"
    assert optimized[1].batch_type == "prefill"

    # Decode batch should have all optimization hints
    decode_batch = optimized[0]
    assert hasattr(decode_batch, "metadata")
    assert decode_batch.metadata["enable_speculation"] is True
    assert decode_batch.metadata["enable_early_exit"] is True


def test_metrics_tracking() -> None:
    """Test strategy metrics tracking."""
    strategy = KVComputeCoOptimizationStrategy(enable_speculative_decoding=True)

    # Simulate observations
    strategy.observe_execution(
        batch_size=4,
        latency_ms=10.0,
        metadata={
            "speculative_accepted": 3,
            "speculative_total": 4,
            "early_exit_triggered": True,
            "recomputation_triggered": False,
        },
    )

    strategy.observe_execution(
        batch_size=4,
        latency_ms=12.0,
        metadata={
            "speculative_accepted": 2,
            "speculative_total": 4,
            "early_exit_triggered": False,
            "recomputation_triggered": True,
        },
    )

    metrics = strategy.get_metrics()

    assert metrics["speculative_accept_rate"] == 5 / 8  # (3+2)/(4+4)
    assert metrics["early_exit_count"] == 1
    assert metrics["recompute_count"] == 1


def test_strategy_can_be_disabled() -> None:
    """Test strategy can be disabled."""
    strategy = KVComputeCoOptimizationStrategy(enabled=False)
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    decode = _decode_batch(1, 2)
    optimized = optimizer.optimize_step([decode])

    # Should remain unchanged when disabled
    assert optimized == [decode]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
