"""Mock Backend 单元测试"""

from __future__ import annotations

import pytest

from sagellm_backend import (
    CapabilityDescriptor,
    DType,
    KernelKind,
    MockBackendProvider,
    create_mock_backend,
)


def test_mock_backend_default_capabilities():
    """测试默认能力矩阵"""
    backend = MockBackendProvider()
    cap = backend.capability()

    assert isinstance(cap, CapabilityDescriptor)
    assert DType.FP16 in cap.supported_dtypes
    assert DType.BF16 in cap.supported_dtypes
    assert KernelKind.MATMUL in cap.supported_kernels
    assert KernelKind.ATTENTION in cap.supported_kernels
    assert cap.has_stream is True
    assert cap.has_event is True
    assert cap.has_collective is True
    assert cap.has_graph_exec is False
    assert cap.max_memory_mb == 16384


def test_mock_backend_custom_capabilities():
    """测试自定义能力矩阵"""
    backend = MockBackendProvider(
        supported_dtypes=[DType.INT8, DType.FP8],
        supported_kernels=[KernelKind.FUSION],
        has_collective=False,
        has_graph_exec=True,
        max_memory_mb=8192,
    )
    cap = backend.capability()

    assert cap.supported_dtypes == [DType.INT8, DType.FP8]
    assert cap.supported_kernels == [KernelKind.FUSION]
    assert cap.has_collective is False
    assert cap.has_graph_exec is True
    assert cap.max_memory_mb == 8192


def test_kv_block_alloc_success():
    """测试 KV 块分配成功"""
    backend = MockBackendProvider()
    block = backend.kv_block_alloc(100, DType.FP16)

    assert block.handle > 0
    assert block.num_tokens == 100
    assert block.dtype == DType.FP16
    assert block.device == "mock:0"


def test_kv_block_alloc_unsupported_dtype():
    """测试不支持的 dtype 分配失败"""
    backend = MockBackendProvider(supported_dtypes=[DType.FP16])

    with pytest.raises(ValueError, match="not supported"):
        backend.kv_block_alloc(100, DType.INT4)


def test_kv_block_free():
    """测试 KV 块释放"""
    backend = MockBackendProvider()
    block = backend.kv_block_alloc(100, DType.FP16)

    # 释放不应抛出异常
    backend.kv_block_free(block)
    # 重复释放也不应抛出异常
    backend.kv_block_free(block)


def test_kv_block_migrate():
    """测试 KV 块迁移"""
    backend = MockBackendProvider()
    block = backend.kv_block_alloc(100, DType.FP16)

    new_block = backend.kv_block_migrate(block, "mock:1")

    assert new_block.handle == block.handle
    assert new_block.num_tokens == block.num_tokens
    assert new_block.dtype == block.dtype
    assert new_block.device == "mock:1"


def test_stream_event():
    """测试 stream 和 event"""
    backend = MockBackendProvider()
    stream = backend.create_stream()
    event1 = backend.create_event()
    event2 = backend.create_event()

    # 记录事件
    stream.record_event(event1)
    stream.record_event(event2)

    # 同步
    stream.synchronize()
    event1.synchronize()
    event2.synchronize()

    # 计算耗时
    elapsed = event2.elapsed_ms(event1)
    assert elapsed >= 0


def test_collective_operations_supported():
    """测试集合操作（支持）"""
    backend = MockBackendProvider(has_collective=True)

    tensor = [1, 2, 3]
    # 不应抛出异常
    backend.all_reduce(tensor, "sum")
    result = backend.all_gather(tensor)
    assert result == tensor


def test_collective_operations_unsupported():
    """测试集合操作（不支持）"""
    backend = MockBackendProvider(has_collective=False)

    tensor = [1, 2, 3]

    with pytest.raises(NotImplementedError, match="collective"):
        backend.all_reduce(tensor, "sum")

    with pytest.raises(NotImplementedError, match="collective"):
        backend.all_gather(tensor)


"""Tests removed: mock providers are forbidden."""
    """测试 kernel 注册"""
