"""Tests for AdaptivePrecisionStrategy (issue #51)."""

from __future__ import annotations

from dataclasses import dataclass

from sagellm_core.runtime_optimizer import AdaptivePrecisionStrategy, RuntimeOptimizer


@dataclass
class _Request:
    request_id: str
    target_latency_ms: float | None = None


@dataclass
class _Batch:
    batch_id: int
    requests: list[_Request]
    batch_type: str = "decode"
    target_latency_ms: float | None = None

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


class _BackendWithPrecision:
    def __init__(self) -> None:
        self.precisions: list[str] = []
        self.layer_maps: list[dict[str, str]] = []
        self.request_maps: list[dict[str, str]] = []

    def get_supported_precisions(self) -> list[str]:
        return ["fp16", "fp8", "int8"]

    def set_precision(self, precision: str) -> None:
        self.precisions.append(precision)

    def set_layer_precision(self, layer_map: dict[str, str]) -> None:
        self.layer_maps.append(layer_map)

    def set_request_precision(self, request_map: dict[str, str]) -> None:
        self.request_maps.append(request_map)


def _step(*, request_latency_ms: float | None = None) -> list[_Batch]:
    req = _Request(request_id="r1", target_latency_ms=request_latency_ms)
    return [_Batch(batch_id=1, requests=[req], batch_type="decode")]


def test_adaptive_precision_degrades_when_latency_over_budget() -> None:
    backend = _BackendWithPrecision()
    strategy = AdaptivePrecisionStrategy(target_latency_ms=50.0)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = _step()
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert strategy.current_precision() == "fp16"

    class _Output:
        forward_time_ms = 70.0

    optimizer.report_execution(step=step, outputs=[_Output()])
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)

    assert strategy.current_precision() == "fp8"
    assert backend.precisions[-1] == "fp8"


def test_adaptive_precision_upgrades_when_quality_low() -> None:
    backend = _BackendWithPrecision()
    strategy = AdaptivePrecisionStrategy(target_latency_ms=50.0)
    strategy.set_quality_score(0.80)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = _step()

    class _OutputHighLatency:
        forward_time_ms = 80.0

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    optimizer.report_execution(step=step, outputs=[_OutputHighLatency()])
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert strategy.current_precision() == "fp8"

    class _OutputLowLatency:
        forward_time_ms = 20.0

    strategy.set_quality_score(0.70)
    optimizer.report_execution(step=step, outputs=[_OutputLowLatency()])
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    assert strategy.current_precision() == "fp16"


def test_adaptive_precision_per_request_policy_and_switch_latency() -> None:
    backend = _BackendWithPrecision()
    strategy = AdaptivePrecisionStrategy(target_latency_ms=50.0)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    fast_req = _Request(request_id="fast", target_latency_ms=20.0)
    normal_req = _Request(request_id="normal", target_latency_ms=80.0)
    step = [_Batch(batch_id=7, requests=[fast_req, normal_req], batch_type="decode")]

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    stats = strategy.stats()

    assert stats["last_switch_latency_ms"] < 10.0
    assert backend.request_maps[-1]["fast"] == "int8"
    assert backend.request_maps[-1]["normal"] in {"fp16", "fp8", "int8"}


def test_adaptive_precision_respects_supported_precisions() -> None:
    class _FP16OnlyBackend(_BackendWithPrecision):
        def get_supported_precisions(self) -> list[str]:
            return ["fp16"]

    backend = _FP16OnlyBackend()
    strategy = AdaptivePrecisionStrategy(target_latency_ms=30.0)
    optimizer = RuntimeOptimizer(backend=backend, strategies=[strategy])

    step = _step()

    class _Output:
        forward_time_ms = 120.0

    optimizer.collect_stats(step)
    optimizer.optimize_step(step)
    optimizer.report_execution(step=step, outputs=[_Output()])
    optimizer.collect_stats(step)
    optimizer.optimize_step(step)

    assert strategy.current_precision() == "fp16"
