"""sageLLM Backend Provider 抽象。

本包提供硬件后端抽象层：
- BackendProvider 接口定义
- Stream/Event/KVBlock 抽象
- get_provider() 工厂函数 - 推荐的获取 Provider 方式
- Kernel 接口和实现（Linear, Embedding, RMSNorm, SiLU 等）
- AttentionBackend 接口和实现
- CPU Backend 实现
- CUDA Backend 实现 (PyTorch)
- Ascend Backend 实现 (PyTorch-NPU)

推荐使用方式 (vLLM v1 style):
    from sagellm_backend import get_provider

    # 自动选择最佳可用后端
    provider = get_provider()

    # 显式指定后端
    provider = get_provider("cuda")
    provider = get_provider("cpu")
    provider = get_provider("ascend")

    # 获取 Kernel
    linear_kernel = provider.get_kernel("linear")
    output = linear_kernel(input, weight)

    # 获取 Attention Backend
    attn_backend = provider.get_attention_backend()
    output = attn_backend.forward(query, key, value, attn_metadata)

能力描述符（DType, KernelKind, CapabilityDescriptor）已迁移至 sagellm-protocol。
请直接从 sagellm_protocol 导入这些类型。

Note: Engine implementations have been moved to sagellm-core.
      This package now focuses solely on hardware abstraction.
"""

from __future__ import annotations

__version__ = "0.5.2.0"

from sagellm_protocol import CapabilityDescriptor, DType, KernelKind

# Attention backend interfaces
from sagellm_backend.attention.base import AttentionBackend, AttentionMetadata
from sagellm_backend.attention.cpu import CPUAttentionBackend
from sagellm_backend.base import BackendProvider, Event, KVBlock, Stream
from sagellm_backend.kernels.activation import CPUSiLUKernel, GELUKernel, ReLUKernel, SiLUKernel

# Kernel interfaces
from sagellm_backend.kernels.base import Kernel, KernelRegistry, RegisteredKernel
from sagellm_backend.kernels.embedding import CPUEmbeddingKernel, EmbeddingKernel
from sagellm_backend.kernels.linear import CPULinearKernel, LinearKernel
from sagellm_backend.kernels.normalization import CPURMSNormKernel, LayerNormKernel, RMSNormKernel

# Memory management
from sagellm_backend.memory import (
    MemoryAllocator,
    MemoryHandle,
    MemoryPool,
    pin_memory,
    unpin_memory,
)
from sagellm_backend.providers.cpu import CPUBackendProvider, create_cpu_backend

# Provider Registry and factory function (vLLM v1 style)
from sagellm_backend.registry import (
    get_provider,
    is_provider_available,
    list_providers,
    register_provider,
)

# Optional PyTorch-based providers (lazy import)
try:
    from sagellm_backend.providers.cuda import CUDABackendProvider, create_cuda_backend

    _CUDA_AVAILABLE = True
except ImportError:
    _CUDA_AVAILABLE = False
    CUDABackendProvider = None  # type: ignore
    create_cuda_backend = None  # type: ignore

try:
    from sagellm_backend.providers.ascend import AscendBackendProvider, create_ascend_backend

    _ASCEND_AVAILABLE = True
except ImportError:
    _ASCEND_AVAILABLE = False
    AscendBackendProvider = None  # type: ignore
    create_ascend_backend = None  # type: ignore

# Domestic accelerator providers (lazy import)
try:
    from sagellm_backend.providers.kunlun import KunlunBackendProvider

    _KUNLUN_AVAILABLE = True
except ImportError:
    _KUNLUN_AVAILABLE = False
    KunlunBackendProvider = None  # type: ignore

try:
    from sagellm_backend.providers.dcu import DCUBackendProvider

    _DCU_AVAILABLE = True
except ImportError:
    _DCU_AVAILABLE = False
    DCUBackendProvider = None  # type: ignore

try:
    from sagellm_backend.providers.mthreads import MThreadsBackendProvider

    _MTHREADS_AVAILABLE = True
except ImportError:
    _MTHREADS_AVAILABLE = False
    MThreadsBackendProvider = None  # type: ignore

__version__ = "0.4.0.14"

__all__ = [
    # Version
    "__version__",
    # =========================================================================
    # Provider Registry (vLLM v1 style) - RECOMMENDED
    # =========================================================================
    "get_provider",  # Factory function to get BackendProvider
    "register_provider",  # Register custom provider
    "list_providers",  # List available providers
    "is_provider_available",  # Check if provider is available
    # =========================================================================
    # Provider interface
    # =========================================================================
    "BackendProvider",
    "Stream",
    "Event",
    "KVBlock",
    # Capability types (re-exported from protocol for convenience)
    "CapabilityDescriptor",
    "DType",
    "KernelKind",
    # =========================================================================
    # Kernel interfaces
    # =========================================================================
    "Kernel",
    "KernelRegistry",
    "RegisteredKernel",
    "LinearKernel",
    "CPULinearKernel",
    "EmbeddingKernel",
    "CPUEmbeddingKernel",
    "RMSNormKernel",
    "CPURMSNormKernel",
    "LayerNormKernel",
    "SiLUKernel",
    "CPUSiLUKernel",
    "GELUKernel",
    "ReLUKernel",
    # =========================================================================
    # Attention backend interfaces
    # =========================================================================
    "AttentionBackend",
    "AttentionMetadata",
    "CPUAttentionBackend",
    # =========================================================================
    # Memory Management
    # =========================================================================
    "MemoryAllocator",
    "MemoryHandle",
    "MemoryPool",
    "pin_memory",
    "unpin_memory",
    # =========================================================================
    # Provider implementations
    # =========================================================================
    # CPU backend implementation
    "CPUBackendProvider",
    "create_cpu_backend",
    # CUDA backend (optional)
    "CUDABackendProvider",
    "create_cuda_backend",
    # Ascend backend (optional)
    "AscendBackendProvider",
    "create_ascend_backend",
]
