"""Tests for Backend Kernel usage verification (issue #30).

This test suite verifies that backend kernels are actually called during
inference, not just obtained but ignored. This addresses the core concern
in issue #30: "Backend Kernels 虽已获取但未在实际推理中使用".

Test Strategy:
1. Patch model layers with backend kernels (via patch_model_with_kernels)
2. Run inference
3. Verify kernel call counts > 0
4. Compare output consistency with/without kernels
"""

from __future__ import annotations

import sys
from types import SimpleNamespace
from typing import Any

import pytest
import torch
import torch.nn as nn

from sagellm_backend import CPUBackendProvider
from sagellm_core.worker.model_runner.kernel_layers import (
    KernelCallStats,
    get_kernel_stats,
    patch_model_with_kernels,
    reset_kernel_stats,
)


class SimpleLLM(nn.Module):
    """Minimal LLM for testing kernel patching."""

    def __init__(
        self,
        vocab_size: int = 100,
        hidden_size: int = 32,
        num_layers: int = 2,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList(
            [nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)]
        )
        self.lm_head = nn.Linear(hidden_size, vocab_size)

    def forward(self, input_ids: torch.Tensor, **kwargs: Any) -> torch.Tensor:
        """Forward pass (accepts attention_mask for compatibility with ModelRunner)."""
        x = self.embedding(input_ids)  # Embedding kernel should be called
        for layer in self.layers:
            x = layer(x)  # Linear kernel should be called per layer
        logits = self.lm_head(x)  # Linear kernel should be called
        return logits


def test_kernel_stats_creation():
    """Test KernelCallStats creation and reset."""
    stats = KernelCallStats()
    assert stats.linear_calls == 0
    assert stats.embedding_calls == 0
    assert stats.rmsnorm_calls == 0
    assert stats.total_calls == 0

    # Modify stats
    stats.linear_calls = 10
    stats.embedding_calls = 5
    stats.total_calls = 15

    # Reset
    stats.reset()
    assert stats.linear_calls == 0
    assert stats.total_calls == 0


def test_kernel_stats_to_dict():
    """Test KernelCallStats to_dict conversion."""
    stats = KernelCallStats()
    stats.linear_calls = 10
    stats.embedding_calls = 5
    stats.rmsnorm_calls = 2
    stats.total_calls = 17

    d = stats.to_dict()
    assert d == {
        "linear_calls": 10,
        "embedding_calls": 5,
        "rmsnorm_calls": 2,
        "total_calls": 17,
    }


def test_global_kernel_stats_access():
    """Test global kernel stats getter/reset."""
    reset_kernel_stats()  # Start fresh
    stats = get_kernel_stats()
    assert stats.total_calls == 0

    # Modify global stats
    stats.linear_calls = 5
    stats.total_calls = 5

    # Get again - should be same object
    stats2 = get_kernel_stats()
    assert stats2.linear_calls == 5
    assert stats2.total_calls == 5

    # Reset
    reset_kernel_stats()
    assert get_kernel_stats().total_calls == 0


def test_patch_model_with_kernels_coverage():
    """Test that patch_model_with_kernels patches all target layers."""
    model = SimpleLLM(vocab_size=100, hidden_size=32, num_layers=3)
    backend = CPUBackendProvider()

    linear_kernel = backend.get_kernel("linear")
    embedding_kernel = backend.get_kernel("embedding")

    patched_count = patch_model_with_kernels(
        model=model,
        backend=backend,
        linear_kernel=linear_kernel,
        embedding_kernel=embedding_kernel,
    )

    # Expected: 1 embedding + 3 layers + 1 lm_head = 5 patched layers
    assert patched_count == 5, f"Expected 5 patched layers, got {patched_count}"


def test_backend_kernels_actually_called():
    """CRITICAL TEST: Verify backend kernels are actually invoked during forward pass.

    This is the core verification for issue #30: demonstrate that backend
    kernels are not just obtained, but actively used in inference.
    """
    reset_kernel_stats()  # Start fresh

    # Create model and backend
    model = SimpleLLM(vocab_size=100, hidden_size=32, num_layers=3)
    backend = CPUBackendProvider()

    linear_kernel = backend.get_kernel("linear")
    embedding_kernel = backend.get_kernel("embedding")

    # Patch model with backend kernels
    patch_model_with_kernels(
        model=model,
        backend=backend,
        linear_kernel=linear_kernel,
        embedding_kernel=embedding_kernel,
    )

    # Prepare input
    input_ids = torch.tensor([[1, 2, 3, 4, 5]])

    # Run forward pass
    model.eval()
    with torch.inference_mode():
        _logits = model(input_ids)

    # Verify kernels were called
    stats = get_kernel_stats()

    # Should have:
    # - 1 embedding call (for embedding layer)
    # - 4 linear calls (3 intermediate layers + 1 lm_head)
    assert stats.embedding_calls > 0, "Embedding kernel was not called"
    assert stats.linear_calls > 0, "Linear kernel was not called"
    assert stats.total_calls > 0, "No backend kernels were called"

    # More specific assertions
    assert stats.embedding_calls == 1, f"Expected 1 embedding call, got {stats.embedding_calls}"
    assert stats.linear_calls == 4, f"Expected 4 linear calls, got {stats.linear_calls}"
    assert stats.total_calls == 5, f"Expected 5 total calls, got {stats.total_calls}"


def test_kernel_call_stats_accumulate_across_multiple_forwards():
    """Test that kernel call stats accumulate across multiple forward passes."""
    reset_kernel_stats()

    model = SimpleLLM(vocab_size=100, hidden_size=32, num_layers=2)
    backend = CPUBackendProvider()

    linear_kernel = backend.get_kernel("linear")
    embedding_kernel = backend.get_kernel("embedding")

    patch_model_with_kernels(
        model=model,
        backend=backend,
        linear_kernel=linear_kernel,
        embedding_kernel=embedding_kernel,
    )

    input_ids = torch.tensor([[1, 2, 3]])

    model.eval()
    with torch.inference_mode():
        # First forward pass
        _logits1 = model(input_ids)
        stats1 = get_kernel_stats()
        first_total = stats1.total_calls

        # Second forward pass
        _logits2 = model(input_ids)
        stats2 = get_kernel_stats()
        second_total = stats2.total_calls

    # Stats should accumulate
    assert second_total == 2 * first_total, "Kernel stats should accumulate across forwards"
    assert stats2.embedding_calls == 2, "Should have 2 embedding calls after 2 forwards"


def test_output_consistency_with_and_without_kernel_patching():
    """Test that patched model produces consistent output with unpatched model.

    This ensures that kernel patching doesn't break correctness.
    """
    torch.manual_seed(42)

    # Create two identical models
    model1 = SimpleLLM(vocab_size=100, hidden_size=32, num_layers=2)
    model2 = SimpleLLM(vocab_size=100, hidden_size=32, num_layers=2)

    # Copy weights to ensure identical
    model2.load_state_dict(model1.state_dict())

    # Patch model2 with backend kernels
    backend = CPUBackendProvider()
    linear_kernel = backend.get_kernel("linear")
    embedding_kernel = backend.get_kernel("embedding")

    patch_model_with_kernels(
        model=model2,
        backend=backend,
        linear_kernel=linear_kernel,
        embedding_kernel=embedding_kernel,
    )

    # Run inference on both
    input_ids = torch.tensor([[1, 2, 3, 4, 5]])

    model1.eval()
    model2.eval()

    with torch.inference_mode():
        logits1 = model1(input_ids)
        logits2 = model2(input_ids)

    # Outputs should be very close (allow small numerical differences)
    assert torch.allclose(logits1, logits2, rtol=1e-4, atol=1e-5), (
        "Patched model output differs from unpatched model"
    )


@pytest.mark.asyncio
async def test_model_runner_reports_kernel_stats(monkeypatch: pytest.MonkeyPatch):
    """Test that ModelRunner.get_kernel_call_stats() returns valid statistics.

    This verifies the integration between kernel_layers and ModelRunner.
    """

    # Use fake transformers to avoid large model loading
    class _FakeTokenizer:
        pad_token = "<pad>"
        eos_token = "<eos>"
        eos_token_id = 2

        @classmethod
        def from_pretrained(cls, *args: Any, **kwargs: Any):
            return cls()

        def __call__(self, text: str, **kwargs: Any):
            # Return fake token IDs
            return {"input_ids": torch.tensor([[1, 2, 3, 4, 5]])}

        def decode(self, token_ids: list[int], **kwargs: Any) -> str:
            return "generated output"

    class _FakeAutoModelForCausalLM:
        @staticmethod
        def from_pretrained(*args: Any, **kwargs: Any):
            return SimpleLLM(vocab_size=100, hidden_size=32, num_layers=2)

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

    # Import after monkeypatching
    from sagellm_core.worker.model_runner.model_runner import ModelRunner

    # Create ModelRunner with CPU backend
    backend = CPUBackendProvider()
    runner = ModelRunner(
        backend=backend,
        comm=None,
        model_path="fake-model",
        max_new_tokens=10,
        dtype="float32",
    )

    # Start runner (this patches the model)
    await runner.start()

    # Run inference
    output = await runner.generate_text("Hello, world!", max_tokens=5)

    # Get kernel stats
    stats = runner.get_kernel_call_stats()

    # Verify stats are populated
    assert stats["total_calls"] > 0, "No kernel calls recorded"
    assert stats["linear_calls"] > 0, "No linear kernel calls recorded"
    assert stats["embedding_calls"] > 0, "No embedding kernel calls recorded"

    # Verify output is valid
    assert output.finished is True
    assert len(output.output_token_ids) > 0

    await runner.stop()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
