"""Tests for DynamicBatchingStrategy (issue #49)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from sagellm_core.runtime_optimizer import DynamicBatchingStrategy, 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_prefill(self) -> bool:
        return self.batch_type == "prefill"

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


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

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


def test_dynamic_batching_strategy_applies_backend_hint() -> None:
    backend = _BackendWithHint()
    strategy = DynamicBatchingStrategy(small_batch_threshold=8, large_batch_threshold=32)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = [
        _Batch(batch_id=1, requests=[1, 2, 3], batch_type="prefill", block_tables=[[], [], []]),
        _Batch(batch_id=2, requests=[1], batch_type="decode", block_tables=[[]]),
    ]
    optimizer.collect_stats(step)
    _ = optimizer.optimize_step(step)

    assert backend.hints
    assert backend.hints[-1]["profile"] == "small"


def test_dynamic_batching_reorders_prefill_before_decode() -> None:
    strategy = DynamicBatchingStrategy(enable_mixed_batching=True)
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    decode = _Batch(batch_id=2, requests=[1], batch_type="decode", block_tables=[[]])
    prefill = _Batch(batch_id=1, requests=[1], batch_type="prefill", block_tables=[[]])

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

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


def test_dynamic_batching_splits_large_decode_batch() -> None:
    strategy = DynamicBatchingStrategy(small_batch_threshold=2, large_batch_threshold=8)
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    decode = _Batch(
        batch_id=10,
        requests=["r1", "r2", "r3", "r4", "r5"],
        batch_type="decode",
        block_tables=[[1], [2], [3], [4], [5]],
    )

    optimized = optimizer.optimize_step([decode])

    assert len(optimized) == 3
    assert [len(item.requests) for item in optimized] == [2, 2, 1]


def test_dynamic_batching_auto_tunes_thresholds() -> None:
    strategy = DynamicBatchingStrategy(
        small_batch_threshold=4,
        large_batch_threshold=16,
        auto_tune_thresholds=True,
    )

    for _ in range(5):
        strategy.observe_execution(
            batch_size=2, latency_ms=30.0, metadata={"dynamic_batch_profile": "small"}
        )
        strategy.observe_execution(
            batch_size=10, latency_ms=20.0, metadata={"dynamic_batch_profile": "medium"}
        )

    assert strategy.small_batch_threshold > 4


def test_runtime_optimizer_report_execution_handles_outputs() -> None:
    strategy = DynamicBatchingStrategy(auto_tune_thresholds=False)
    optimizer = RuntimeOptimizer(backend=object(), strategies=[strategy])

    step = [_Batch(batch_id=1, requests=[1, 2], batch_type="decode", block_tables=[[], []])]
    optimizer.collect_stats(step)

    class _Output:
        forward_time_ms = 12.5

    optimizer.report_execution(step=step, outputs=[_Output(), _Output()])
    assert strategy._latency_by_profile["small"]
