"""Tests for ModelRunner SageAttention replacement wiring (issue #37)."""

from __future__ import annotations

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

import pytest
import torch.nn as nn

from sagellm_protocol import CapabilityDescriptor


class _FakeAttentionBackend:
    name = "fake-attention"


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

    def get_kernel(self, name: str) -> Any:
        if name in {"linear", "embedding", "rmsnorm", "silu"}:
            return lambda *args, **kwargs: None
        raise KeyError(name)

    def get_attention_backend(self, name: str | None = None) -> _FakeAttentionBackend:
        return _FakeAttentionBackend()


class _FakeTokenizer:
    pad_token = "<pad>"
    eos_token = "<eos>"

    @classmethod
    def from_pretrained(cls, *args: Any, **kwargs: Any) -> "_FakeTokenizer":
        return cls()


class _FakeModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.proj = nn.Linear(4, 4)

    def eval(self) -> "_FakeModel":
        return self


class _FakeAutoModelForCausalLM:
    @staticmethod
    def from_pretrained(*args: Any, **kwargs: Any) -> _FakeModel:
        return _FakeModel()


@dataclass
class _FakeKVManager:
    max_tokens: int = 1024


@pytest.mark.asyncio
async def test_model_runner_starts_with_attention_replacement(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """ModelRunner.start should wire replace_attention_layers with backend + KV manager."""
    from sagellm_core.worker.model_runner.model_runner import ModelRunner

    if "_replace_attention_layers_with_backend" not in inspect.getsource(ModelRunner.start):
        pytest.skip("Imported ModelRunner.start does not include attention replacement hook")

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        SimpleNamespace(
            AutoTokenizer=_FakeTokenizer,
            AutoModelForCausalLM=_FakeAutoModelForCausalLM,
        ),
    )

    monkeypatch.setattr(
        "sagellm_core.worker.model_runner.kernel_layers.patch_model_with_kernels",
        lambda **kwargs: 0,
    )

    captured: dict[str, Any] = {}

    def _fake_replace_attention_layers(
        model: Any,
        attention_backend: Any,
        kv_cache_manager: Any | None = None,
        **kwargs: Any,
    ) -> Any:
        captured["model"] = model
        captured["attention_backend"] = attention_backend
        captured["kv_cache_manager"] = kv_cache_manager
        captured["kwargs"] = kwargs
        return model

    monkeypatch.setattr(
        "sagellm_core.attention.replace_attention_layers",
        _fake_replace_attention_layers,
    )
    monkeypatch.setattr(
        "sagellm_core.attention.layer_registry.replace_attention_layers",
        _fake_replace_attention_layers,
    )

    kv_manager = _FakeKVManager()
    runner = ModelRunner(
        backend=_FakeBackend(),
        model_path="fake/model",
        kv_cache_manager=kv_manager,
    )

    await runner.start()
    try:
        assert runner.is_running
        assert captured["model"] is not None
        assert captured["attention_backend"].name == "fake-attention"
        assert captured["kv_cache_manager"] is kv_manager
        assert captured["kwargs"]["inplace"] is True
    finally:
        await runner.stop()
