"""Tests for LLMEngine backend runtime integration (issue #29)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import pytest

from sagellm_core import LLMEngine, LLMEngineConfig
from sagellm_protocol import CapabilityDescriptor, DType


@dataclass
class _FakeKVBlock:
    handle: int
    num_tokens: int
    dtype: DType
    device: str


class _FakeEvent:
    def __init__(self, event_id: int) -> None:
        self.event_id = event_id
        self.timestamp = float(event_id)

    def synchronize(self) -> None:
        return None

    def elapsed_ms(self, start: Any) -> float:
        return max((self.timestamp - getattr(start, "timestamp", 0.0)) * 1000.0, 0.0)


class _FakeStream:
    def __init__(self) -> None:
        self.recorded: list[int] = []
        self.synced = False

    def synchronize(self) -> None:
        self.synced = True

    def record_event(self, event: Any) -> None:
        self.recorded.append(getattr(event, "event_id", -1))


class _FakeBackend:
    def __init__(self) -> None:
        self._stream = _FakeStream()
        self._event_id = 0
        self._kv_id = 0
        self.kv_alloc_tokens: list[int] = []
        self.kv_freed_handles: list[int] = []
        self.attn_backend = object()
        self.kernel_hits: list[str] = []

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

    def create_stream(self) -> _FakeStream:
        return self._stream

    def create_event(self) -> _FakeEvent:
        self._event_id += 1
        return _FakeEvent(self._event_id)

    def kv_block_alloc(self, num_tokens: int, dtype: DType) -> _FakeKVBlock:
        self._kv_id += 1
        self.kv_alloc_tokens.append(num_tokens)
        return _FakeKVBlock(handle=self._kv_id, num_tokens=num_tokens, dtype=dtype, device="cpu")

    def kv_block_free(self, block: _FakeKVBlock) -> None:
        self.kv_freed_handles.append(block.handle)

    def kv_block_copy(self, src: Any, dst: Any) -> None:
        return None

    def kv_block_migrate(self, block: Any, target_device: str) -> Any:
        return block

    def memory_stats(self) -> dict[str, int]:
        return {"used": 16 * 1024 * 1024}

    def clear_cache(self) -> None:
        return None

    def register_kernel(self, name: str, impl: Any) -> None:
        return None

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

    def get_attention_backend(self, name: str | None = None) -> Any:
        return self.attn_backend

    def get_supported_attention_backends(self) -> list[str]:
        return ["cpu"]

    def get_torch_device(self) -> str:
        return "cpu"

    def to_device(self, tensor: Any) -> Any:
        return tensor


@pytest.mark.asyncio
async def test_start_initializes_backend_runtime_primitives(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Engine start should initialize stream/attention/kernel runtime handles."""
    config = LLMEngineConfig(model_path="dummy", backend_type="cpu")
    engine = LLMEngine(config)
    fake_backend = _FakeBackend()

    monkeypatch.setattr(engine, "_init_backend", lambda: fake_backend)

    async def _fake_load_model() -> None:
        engine._model = object()
        engine._tokenizer = object()
        engine._torch = None

    monkeypatch.setattr(engine, "_load_model", _fake_load_model)

    await engine.start()
    try:
        assert engine._backend_default_stream is not None
        assert engine._backend_attention_backend is fake_backend.attn_backend
        assert "linear" in engine._backend_kernels
        assert "embedding" in engine._backend_kernels
    finally:
        await engine.stop()


def test_backend_request_context_alloc_and_free() -> None:
    """Backend request context should allocate and free KV blocks."""
    config = LLMEngineConfig(model_path="dummy", backend_type="cpu")
    engine = LLMEngine(config)
    fake_backend = _FakeBackend()

    engine._backend = fake_backend
    engine._initialize_backend_runtime_primitives()

    ctx = engine._begin_backend_request(request_id="req-ctx", prompt_tokens=12, max_new_tokens=8)
    assert fake_backend.kv_alloc_tokens == [20]
    assert ctx.kv_block is not None

    engine._end_backend_request(ctx)
    assert fake_backend.kv_freed_handles == [ctx.kv_block.handle]
    assert engine._backend_default_stream is not None
    assert engine._backend_default_stream.synced is True


@pytest.mark.asyncio
async def test_generate_uses_backend_request_lifecycle(monkeypatch: pytest.MonkeyPatch) -> None:
    """generate() should enter and exit backend request context exactly once."""
    torch = pytest.importorskip("torch")

    class _Tokenizer:
        pad_token_id = 0

        def __call__(
            self, prompt: str, return_tensors: str = "pt", padding: bool = True
        ) -> dict[str, Any]:
            return {
                "input_ids": torch.tensor([[1, 2]], dtype=torch.long),
                "attention_mask": torch.tensor([[1, 1]], dtype=torch.long),
            }

        def decode(self, token_ids: Any, skip_special_tokens: bool = True) -> str:
            return "ok"

    class _Model:
        def generate(self, input_ids: Any, **kwargs: Any) -> Any:
            del kwargs
            return type("_Out", (), {"sequences": torch.tensor([[1, 2, 3]], dtype=torch.long)})()

    config = LLMEngineConfig(model_path="dummy", backend_type="cpu")
    engine = LLMEngine(config)
    fake_backend = _FakeBackend()

    engine._is_running = True
    engine._backend = fake_backend
    engine._torch = torch
    engine._model = _Model()
    engine._tokenizer = _Tokenizer()
    engine._initialize_backend_runtime_primitives()

    begin_calls: list[str] = []
    end_calls: list[str] = []
    original_begin = engine._begin_backend_request
    original_end = engine._end_backend_request

    def _wrapped_begin(*, request_id: str, prompt_tokens: int, max_new_tokens: int):
        begin_calls.append(request_id)
        return original_begin(
            request_id=request_id,
            prompt_tokens=prompt_tokens,
            max_new_tokens=max_new_tokens,
        )

    def _wrapped_end(ctx: Any) -> None:
        end_calls.append(ctx.request_id)
        original_end(ctx)

    monkeypatch.setattr(engine, "_begin_backend_request", _wrapped_begin)
    monkeypatch.setattr(engine, "_end_backend_request", _wrapped_end)

    response = await engine.generate(prompt="hi", max_tokens=4, request_id="req-generate")

    assert response.request_id == "req-generate"
    assert begin_calls == ["req-generate"]
    assert end_calls == ["req-generate"]
    assert fake_backend.kv_alloc_tokens == [6]
    assert len(fake_backend.kv_freed_handles) == 1
    assert response.metrics.peak_mem_mb == 16
