"""Attention backend interfaces and implementations.

This module provides:
- AttentionBackend: Base class for attention implementations
- AttentionMetadata: Metadata for attention computation
- AttentionType: Enum for attention types
- CPUAttentionBackend: CPU-based attention
- FlashAttentionBackend: Flash Attention v2 (requires flash-attn)
- PagedAttentionBackend: Paged attention for KV cache

Registry functions:
- get_attention_backend: Get backend by name
- register_attention_backend: Register custom backend
- list_attention_backends: List available backends

The attention backend abstracts hardware-specific attention implementations:
- Flash Attention (CUDA)
- PagedAttention (CUDA, vLLM style)
- Vanilla attention (CPU)
- ACL attention (Ascend)

Example:
    from sagellm_backend.attention import get_attention_backend, AttentionMetadata

    # Get best available backend
    backend = get_attention_backend("flash")

    # Create metadata for prefill
    metadata = AttentionMetadata.for_prefill(seq_lens=[128, 256])

    # Run attention
    output = backend.forward(query, key, value, metadata)
"""

from __future__ import annotations

from sagellm_backend.attention.base import (
    AttentionBackend,
    AttentionMetadata,
    AttentionType,
)
from sagellm_backend.attention.cpu import CPUAttentionBackend
from sagellm_backend.attention.registry import (
    attention_backend,
    get_attention_backend,
    is_attention_backend_available,
    list_attention_backends,
    register_attention_backend,
    unregister_attention_backend,
)

__all__ = [
    # Base classes
    "AttentionBackend",
    "AttentionMetadata",
    "AttentionType",
    # CPU backend (always available)
    "CPUAttentionBackend",
    # Registry functions
    "get_attention_backend",
    "register_attention_backend",
    "unregister_attention_backend",
    "list_attention_backends",
    "is_attention_backend_available",
    "attention_backend",
]

# Conditionally export FlashAttention if available
try:
    from sagellm_backend.attention.flash_attention import (
        FlashAttentionBackend,
        FlashAttentionMetadataBuilder,
        get_flash_attn_version,
        is_flash_attn_available,
    )

    __all__.extend(
        [
            "FlashAttentionBackend",
            "FlashAttentionMetadataBuilder",
            "is_flash_attn_available",
            "get_flash_attn_version",
        ]
    )
except ImportError:
    pass

# Conditionally export PagedAttention if available
try:
    from sagellm_backend.attention.paged_attention import (
        PagedAttentionBackend,
        PagedAttentionMetadataBuilder,
        is_vllm_available,
    )

    __all__.extend(
        [
            "PagedAttentionBackend",
            "PagedAttentionMetadataBuilder",
            "is_vllm_available",
        ]
    )
except ImportError:
    pass
