"""Tests for runtime optimizer framework (issue #48)."""

from __future__ import annotations

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any

import pytest

from sagellm_protocol import CapabilityDescriptor

from sagellm_core.runtime_optimizer import BaseStrategy, RuntimeOptimizer, RuntimeStats
from sagellm_core.worker.worker import Worker


class _AddTagStrategy(BaseStrategy):
    def __init__(self, tag: str, *, enabled: bool = True) -> None:
        super().__init__(name=f"add-{tag}", enabled=enabled)
        self._tag = tag

    def optimize(self, step: Any, stats: RuntimeStats) -> Any:
        return list(step) + [self._tag]


class _FakeBackend:
    def __init__(self, util: float = 42.0) -> None:
        self._util = util

    def capability(self) -> CapabilityDescriptor:
        return CapabilityDescriptor(device_name="fake-cpu", device_type="cpu")

    def memory_stats(self) -> dict[str, float]:
        return {"utilization": self._util}

    def get_kernel(self, name: str) -> Any:
        return lambda *args, **kwargs: None

    def get_attention_backend(self, name: str | None = None) -> Any:
        return SimpleNamespace(name="fake-attention")


class _FakeKVManager:
    def get_stats(self) -> dict[str, Any]:
        return {
            "prefix_cache": {
                "hit_count": 7,
                "query_count": 10,
            }
        }


class _FakeModelRunner:
    def __init__(self) -> None:
        self.last_groups: Any = None

    async def forward(self, groups: Any) -> list[Any]:
        self.last_groups = groups
        return ["ok"]


@dataclass
class _Group:
    request_id: str
    prompt: str
    max_tokens: int = 8


def test_strategy_register_remove_enable_disable() -> None:
    optimizer = RuntimeOptimizer(backend=_FakeBackend())

    strategy = _AddTagStrategy("x")
    optimizer.register_strategy(strategy)
    assert optimizer.list_strategies() == ["add-x"]

    optimizer.disable_strategy("add-x")
    assert strategy.enabled is False

    optimizer.enable_strategy("add-x")
    assert strategy.enabled is True

    removed = optimizer.remove_strategy("add-x")
    assert removed is strategy
    assert optimizer.list_strategies() == []


def test_optimize_step_pipeline_order() -> None:
    optimizer = RuntimeOptimizer(
        backend=_FakeBackend(),
        strategies=[_AddTagStrategy("a"), _AddTagStrategy("b")],
    )

    result = optimizer.optimize_step(["seed"])
    assert result == ["seed", "a", "b"]


def test_runtime_stats_collects_fields() -> None:
    stats = RuntimeStats()

    snapshot = stats.collect_from_step(
        [_Group(request_id="r1", prompt="hello world")],
        backend=_FakeBackend(util=66.0),
        kv_cache_manager=_FakeKVManager(),
    )

    assert snapshot.batch_size == 1
    assert snapshot.seq_len >= 2
    assert snapshot.gpu_utilization == 66.0
    assert snapshot.kv_cache_hit_rate == pytest.approx(0.7)

    summary = stats.summary()
    assert summary["num_samples"] == 1
    assert summary["avg_batch_size"] == 1


@pytest.mark.asyncio
async def test_worker_execute_model_uses_runtime_optimizer() -> None:
    class _InjectStrategy(BaseStrategy):
        def optimize(self, step: Any, stats: RuntimeStats) -> Any:
            return list(step) + ["optimized"]

    optimizer = RuntimeOptimizer(
        backend=_FakeBackend(),
        strategies=[_InjectStrategy("inject")],
    )

    worker = Worker(
        backend=_FakeBackend(),
        model_path="fake/model",
        runtime_optimizer=optimizer,
    )

    worker._is_running = True
    worker._model_runner = _FakeModelRunner()
    worker._kv_cache_manager = _FakeKVManager()

    groups = [SimpleNamespace(request_id="r1", prompt="one two")]
    result = await worker.execute_model(groups)

    assert result == ["ok"]
    assert worker._model_runner.last_groups[-1] == "optimized"
    assert worker.runtime_optimizer.stats.latest() is not None
