"""sagellm-kv-cache: KV Cache Management Module for sageLLM.

本模块提供 KV Cache 管理功能，包括：
- 本地模型类：KVHandle 及其相关类型
- 错误体系：KVCacheError 及其子类
- MVP: KV Transfer Engine（Task1.3）
- Prefix Cache: Task2.1 ✅
- KV Pool: Task2.2 ✅
- Eviction Policy: Task2.3 ✅
- Lifetime Predictor: Task2.5 ✅ (CPU-only, 统计模型)
- Scheduler Bridge: Task2.6/2.7 ✅

使用示例：
    >>> from sagellm_kv_cache import KVHandle, KVBudgetExceededError, KVTransferEngine
    >>> from sagellm_kv_cache import PrefixCache, KVPool, EvictionManager, SchedulerBridge
    >>> from sagellm_kv_cache import HistoricalPredictor, PredictorFeatures
    >>> handle = KVHandle.create(num_tokens=128, device="cpu")
    >>> print(handle.to_dict())
"""

from __future__ import annotations

from sagellm_kv_cache._version import __version__

# ============================================================
# Task 2.1: Prefix Cache
# ============================================================
# ============================================================
# 本地模型类 (Agent 2)
# ============================================================
# ============================================================
# 错误体系 (Agent 2)
# ============================================================
# ============================================================
# 分布式池化接口（预留，未来实现）
# ============================================================
from sagellm_kv_cache.distributed import (
    GlobalKVPoolManager,  # 全局池管理器抽象接口
    NodeContribution,  # 节点内存贡献配置
    RemoteKVAccessor,  # 远程 KV 访问层接口
)
from sagellm_kv_cache.errors import (
    KVAllPinnedError,
    # 预算/资源错误
    KVBudgetExceededError,
    # 基类
    KVCacheError,
    # 缓存错误
    KVCacheMissError,
    KVChecksumMismatchError,
    KVConfigInvalidError,
    # 配置错误
    KVConfigMissingError,
    KVDoubleFreeError,
    KVEvictionFailedError,
    KVHandleNotFoundError,
    KVHandlePinnedError,
    # 句柄操作错误
    KVInvalidHandleError,
    KVMemoryExhaustedError,
    KVMetadataMismatchError,
    KVNotImplementedError,
    # 驱逐错误
    KVNoVictimError,
    KVPoolFullError,
    KVPrefixTooLongError,
    KVRefCountUnderflowError,
    # 传输错误
    KVTransferFailedError,
    KVTransferTimeoutError,
)

# ============================================================
# Task 2.3: Eviction Policy
# ============================================================
from sagellm_kv_cache.eviction import EvictionManager, EvictionStrategy, FIFOEviction, LRUEviction
from sagellm_kv_cache.kv_cache_manager import KVCacheManager
from sagellm_kv_cache.load_aware_scheduler import LoadAwareScheduler, PlacementStrategy
from sagellm_kv_cache.load_metrics import DeviceLoadMetrics, LoadMetricsCollector
from sagellm_kv_cache.models import (
    DType,
    EvictionCandidate,
    EvictionPolicy,
    KVBlock,
    KVBlockState,
    KVHandle,
    KVPoolStats,
    KVTransferMetadata,
    Layout,
    LifetimePrediction,
    MemoryTier,
    PrefixCacheEntry,
    SchedulerPlan,
    SchedulerRequest,
)

# ============================================================
# Task 2.2: KV Pool
# ============================================================
from sagellm_kv_cache.pool.kv_pool import KVPool
from sagellm_kv_cache.predictor import (
    FeatureExtractor,
    # ============================================================
    # Task 2.5: Lifetime Predictor
    # ============================================================
    HistoricalPredictor,
    LifetimePredictor,
    PredictorFeatures,
)
from sagellm_kv_cache.prefix import PrefixIndex
from sagellm_kv_cache.prefix_cache import PrefixCache

# ============================================================
# Task 2.6/2.7: Scheduler Bridge
# ============================================================
from sagellm_kv_cache.scheduler_bridge import SchedulerBridge

# ============================================================
# MVP: KV Transfer Engine (Task1.3)
# ============================================================
from sagellm_kv_cache.transfer import (
    CommTransferBackendAdapter,
    KVTransferEngine,
    SocketTransferBackend,
)

__all__ = [
    "__version__",
    # === Task 2.1: Prefix Cache ===
    "PrefixCache",
    "PrefixIndex",
    # === Task 2.2: KV Pool ===
    "KVPool",
    # === Task 2.3: Eviction Policy ===
    "EvictionManager",
    "EvictionStrategy",
    "LRUEviction",
    "FIFOEviction",
    # === Task 2.5: Lifetime Predictor ===
    "LifetimePredictor",
    "HistoricalPredictor",
    "PredictorFeatures",
    "FeatureExtractor",
    # === Task 2.6/2.7: Scheduler Bridge ===
    "SchedulerBridge",
    "LoadAwareScheduler",
    "PlacementStrategy",
    "DeviceLoadMetrics",
    "LoadMetricsCollector",
    "KVCacheManager",
    # === 本地模型类 ===
    "KVHandle",
    "KVBlock",
    "KVTransferMetadata",
    "PrefixCacheEntry",
    "EvictionCandidate",
    "SchedulerRequest",
    "SchedulerPlan",
    "LifetimePrediction",
    "KVPoolStats",
    # === MVP: KV Transfer Engine ===
    "KVTransferEngine",
    "CommTransferBackendAdapter",
    "SocketTransferBackend",
    # === 分布式池化接口（预留） ===
    "GlobalKVPoolManager",
    "NodeContribution",
    "RemoteKVAccessor",
    # === 枚举类型 ===
    "KVBlockState",
    "DType",
    "Layout",
    "MemoryTier",
    "EvictionPolicy",
    # === 错误基类 ===
    "KVCacheError",
    # === 预算/资源错误 ===
    "KVBudgetExceededError",
    "KVMemoryExhaustedError",
    "KVPoolFullError",
    # === 句柄操作错误 ===
    "KVInvalidHandleError",
    "KVDoubleFreeError",
    "KVHandlePinnedError",
    "KVHandleNotFoundError",
    "KVRefCountUnderflowError",
    # === 缓存错误 ===
    "KVCacheMissError",
    "KVPrefixTooLongError",
    # === 驱逐错误 ===
    "KVNoVictimError",
    "KVEvictionFailedError",
    "KVAllPinnedError",
    # === 传输错误 ===
    "KVTransferFailedError",
    "KVTransferTimeoutError",
    "KVChecksumMismatchError",
    "KVMetadataMismatchError",
    # === 配置错误 ===
    "KVConfigMissingError",
    "KVConfigInvalidError",
    "KVNotImplementedError",
]
