"""Unit tests for Operator Mapping Protocol.

Tests the logical-to-physical operator mapping types and validation.
"""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from sagellm_protocol import (
    DType,
    ErrorCode,
    KernelKind,
    LogicalOperator,
    OperatorConstraint,
    OperatorMapping,
    OperatorRegistryQuery,
    OperatorRegistrySnapshot,
    OperatorType,
    PhysicalOperator,
)
from sagellm_protocol.operator_mapping import OperatorRegistryEntry


class TestOperatorType:
    """Test OperatorType enum."""

    def test_operator_types_exist(self) -> None:
        """Test all expected operator types are defined."""
        assert OperatorType.LINEAR == "linear"
        assert OperatorType.ATTENTION == "attention"
        assert OperatorType.FLASH_ATTENTION == "flash_attention"
        assert OperatorType.LAYERNORM == "layernorm"
        assert OperatorType.RMSNORM == "rmsnorm"
        assert OperatorType.RELU == "relu"
        assert OperatorType.GELU == "gelu"
        assert OperatorType.SILU == "silu"
        assert OperatorType.SOFTMAX == "softmax"
        assert OperatorType.FUSED_QKV_PROJECTION == "fused_qkv_projection"
        assert OperatorType.FUSED_SILU_MUL == "fused_silu_mul"
        assert OperatorType.SAMPLE == "sample"
        assert OperatorType.CUSTOM == "custom"

    def test_operator_type_string_value(self) -> None:
        """Test operator types have correct string values."""
        op_type = OperatorType.LINEAR
        assert isinstance(op_type, str)
        assert op_type == "linear"


class TestOperatorConstraint:
    """Test OperatorConstraint validation."""

    def test_constraint_creation(self) -> None:
        """Test creating operator constraints."""
        constraint = OperatorConstraint(
            device_type="cuda",
            dtype=DType.FP16,
            min_compute_capability="8.0",
            kernel_kind=KernelKind.MATMUL,
        )
        assert constraint.device_type == "cuda"
        assert constraint.dtype == DType.FP16
        assert constraint.min_compute_capability == "8.0"
        assert constraint.kernel_kind == KernelKind.MATMUL

    def test_constraint_optional_fields(self) -> None:
        """Test constraints with optional fields."""
        constraint = OperatorConstraint(device_type="cuda")
        assert constraint.device_type == "cuda"
        assert constraint.dtype is None
        assert constraint.min_compute_capability is None

    def test_constraint_custom_fields(self) -> None:
        """Test custom constraint fields."""
        constraint = OperatorConstraint(
            device_type="cuda",
            custom_constraints={
                "min_memory_gb": 16,
                "requires_tensor_cores": True,
            },
        )
        assert constraint.custom_constraints["min_memory_gb"] == 16
        assert constraint.custom_constraints["requires_tensor_cores"] is True

    def test_constraint_batch_size_validation(self) -> None:
        """Test batch size constraints must be positive."""
        # Valid
        constraint = OperatorConstraint(min_batch_size=1, max_batch_size=32)
        assert constraint.min_batch_size == 1
        assert constraint.max_batch_size == 32

        # Invalid: batch size must be >= 1
        with pytest.raises(ValidationError):
            OperatorConstraint(min_batch_size=0)


class TestLogicalOperator:
    """Test LogicalOperator definition."""

    def test_linear_operator(self) -> None:
        """Test defining a linear operator."""
        op = LogicalOperator(
            op_type=OperatorType.LINEAR,
            name="fc1",
            input_shapes=[(1024,)],
            output_shape=(4096,),
            params={"in_features": 1024, "out_features": 4096},
        )
        assert op.op_type == OperatorType.LINEAR
        assert op.name == "fc1"
        assert op.input_shapes == [(1024,)]
        assert op.output_shape == (4096,)
        assert op.params["in_features"] == 1024

    def test_attention_operator(self) -> None:
        """Test defining an attention operator."""
        op = LogicalOperator(
            op_type=OperatorType.ATTENTION,
            name="self_attn",
            input_shapes=[(128, 1024), (128, 1024), (128, 1024)],
            output_shape=(128, 1024),
            params={
                "num_heads": 16,
                "head_dim": 64,
                "seq_len": 128,
            },
        )
        assert op.op_type == OperatorType.ATTENTION
        assert len(op.input_shapes) == 3  # Q, K, V
        assert op.params["num_heads"] == 16

    def test_operator_with_metadata(self) -> None:
        """Test operator with metadata."""
        op = LogicalOperator(
            op_type=OperatorType.LAYERNORM,
            name="ln1",
            input_shapes=[(1024,)],
            output_shape=(1024,),
            params={"normalized_shape": 1024, "eps": 1e-5},
            metadata={
                "layer_id": 0,
                "module_path": "model.layers.0.ln1",
            },
        )
        assert op.metadata["layer_id"] == 0
        assert op.metadata["module_path"] == "model.layers.0.ln1"


class TestPhysicalOperator:
    """Test PhysicalOperator definition."""

    def test_cuda_kernel(self) -> None:
        """Test defining a CUDA kernel."""
        kernel = PhysicalOperator(
            kernel_name="CUDALinearKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.MATMUL,
            priority=100,
        )
        assert kernel.kernel_name == "CUDALinearKernel"
        assert kernel.backend == "cuda"
        assert kernel.device == "cuda:0"
        assert kernel.dtype == DType.FP16
        assert kernel.kernel_kind == KernelKind.MATMUL
        assert kernel.priority == 100

    def test_ascend_kernel(self) -> None:
        """Test defining an Ascend NPU kernel."""
        kernel = PhysicalOperator(
            kernel_name="AscendAttentionKernel",
            backend="ascend",
            device="ascend:0",
            dtype=DType.BF16,
            kernel_kind=KernelKind.ATTENTION,
            priority=90,
        )
        assert kernel.backend == "ascend"
        assert kernel.device == "ascend:0"
        assert kernel.dtype == DType.BF16

    def test_kernel_with_constraints(self) -> None:
        """Test kernel with constraints."""
        kernel = PhysicalOperator(
            kernel_name="FlashAttentionV2",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.ATTENTION,
            priority=200,
            constraints=OperatorConstraint(
                device_type="cuda",
                min_compute_capability="8.0",
            ),
        )
        assert kernel.constraints is not None
        assert kernel.constraints.device_type == "cuda"
        assert kernel.constraints.min_compute_capability == "8.0"

    def test_kernel_default_priority(self) -> None:
        """Test kernel default priority is 0."""
        kernel = PhysicalOperator(
            kernel_name="DefaultKernel",
            backend="cpu",
            device="cpu",
            dtype=DType.FP32,
            kernel_kind=KernelKind.MATMUL,
        )
        assert kernel.priority == 0


class TestOperatorMapping:
    """Test OperatorMapping rules."""

    def test_basic_mapping(self) -> None:
        """Test creating a basic operator mapping."""
        mapping = OperatorMapping(
            logical_op_type=OperatorType.LINEAR,
            physical_kernel="CUDALinearKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            priority=100,
        )
        assert mapping.logical_op_type == OperatorType.LINEAR
        assert mapping.physical_kernel == "CUDALinearKernel"
        assert mapping.backend == "cuda"
        assert mapping.device == "cuda:0"
        assert mapping.dtype == DType.FP16
        assert mapping.priority == 100

    def test_mapping_with_constraints(self) -> None:
        """Test mapping with constraints."""
        mapping = OperatorMapping(
            logical_op_type=OperatorType.FLASH_ATTENTION,
            physical_kernel="FlashAttentionV2",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            constraints=OperatorConstraint(
                device_type="cuda",
                min_compute_capability="8.0",
                kernel_kind=KernelKind.ATTENTION,
            ),
            priority=200,
        )
        assert mapping.constraints is not None
        assert mapping.constraints.min_compute_capability == "8.0"

    def test_mapping_default_priority(self) -> None:
        """Test mapping default priority is 0."""
        mapping = OperatorMapping(
            logical_op_type=OperatorType.LINEAR,
            physical_kernel="CPULinear",
            backend="cpu",
            device="cpu",
            dtype=DType.FP32,
        )
        assert mapping.priority == 0


class TestOperatorRegistryQuery:
    """Test OperatorRegistryQuery."""

    def test_query_creation(self) -> None:
        """Test creating a registry query."""
        query = OperatorRegistryQuery(
            op_type=OperatorType.LINEAR,
            device="cuda:0",
            dtype=DType.FP16,
            backend="cuda",
        )
        assert query.op_type == OperatorType.LINEAR
        assert query.device == "cuda:0"
        assert query.dtype == DType.FP16
        assert query.backend == "cuda"

    def test_query_minimal(self) -> None:
        """Test minimal query with only op_type."""
        query = OperatorRegistryQuery(op_type=OperatorType.ATTENTION)
        assert query.op_type == OperatorType.ATTENTION
        assert query.device is None
        assert query.dtype is None
        assert query.backend is None

    def test_query_with_constraints(self) -> None:
        """Test query with constraints."""
        query = OperatorRegistryQuery(
            op_type=OperatorType.ATTENTION,
            device="cuda:0",
            constraints=OperatorConstraint(
                device_type="cuda",
                min_compute_capability="8.0",
            ),
        )
        assert query.constraints is not None
        assert query.constraints.device_type == "cuda"


class TestOperatorRegistrySnapshot:
    """Test OperatorRegistrySnapshot."""

    def test_snapshot_creation(self) -> None:
        """Test creating a registry snapshot."""
        snapshot = OperatorRegistrySnapshot(
            total_operators=10,
            operators_by_type={"linear": 3, "attention": 2},
            operators_by_backend={"cuda": 7, "ascend": 3},
        )
        assert snapshot.total_operators == 10
        assert snapshot.operators_by_type["linear"] == 3
        assert snapshot.operators_by_backend["cuda"] == 7

    def test_snapshot_with_operators(self) -> None:
        """Test snapshot with registered operators."""
        kernel = PhysicalOperator(
            kernel_name="TestKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.MATMUL,
        )
        mapping = OperatorMapping(
            logical_op_type=OperatorType.LINEAR,
            physical_kernel="TestKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
        )

        snapshot = OperatorRegistrySnapshot(
            total_operators=1,
            registered_operators=[kernel],
            active_mappings=[mapping],
        )
        assert len(snapshot.registered_operators) == 1
        assert len(snapshot.active_mappings) == 1
        assert snapshot.registered_operators[0].kernel_name == "TestKernel"


class TestOperatorRegistryEntry:
    """Test OperatorRegistryEntry."""

    def test_entry_creation(self) -> None:
        """Test creating a registry entry."""
        physical_op = PhysicalOperator(
            kernel_name="TestKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.MATMUL,
        )

        def dummy_kernel(*_args, **_kwargs):
            return "result"

        entry = OperatorRegistryEntry(
            physical_op=physical_op,
            kernel_callable=dummy_kernel,
            registration_metadata={"registered_at": "2026-02-15"},
        )

        assert entry.physical_op.kernel_name == "TestKernel"
        assert entry.kernel_callable is not None
        assert entry.kernel_callable() == "result"
        assert entry.registration_metadata["registered_at"] == "2026-02-15"

    def test_entry_without_callable(self) -> None:
        """Test entry can be created without callable."""
        physical_op = PhysicalOperator(
            kernel_name="TestKernel",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.MATMUL,
        )

        entry = OperatorRegistryEntry(physical_op=physical_op)
        assert entry.physical_op.kernel_name == "TestKernel"
        assert entry.kernel_callable is None
        assert entry.registration_metadata == {}


class TestErrorCodes:
    """Test operator mapping error codes."""

    def test_operator_error_codes_exist(self) -> None:
        """Test operator-related error codes are defined."""
        assert ErrorCode.OPERATOR_NOT_FOUND == "operator_not_found"
        assert ErrorCode.INVALID_OPERATOR_MAPPING == "invalid_operator_mapping"
        assert ErrorCode.OPERATOR_REGISTRATION_FAILED == "operator_registration_failed"


class TestIntegrationScenarios:
    """Integration tests for real-world scenarios."""

    def test_multi_backend_scenario(self) -> None:
        """Test scenario with multiple backends."""
        # Define logical operator
        logical_op = LogicalOperator(
            op_type=OperatorType.LINEAR,
            name="fc1",
            input_shapes=[(1024,)],
            output_shape=(4096,),
            params={"in_features": 1024, "out_features": 4096},
        )

        # Define physical implementations
        cuda_kernel = PhysicalOperator(
            kernel_name="CUDALinear",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.MATMUL,
            priority=100,
        )

        ascend_kernel = PhysicalOperator(
            kernel_name="AscendLinear",
            backend="ascend",
            device="ascend:0",
            dtype=DType.BF16,
            kernel_kind=KernelKind.MATMUL,
            priority=90,
        )

        cpu_kernel = PhysicalOperator(
            kernel_name="CPULinear",
            backend="cpu",
            device="cpu",
            dtype=DType.FP32,
            kernel_kind=KernelKind.MATMUL,
            priority=50,
        )

        # Create mappings
        cuda_mapping = OperatorMapping(
            logical_op_type=logical_op.op_type,
            physical_kernel=cuda_kernel.kernel_name,
            backend=cuda_kernel.backend,
            device=cuda_kernel.device,
            dtype=cuda_kernel.dtype,
            priority=cuda_kernel.priority,
        )

        # Verify priority ordering
        assert cuda_kernel.priority > ascend_kernel.priority > cpu_kernel.priority
        assert cuda_mapping.logical_op_type == logical_op.op_type

    def test_fused_operator_scenario(self) -> None:
        """Test scenario with fused operators."""
        # Fused operator: SiLU + Multiply (common in LLaMA)
        fused_op = LogicalOperator(
            op_type=OperatorType.FUSED_SILU_MUL,
            name="mlp.act_mul",
            input_shapes=[(4096,), (4096,)],
            output_shape=(4096,),
            params={"activation": "silu", "fusion": "multiply"},
        )

        fused_kernel = PhysicalOperator(
            kernel_name="FusedSiLUMul",
            backend="cuda",
            device="cuda:0",
            dtype=DType.FP16,
            kernel_kind=KernelKind.FUSION,
            priority=150,  # Higher priority for optimized fused kernel
        )

        mapping = OperatorMapping(
            logical_op_type=fused_op.op_type,
            physical_kernel=fused_kernel.kernel_name,
            backend=fused_kernel.backend,
            device=fused_kernel.device,
            dtype=fused_kernel.dtype,
            priority=fused_kernel.priority,
        )

        assert mapping.logical_op_type == OperatorType.FUSED_SILU_MUL
        assert fused_kernel.kernel_kind == KernelKind.FUSION
        assert fused_kernel.priority > 100  # Higher than normal kernels
