"""Integration test for logical→physical operator mapping (issue #28).

This test verifies that the kernel injection framework properly connects
backend kernels to core execution layers.
"""

from __future__ import annotations

import pytest


@pytest.mark.asyncio
async def test_model_runner_kernel_injection_setup():
    """Test that ModelRunner properly sets up kernel injection framework."""
    from sagellm_backend import get_provider
    from sagellm_core.worker.model_runner import ModelRunner

    backend = get_provider("cpu")
    runner = ModelRunner(
        backend=backend,
        model_path="sshleifer/tiny-gpt2",
        dtype="float32",
        trust_remote_code=True,
        max_new_tokens=10,
    )

    await runner.start()

    # Verify that kernel injection components are initialized
    assert runner._operator_registry is not None, "OperatorRegistry should be initialized"
    assert runner._kernel_injector is not None, "KernelInjector should be initialized"

    # Verify that backend kernels are registered
    registry = runner._operator_registry
    assert registry.get_kernel_count() > 0, "Backend kernels should be registered"

    # List registered operators
    operators = registry.list_operators()
    print(f"\nRegistered operators: {[op.value for op in operators]}")

    await runner.stop()


@pytest.mark.asyncio
async def test_kernel_injection_with_model():
    """Test that kernel injection is attempted on model loading."""
    from sagellm_backend import get_provider
    from sagellm_core.worker.model_runner import ModelRunner

    backend = get_provider("cpu")
    runner = ModelRunner(
        backend=backend,
        model_path="sshleifer/tiny-gpt2",
        dtype="float32",
        trust_remote_code=True,
        max_new_tokens=10,
    )

    await runner.start()

    # Verify model is loaded
    assert runner._model is not None, "Model should be loaded"
    assert runner.is_running, "ModelRunner should be running"

    # The _inject_backend_kernels_full method should have been called
    # For standard HF models, it returns empty dict (no injectable layers)
    # but the framework is properly set up

    await runner.stop()


@pytest.mark.asyncio
async def test_backend_kernels_registration():
    """Test that backend properly registers its kernels into OperatorRegistry."""
    from sagellm_backend import get_provider
    from sagellm_core.kernel_injection.integration import setup_kernel_injection
    from sagellm_protocol import OperatorType

    backend = get_provider("cpu")
    registry, injector = setup_kernel_injection(backend)

    # Check that common operators are registered
    assert registry.has_operator(OperatorType.LINEAR), "LINEAR operator should be registered"
    assert registry.has_operator(OperatorType.EMBEDDING), "EMBEDDING operator should be registered"

    # Query for CPU kernels
    linear_spec = registry.query_kernel(OperatorType.LINEAR, device_type="cpu")
    print(f"\nLinear kernel: {linear_spec}")

    embedding_spec = registry.query_kernel(OperatorType.EMBEDDING, device_type="cpu")
    print(f"Embedding kernel: {embedding_spec}")


@pytest.mark.asyncio
async def test_operator_mapping_with_custom_layer():
    """Test operator mapping with custom layers that support KernelInjectable protocol."""
    from sagellm_backend import get_provider
    from sagellm_core.kernel_injection import KernelInjector, OperatorRegistry
    from sagellm_protocol import OperatorConstraint, OperatorType

    backend = get_provider("cpu")
    registry = OperatorRegistry()

    # Register a custom kernel
    def custom_linear_kernel(x, weight, bias):
        """Custom linear kernel for testing."""
        import torch

        return torch.nn.functional.linear(x, weight, bias)

    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=custom_linear_kernel,
        constraint=OperatorConstraint(device_type="cpu"),
        priority=10,
        description="Test linear kernel",
    )

    injector = KernelInjector(registry, backend)

    # Create a custom layer that implements KernelInjectable protocol
    class CustomLinearLayer:
        def __init__(self):
            self._backend = None
            self._kernel = None

        def set_backend(self, backend):
            self._backend = backend

        def set_kernel(self, kernel):
            self._kernel = kernel

    layer = CustomLinearLayer()

    # Inject kernel
    success = injector.inject_into_layer(layer, OperatorType.LINEAR, device_type="cpu")

    assert success, "Kernel injection should succeed"
    assert layer._kernel is not None, "Kernel should be injected"
    assert layer._backend is not None, "Backend should be injected"


def test_operator_registry_priority():
    """Test that operator registry respects kernel priority."""
    from sagellm_core.kernel_injection import OperatorRegistry
    from sagellm_protocol import OperatorConstraint, OperatorType

    registry = OperatorRegistry()

    # Register low-priority kernel
    def low_priority_kernel():
        return "low"

    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=low_priority_kernel,
        constraint=OperatorConstraint(device_type="cpu"),
        priority=5,
    )

    # Register high-priority kernel
    def high_priority_kernel():
        return "high"

    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=high_priority_kernel,
        constraint=OperatorConstraint(device_type="cpu"),
        priority=10,
    )

    # Query should return high-priority kernel
    spec = registry.query_kernel(OperatorType.LINEAR, device_type="cpu")
    assert spec is not None
    assert spec.priority == 10
    assert spec.kernel_impl() == "high"


@pytest.mark.asyncio
async def test_full_inference_with_kernel_mapping():
    """Test full inference flow with kernel mapping."""
    from sagellm_backend import get_provider
    from sagellm_core.worker.model_runner import ModelRunner

    backend = get_provider("cpu")
    runner = ModelRunner(
        backend=backend,
        model_path="sshleifer/tiny-gpt2",
        dtype="float32",
        trust_remote_code=True,
        max_new_tokens=5,
    )

    await runner.start()

    # Create a simple sequence group
    class SimpleSeqGroup:
        def __init__(self):
            self.request_id = "test-001"
            self.prompt = "Hello, world!"
            self.max_tokens = 5

    seq_group = SimpleSeqGroup()

    # Run inference
    outputs = await runner.forward([seq_group])

    assert len(outputs) == 1
    output = outputs[0]
    assert output.request_id == "test-001"
    assert output.output_text is not None
    assert len(output.output_token_ids) > 0

    print(f"\nGenerated text: {output.output_text}")
    print(f"Generated tokens: {output.output_token_ids}")

    # Check kernel call stats (should be > 0 if monkey patching worked)
    kernel_stats = runner.get_kernel_call_stats()
    print(f"\nKernel call stats: {kernel_stats}")

    # For standard HF models, monkey patching should have worked
    assert kernel_stats["total_calls"] > 0, "Backend kernels should have been called"

    await runner.stop()


if __name__ == "__main__":
    import asyncio

    print("Running operator mapping integration tests...")
    asyncio.run(test_model_runner_kernel_injection_setup())
    asyncio.run(test_kernel_injection_with_model())
    asyncio.run(test_backend_kernels_registration())
    asyncio.run(test_operator_mapping_with_custom_layer())
    test_operator_registry_priority()
    asyncio.run(test_full_inference_with_kernel_mapping())
    print("\n✓ All tests passed!")
