"""Tests for CUDAGraphCaptureStrategy (issue #50)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from sagellm_core.runtime_optimizer import CUDAGraphCaptureStrategy, RuntimeOptimizer


@dataclass
class _Batch:
    batch_id: int
    requests: list[Any]
    batch_type: str
    block_tables: list[list[int]]

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

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


class _BackendWithGraphCalls:
    def __init__(self) -> None:
        self.captures: list[dict[str, Any]] = []
        self.replays: list[dict[str, Any]] = []
        self.hints: list[dict[str, Any]] = []

    def capture_cuda_graph(self, payload: dict[str, Any]) -> None:
        self.captures.append(payload)

    def replay_cuda_graph(self, payload: dict[str, Any]) -> None:
        self.replays.append(payload)

    def set_runtime_hint(self, payload: dict[str, Any]) -> None:
        self.hints.append(payload)


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


def test_cuda_graph_strategy_capture_then_replay() -> None:
    backend = _BackendWithGraphCalls()
    strategy = CUDAGraphCaptureStrategy(cache_size=8, warmup_runs=2)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = [_decode_batch(batch_id=1, batch_size=4)]

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert not backend.captures

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert len(backend.captures) == 1

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert len(backend.replays) == 1
    assert strategy.cache_hit_rate() > 0.0


def test_cuda_graph_strategy_cache_hit_rate_for_stable_workload() -> None:
    backend = _BackendWithGraphCalls()
    strategy = CUDAGraphCaptureStrategy(cache_size=4, warmup_runs=1)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = [_decode_batch(batch_id=3, batch_size=8)]
    for _ in range(20):
        optimizer.collect_stats(step)
        optimizer.optimize_step(step)

    stats = strategy.cache_stats()
    assert stats["hit_rate"] >= 0.8


def test_cuda_graph_strategy_supports_dynamic_toggle() -> None:
    backend = _BackendWithGraphCalls()
    strategy = CUDAGraphCaptureStrategy(cache_size=4, warmup_runs=1)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = [_decode_batch(batch_id=5, batch_size=2)]
    strategy.set_capture_enabled(False)
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert not backend.captures

    strategy.set_capture_enabled(True)
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert backend.captures


def test_cuda_graph_strategy_lru_eviction() -> None:
    strategy = CUDAGraphCaptureStrategy(cache_size=2, warmup_runs=1)
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    step_a = [_decode_batch(batch_id=1, batch_size=2)]
    step_b = [_decode_batch(batch_id=2, batch_size=3)]
    step_c = [_decode_batch(batch_id=3, batch_size=4)]

    optimizer.collect_stats(step_a)
    optimizer.optimize_step(step_a)
    optimizer.collect_stats(step_b)
    optimizer.optimize_step(step_b)
    optimizer.collect_stats(step_c)
    optimizer.optimize_step(step_c)

    stats = strategy.cache_stats()
    assert stats["cache_entries"] == 2
