//! 统一 GPU 设备抽象层 — Phase-61
//!
//! 所有 GPU/计算功能共享此单一设备检测与选择层。
//! 取代原有三套独立的 GPU 检测路径：
//!
//! 1. `gpu::device::probe_gpu_availability()` — 索引构建 (gpu-index feature)
//! 2. `CandleEmbeddingEngine::get_device()` — Embedding 推理 (candle feature)
//! 3. `ComputeBackend::select_backend() → CandleCudaBackend::try_new()` — SiFu 矩阵运算
//!
//! ## 设计原则
//!
//! - **统一入口**：所有 GPU 功能通过 `DeviceSelector` 获取设备
//! - **双层安全降级**：编译期 feature gate + 运行时 GPU 探测
//! - **OnceLock 缓存**：全进程仅探测一次，零重复开销
//! - **零侵入**：无 GPU 时所有功能正常运行在 CPU 上
//!
//! ## 使用方式
//!
//! ```ignore
//! use aifastdb_core::device::{DeviceSelector, GpuAccelConfig};
//!
//! // 默认 auto 模式
//! let selector = DeviceSelector::new(GpuAccelConfig::default());
//!
//! // 获取 candle Device（自动选择最优设备）
//! let device = selector.candle_device();
//!
//! // 判断给定工作负载是否应使用 GPU
//! let hint = WorkloadHint::new(100_000, 768);
//! if selector.should_use_gpu(&hint) {
//!     // GPU 加速路径
//! } else {
//!     // CPU 路径
//! }
//! ```

use std::fmt;
use std::sync::OnceLock;

// ============================================================================
// DevicePreference — 用户设备偏好
// ============================================================================

/// 设备偏好选择
///
/// 控制 `DeviceSelector` 如何选择计算设备。
/// `Auto` 是推荐默认值，会按 CUDA → Metal → CPU 顺序尝试。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DevicePreference {
    /// 自动检测最优设备：CUDA → Metal → CPU
    ///
    /// **推荐默认值。** 在无 GPU 环境下自动降级到 CPU，零配置。
    Auto,

    /// 指定 CUDA GPU 设备（按 ordinal 编号）
    ///
    /// `Cuda(0)` = 第一块 GPU，`Cuda(1)` = 第二块 GPU
    /// 如果指定的设备不可用，会自动 fallback 到 CPU 并记录警告。
    Cuda(usize),

    /// 指定 Apple Metal GPU
    ///
    /// 仅在 macOS/iOS 上可用。不可用时 fallback 到 CPU。
    Metal,

    /// 强制使用 CPU
    ///
    /// 即使 GPU 可用也不使用。适用于调试或资源受限场景。
    Cpu,
}

impl Default for DevicePreference {
    fn default() -> Self {
        DevicePreference::Auto
    }
}

impl DevicePreference {
    /// 从字符串解析设备偏好
    ///
    /// 支持格式：
    /// - `"auto"` → `Auto`
    /// - `"cpu"` → `Cpu`
    /// - `"cuda"` / `"cuda:0"` → `Cuda(0)`
    /// - `"cuda:1"` → `Cuda(1)`
    /// - `"metal"` → `Metal`
    pub fn from_str_flexible(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "auto" => DevicePreference::Auto,
            "cpu" => DevicePreference::Cpu,
            "cuda" | "cuda:0" | "gpu" => DevicePreference::Cuda(0),
            "metal" => DevicePreference::Metal,
            s if s.starts_with("cuda:") => {
                if let Ok(ordinal) = s[5..].parse::<usize>() {
                    DevicePreference::Cuda(ordinal)
                } else {
                    tracing::warn!(
                        "Invalid CUDA device ordinal '{}', falling back to Auto",
                        s
                    );
                    DevicePreference::Auto
                }
            }
            _ => {
                tracing::warn!("Unknown device preference '{}', falling back to Auto", s);
                DevicePreference::Auto
            }
        }
    }
}

impl fmt::Display for DevicePreference {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DevicePreference::Auto => write!(f, "auto"),
            DevicePreference::Cuda(ordinal) => write!(f, "cuda:{ordinal}"),
            DevicePreference::Metal => write!(f, "metal"),
            DevicePreference::Cpu => write!(f, "cpu"),
        }
    }
}

// ============================================================================
// GpuAccelConfig — 全局 GPU 加速配置
// ============================================================================

/// 全局 GPU 加速策略配置
///
/// 控制所有 GPU 功能的统一行为：设备选择、内存预算、降级策略、日志等。
/// 所有字段均有合理默认值，最简用法是 `GpuAccelConfig::default()`。
#[derive(Debug, Clone)]
pub struct GpuAccelConfig {
    /// 设备偏好（默认 Auto）
    pub device_preference: DevicePreference,

    /// GPU 内存预算（字节），`None` = 自动使用 80% 可用显存
    pub memory_budget: Option<usize>,

    /// 是否在 GPU→CPU 降级时记录 tracing 日志（默认 true）
    pub log_fallback: bool,

    /// Auto 模式下触发 GPU 加速的最小数据量（字节）
    ///
    /// 低于此阈值时直接使用 CPU（避免 GPU 传输开销）。
    /// 默认 1MB — 小于 1MB 的数据在 CPU 上通常更快。
    pub auto_min_data_bytes: usize,

    /// Auto 模式下触发 GPU 加速的最小向量数
    ///
    /// 默认 1000 — 少于 1000 个向量时 GPU 传输开销不划算。
    pub auto_min_vectors: usize,
}

impl Default for GpuAccelConfig {
    fn default() -> Self {
        Self {
            device_preference: DevicePreference::Auto,
            memory_budget: None,
            log_fallback: true,
            auto_min_data_bytes: 1024 * 1024, // 1MB
            auto_min_vectors: 1_000,
        }
    }
}

impl GpuAccelConfig {
    /// 创建纯 CPU 配置（禁用所有 GPU 加速）
    pub fn cpu_only() -> Self {
        Self {
            device_preference: DevicePreference::Cpu,
            ..Default::default()
        }
    }

    /// 创建指定 CUDA 设备的配置
    pub fn cuda(ordinal: usize) -> Self {
        Self {
            device_preference: DevicePreference::Cuda(ordinal),
            ..Default::default()
        }
    }

    /// 设置内存预算（MB）
    pub fn with_memory_budget_mb(mut self, mb: usize) -> Self {
        self.memory_budget = Some(mb * 1024 * 1024);
        self
    }
}

// ============================================================================
// WorkloadHint — 工作负载描述
// ============================================================================

/// 工作负载描述，用于 `should_use_gpu()` 决策
///
/// 提供足够信息让 DeviceSelector 判断 GPU 加速是否有价值。
#[derive(Debug, Clone)]
pub struct WorkloadHint {
    /// 数据项数量（如向量数、样本数、batch size）
    pub num_items: usize,

    /// 每项数据的维度（如 embedding dimension）
    pub dimension: usize,

    /// 操作类型标签（用于日志，如 "kmeans", "brute_force_search"）
    pub operation: &'static str,
}

impl WorkloadHint {
    /// 创建工作负载描述
    pub fn new(num_items: usize, dimension: usize) -> Self {
        Self {
            num_items,
            dimension,
            operation: "unknown",
        }
    }

    /// 创建带操作标签的工作负载描述
    pub fn with_operation(num_items: usize, dimension: usize, operation: &'static str) -> Self {
        Self {
            num_items,
            dimension,
            operation,
        }
    }

    /// 估算数据量（字节）
    pub fn estimated_data_bytes(&self) -> usize {
        self.num_items * self.dimension * std::mem::size_of::<f32>()
    }
}

// ============================================================================
// DeviceAvailability — 设备可用性探测结果
// ============================================================================

/// 设备可用性探测结果
#[derive(Debug, Clone)]
pub enum DeviceAvailability {
    /// CUDA GPU 可用
    CudaAvailable {
        /// 设备编号
        ordinal: usize,
        /// 设备名称
        name: String,
        /// 可用内存（字节）
        free_memory_bytes: usize,
        /// 总内存（字节）
        total_memory_bytes: usize,
    },

    /// Metal GPU 可用 (macOS)
    MetalAvailable,

    /// 仅 CPU 可用（附带原因）
    CpuOnly(String),
}

impl DeviceAvailability {
    /// GPU 是否可用
    pub fn has_gpu(&self) -> bool {
        !matches!(self, DeviceAvailability::CpuOnly(_))
    }
}

impl fmt::Display for DeviceAvailability {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DeviceAvailability::CudaAvailable {
                ordinal,
                name,
                free_memory_bytes,
                ..
            } => {
                let free_mb = *free_memory_bytes as f64 / (1024.0 * 1024.0);
                write!(f, "CUDA:{ordinal} ({name}, {free_mb:.0}MB free)")
            }
            DeviceAvailability::MetalAvailable => write!(f, "Metal GPU"),
            DeviceAvailability::CpuOnly(reason) => write!(f, "CPU only ({reason})"),
        }
    }
}

// ============================================================================
// GpuFallbackEvent — 降级事件（用于日志和指标）
// ============================================================================

/// GPU 降级事件
///
/// 当 GPU 不可用或操作在 GPU 上失败时记录此事件。
/// 用于 tracing 日志和可选的 metrics 上报。
#[derive(Debug, Clone)]
pub struct GpuFallbackEvent {
    /// 降级原因分类
    pub reason: FallbackReason,
    /// 详细描述
    pub detail: String,
    /// 触发降级的操作
    pub operation: String,
    /// 原始请求的设备偏好
    pub requested_device: DevicePreference,
}

/// 降级原因分类
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FallbackReason {
    /// 编译时未启用 GPU feature
    FeatureNotEnabled,
    /// CUDA 驱动未安装
    DriverNotFound,
    /// 没有兼容的 GPU 设备
    NoCompatibleDevice,
    /// GPU 显存不足
    InsufficientMemory,
    /// 指定设备不存在
    DeviceNotFound,
    /// GPU 运算出错，回退到 CPU
    RuntimeError,
    /// 数据量太小，不值得用 GPU
    DataTooSmall,
    /// 用户显式选择 CPU
    UserRequestedCpu,
}

impl fmt::Display for FallbackReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FallbackReason::FeatureNotEnabled => write!(f, "GPU feature not enabled at compile time"),
            FallbackReason::DriverNotFound => write!(f, "CUDA driver not found"),
            FallbackReason::NoCompatibleDevice => write!(f, "No compatible GPU device"),
            FallbackReason::InsufficientMemory => write!(f, "Insufficient GPU memory"),
            FallbackReason::DeviceNotFound => write!(f, "Specified GPU device not found"),
            FallbackReason::RuntimeError => write!(f, "GPU runtime error"),
            FallbackReason::DataTooSmall => write!(f, "Data too small for GPU overhead"),
            FallbackReason::UserRequestedCpu => write!(f, "CPU explicitly requested by user"),
        }
    }
}

impl GpuFallbackEvent {
    /// 记录降级事件到 tracing
    pub fn log(&self) {
        match self.reason {
            FallbackReason::UserRequestedCpu | FallbackReason::DataTooSmall => {
                tracing::debug!(
                    operation = %self.operation,
                    reason = %self.reason,
                    detail = %self.detail,
                    "GPU fallback to CPU (expected)"
                );
            }
            FallbackReason::FeatureNotEnabled => {
                tracing::debug!(
                    operation = %self.operation,
                    reason = %self.reason,
                    "GPU fallback to CPU (feature not compiled)"
                );
            }
            _ => {
                tracing::warn!(
                    operation = %self.operation,
                    reason = %self.reason,
                    detail = %self.detail,
                    requested_device = %self.requested_device,
                    "GPU fallback to CPU"
                );
            }
        }
    }
}

// ============================================================================
// DeviceSelector — 统一设备选择器
// ============================================================================

/// 全进程缓存的设备可用性
static DEVICE_AVAILABILITY_CACHE: OnceLock<DeviceAvailability> = OnceLock::new();

/// 全进程降级事件计数器
static FALLBACK_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// 统一设备选择器 — 所有 GPU 功能的入口
///
/// 负责：
/// - 探测 GPU 可用性（OnceLock 缓存，全进程仅一次）
/// - 根据配置和工作负载选择最优设备
/// - 返回 `candle_core::Device` 供推理/计算使用
/// - 在 GPU 不可用时优雅降级到 CPU 并记录日志
///
/// ## 线程安全
///
/// `DeviceSelector` 是 `Send + Sync` 的，可安全跨线程共享。
/// 内部 GPU 探测使用 `OnceLock`，保证仅执行一次。
#[derive(Debug, Clone)]
pub struct DeviceSelector {
    config: GpuAccelConfig,
}

impl DeviceSelector {
    /// 创建设备选择器
    pub fn new(config: GpuAccelConfig) -> Self {
        Self { config }
    }

    /// 使用默认配置（auto 模式）创建
    pub fn auto() -> Self {
        Self::new(GpuAccelConfig::default())
    }

    /// 获取当前配置
    pub fn config(&self) -> &GpuAccelConfig {
        &self.config
    }

    /// 获取设备可用性信息（OnceLock 缓存）
    pub fn availability(&self) -> &DeviceAvailability {
        DEVICE_AVAILABILITY_CACHE.get_or_init(|| probe_device_availability())
    }

    /// 获取 candle Device
    ///
    /// 根据 `DevicePreference` 和实际硬件可用性返回最优设备。
    /// 如果请求的 GPU 不可用，自动 fallback 到 CPU。
    ///
    /// # 行为
    ///
    /// | DevicePreference | GPU 可用 | GPU 不可用 |
    /// |---|---|---|
    /// | Auto | 返回 GPU Device | 返回 CPU Device |
    /// | Cuda(n) | 返回 cuda:n Device | fallback CPU + warn |
    /// | Metal | 返回 Metal Device | fallback CPU + warn |
    /// | Cpu | 返回 CPU Device | 返回 CPU Device |
    #[cfg(feature = "candle")]
    pub fn candle_device(&self) -> candle_core::Device {
        match &self.config.device_preference {
            DevicePreference::Cpu => {
                return candle_core::Device::Cpu;
            }
            DevicePreference::Cuda(ordinal) => {
                match candle_core::Device::new_cuda(*ordinal) {
                    Ok(device) => return device,
                    Err(e) => {
                        self.emit_fallback(GpuFallbackEvent {
                            reason: FallbackReason::DeviceNotFound,
                            detail: format!("CUDA:{ordinal} failed: {e}"),
                            operation: "candle_device".into(),
                            requested_device: self.config.device_preference.clone(),
                        });
                        return candle_core::Device::Cpu;
                    }
                }
            }
            DevicePreference::Metal => {
                match candle_core::Device::new_metal(0) {
                    Ok(device) => return device,
                    Err(e) => {
                        self.emit_fallback(GpuFallbackEvent {
                            reason: FallbackReason::DeviceNotFound,
                            detail: format!("Metal failed: {e}"),
                            operation: "candle_device".into(),
                            requested_device: self.config.device_preference.clone(),
                        });
                        return candle_core::Device::Cpu;
                    }
                }
            }
            DevicePreference::Auto => {
                // Auto 模式: CUDA → Metal → CPU
                // 先尝试 CUDA
                if let Ok(device) = candle_core::Device::new_cuda(0) {
                    tracing::info!("DeviceSelector: auto-selected CUDA:0");
                    return device;
                }

                // 再尝试 Metal
                #[cfg(target_os = "macos")]
                if let Ok(device) = candle_core::Device::new_metal(0) {
                    tracing::info!("DeviceSelector: auto-selected Metal");
                    return device;
                }

                // fallback CPU
                self.emit_fallback(GpuFallbackEvent {
                    reason: FallbackReason::NoCompatibleDevice,
                    detail: "No CUDA or Metal device available".into(),
                    operation: "candle_device".into(),
                    requested_device: DevicePreference::Auto,
                });
                candle_core::Device::Cpu
            }
        }
    }

    /// 判断给定工作负载是否应使用 GPU
    ///
    /// 综合考虑：
    /// 1. 用户配置的设备偏好
    /// 2. GPU 硬件是否可用
    /// 3. 数据量是否值得 GPU 传输开销
    /// 4. GPU 内存是否足够
    ///
    /// 返回 `false` 时，调用方应使用 CPU 路径。
    pub fn should_use_gpu(&self, hint: &WorkloadHint) -> bool {
        // 1. 用户显式选择 CPU → 直接返回 false
        if self.config.device_preference == DevicePreference::Cpu {
            return false;
        }

        // 2. 检查数据量阈值（Auto 模式下）
        if self.config.device_preference == DevicePreference::Auto {
            if hint.num_items < self.config.auto_min_vectors {
                if self.config.log_fallback {
                    self.emit_fallback(GpuFallbackEvent {
                        reason: FallbackReason::DataTooSmall,
                        detail: format!(
                            "{} items < min {} for GPU",
                            hint.num_items, self.config.auto_min_vectors
                        ),
                        operation: hint.operation.into(),
                        requested_device: self.config.device_preference.clone(),
                    });
                }
                return false;
            }
            if hint.estimated_data_bytes() < self.config.auto_min_data_bytes {
                if self.config.log_fallback {
                    self.emit_fallback(GpuFallbackEvent {
                        reason: FallbackReason::DataTooSmall,
                        detail: format!(
                            "{}KB data < min {}KB for GPU",
                            hint.estimated_data_bytes() / 1024,
                            self.config.auto_min_data_bytes / 1024
                        ),
                        operation: hint.operation.into(),
                        requested_device: self.config.device_preference.clone(),
                    });
                }
                return false;
            }
        }

        // 3. 检查 GPU 硬件可用性
        let availability = self.availability();
        match availability {
            DeviceAvailability::CudaAvailable {
                free_memory_bytes, ..
            } => {
                // 4. 检查内存是否足够（需要约 3x 数据量用于工作空间）
                let required = hint.estimated_data_bytes() * 3;
                let budget = self.config.memory_budget.unwrap_or(*free_memory_bytes);
                if required > budget {
                    if self.config.log_fallback {
                        self.emit_fallback(GpuFallbackEvent {
                            reason: FallbackReason::InsufficientMemory,
                            detail: format!(
                                "Need {}MB, budget {}MB for '{}'",
                                required / (1024 * 1024),
                                budget / (1024 * 1024),
                                hint.operation
                            ),
                            operation: hint.operation.into(),
                            requested_device: self.config.device_preference.clone(),
                        });
                    }
                    return false;
                }
                true
            }
            DeviceAvailability::MetalAvailable => {
                // Metal 可用（macOS），内存预算由系统管理
                true
            }
            DeviceAvailability::CpuOnly(reason) => {
                if self.config.log_fallback
                    && self.config.device_preference != DevicePreference::Cpu
                {
                    self.emit_fallback(GpuFallbackEvent {
                        reason: FallbackReason::NoCompatibleDevice,
                        detail: reason.clone(),
                        operation: hint.operation.into(),
                        requested_device: self.config.device_preference.clone(),
                    });
                }
                false
            }
        }
    }

    /// 尝试获取 GPU candle Device，失败时返回 CPU 并记录降级
    ///
    /// 这是推荐的获取设备方式：调用者无需处理错误。
    /// 始终返回一个可用的 Device（最差情况是 CPU）。
    #[cfg(feature = "candle")]
    pub fn candle_device_with_fallback(&self, operation: &str) -> (candle_core::Device, bool) {
        let device = self.candle_device();
        let is_gpu = !matches!(device, candle_core::Device::Cpu);
        if !is_gpu && self.config.device_preference != DevicePreference::Cpu {
            tracing::debug!(
                operation = operation,
                "DeviceSelector: using CPU for '{operation}'"
            );
        } else if is_gpu {
            tracing::debug!(
                operation = operation,
                device = ?device,
                "DeviceSelector: using GPU for '{operation}'"
            );
        }
        (device, is_gpu)
    }

    /// 获取全局降级事件计数
    pub fn fallback_count() -> u64 {
        FALLBACK_COUNT.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// 重置全局降级计数（用于测试）
    #[cfg(test)]
    pub fn reset_fallback_count() {
        FALLBACK_COUNT.store(0, std::sync::atomic::Ordering::Relaxed);
    }

    /// 发出降级事件（公共接口，供 gpu 模块使用）
    pub fn emit_fallback_public(&self, event: GpuFallbackEvent) {
        self.emit_fallback(event);
    }

    // ---- internal ----

    /// 发出降级事件
    fn emit_fallback(&self, event: GpuFallbackEvent) {
        FALLBACK_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if self.config.log_fallback {
            event.log();
        }
    }
}

impl Default for DeviceSelector {
    fn default() -> Self {
        Self::auto()
    }
}

// ============================================================================
// 全局便捷函数
// ============================================================================

/// 全局默认 DeviceSelector（auto 模式）
///
/// 用于不需要自定义配置的场景。
static GLOBAL_SELECTOR: OnceLock<DeviceSelector> = OnceLock::new();

/// 获取全局默认 DeviceSelector
pub fn global_selector() -> &'static DeviceSelector {
    GLOBAL_SELECTOR.get_or_init(|| DeviceSelector::auto())
}

/// 快捷方式：获取 candle Device（使用全局默认配置）
#[cfg(feature = "candle")]
pub fn auto_device() -> candle_core::Device {
    global_selector().candle_device()
}

/// 快捷方式：判断是否应使用 GPU
pub fn should_use_gpu(num_items: usize, dimension: usize) -> bool {
    global_selector().should_use_gpu(&WorkloadHint::new(num_items, dimension))
}

// ============================================================================
// probe_device_availability — 核心探测函数
// ============================================================================

/// 探测设备可用性
///
/// 使用 candle 作为轻量级 CUDA 探测手段。
/// 此函数永远不会 panic — 任何错误都返回 `CpuOnly`。
fn probe_device_availability() -> DeviceAvailability {
    #[cfg(feature = "candle")]
    {
        // 尝试 CUDA
        match candle_core::Device::new_cuda(0) {
            Ok(device) => {
                // 探测可用内存
                let (total_mem, free_mem) = estimate_gpu_memory(&device);
                let name = format!("CUDA Device 0"); // candle 不直接暴露设备名称

                tracing::info!(
                    "DeviceSelector probe: CUDA available ({}, {}MB free / {}MB total)",
                    name,
                    free_mem / (1024 * 1024),
                    total_mem / (1024 * 1024)
                );

                return DeviceAvailability::CudaAvailable {
                    ordinal: 0,
                    name,
                    free_memory_bytes: free_mem,
                    total_memory_bytes: total_mem,
                };
            }
            Err(e) => {
                tracing::debug!("DeviceSelector probe: CUDA not available: {e}");
            }
        }

        // 尝试 Metal (macOS)
        #[cfg(target_os = "macos")]
        {
            match candle_core::Device::new_metal(0) {
                Ok(_) => {
                    tracing::info!("DeviceSelector probe: Metal GPU available");
                    return DeviceAvailability::MetalAvailable;
                }
                Err(e) => {
                    tracing::debug!("DeviceSelector probe: Metal not available: {e}");
                }
            }
        }

        DeviceAvailability::CpuOnly("No CUDA or Metal device detected".into())
    }

    #[cfg(not(feature = "candle"))]
    {
        DeviceAvailability::CpuOnly(
            "GPU support requires the 'candle' feature to be enabled at compile time".into(),
        )
    }
}

/// 通过 candle 估算 GPU 可用内存
#[cfg(feature = "candle")]
fn estimate_gpu_memory(device: &candle_core::Device) -> (usize, usize) {
    let mut max_successful_bytes = 0usize;

    // 尝试递增分配来估算可用内存
    for size_mb in [256, 512, 1024, 2048, 4096, 8192, 16384, 24576] {
        let num_floats = size_mb * 1024 * 1024 / 4;
        match candle_core::Tensor::zeros(num_floats, candle_core::DType::F32, device) {
            Ok(_tensor) => {
                max_successful_bytes = size_mb * 1024 * 1024;
            }
            Err(_) => break,
        }
    }

    if max_successful_bytes == 0 {
        // 无法确定内存，使用保守默认值
        (4 * 1024 * 1024 * 1024, 2 * 1024 * 1024 * 1024)
    } else {
        let estimated_total = (max_successful_bytes as f64 / 0.8) as usize;
        (estimated_total, max_successful_bytes)
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // ----- DevicePreference tests -----

    #[test]
    fn test_device_preference_default_is_auto() {
        assert_eq!(DevicePreference::default(), DevicePreference::Auto);
    }

    #[test]
    fn test_device_preference_from_str() {
        assert_eq!(
            DevicePreference::from_str_flexible("auto"),
            DevicePreference::Auto
        );
        assert_eq!(
            DevicePreference::from_str_flexible("cpu"),
            DevicePreference::Cpu
        );
        assert_eq!(
            DevicePreference::from_str_flexible("cuda"),
            DevicePreference::Cuda(0)
        );
        assert_eq!(
            DevicePreference::from_str_flexible("cuda:0"),
            DevicePreference::Cuda(0)
        );
        assert_eq!(
            DevicePreference::from_str_flexible("cuda:2"),
            DevicePreference::Cuda(2)
        );
        assert_eq!(
            DevicePreference::from_str_flexible("metal"),
            DevicePreference::Metal
        );
        assert_eq!(
            DevicePreference::from_str_flexible("gpu"),
            DevicePreference::Cuda(0)
        );
        // Unknown → Auto
        assert_eq!(
            DevicePreference::from_str_flexible("unknown"),
            DevicePreference::Auto
        );
        // Case insensitive
        assert_eq!(
            DevicePreference::from_str_flexible("CUDA:1"),
            DevicePreference::Cuda(1)
        );
        assert_eq!(
            DevicePreference::from_str_flexible("AUTO"),
            DevicePreference::Auto
        );
    }

    #[test]
    fn test_device_preference_display() {
        assert_eq!(format!("{}", DevicePreference::Auto), "auto");
        assert_eq!(format!("{}", DevicePreference::Cpu), "cpu");
        assert_eq!(format!("{}", DevicePreference::Cuda(0)), "cuda:0");
        assert_eq!(format!("{}", DevicePreference::Cuda(2)), "cuda:2");
        assert_eq!(format!("{}", DevicePreference::Metal), "metal");
    }

    // ----- GpuAccelConfig tests -----

    #[test]
    fn test_gpu_accel_config_default() {
        let config = GpuAccelConfig::default();
        assert_eq!(config.device_preference, DevicePreference::Auto);
        assert!(config.memory_budget.is_none());
        assert!(config.log_fallback);
        assert_eq!(config.auto_min_data_bytes, 1024 * 1024);
        assert_eq!(config.auto_min_vectors, 1_000);
    }

    #[test]
    fn test_gpu_accel_config_cpu_only() {
        let config = GpuAccelConfig::cpu_only();
        assert_eq!(config.device_preference, DevicePreference::Cpu);
    }

    #[test]
    fn test_gpu_accel_config_with_budget() {
        let config = GpuAccelConfig::default().with_memory_budget_mb(4096);
        assert_eq!(config.memory_budget, Some(4096 * 1024 * 1024));
    }

    // ----- WorkloadHint tests -----

    #[test]
    fn test_workload_hint_data_bytes() {
        let hint = WorkloadHint::new(100_000, 768);
        // 100K * 768 * 4 = 307,200,000 bytes ≈ 293MB
        assert_eq!(hint.estimated_data_bytes(), 100_000 * 768 * 4);
    }

    #[test]
    fn test_workload_hint_with_operation() {
        let hint = WorkloadHint::with_operation(50_000, 384, "kmeans");
        assert_eq!(hint.operation, "kmeans");
        assert_eq!(hint.num_items, 50_000);
        assert_eq!(hint.dimension, 384);
    }

    // ----- DeviceAvailability tests -----

    #[test]
    fn test_device_availability_cpu_only() {
        let avail = DeviceAvailability::CpuOnly("No GPU".into());
        assert!(!avail.has_gpu());
        assert!(format!("{avail}").contains("CPU only"));
    }

    #[test]
    fn test_device_availability_cuda() {
        let avail = DeviceAvailability::CudaAvailable {
            ordinal: 0,
            name: "RTX 4090".into(),
            free_memory_bytes: 20 * 1024 * 1024 * 1024,
            total_memory_bytes: 24 * 1024 * 1024 * 1024,
        };
        assert!(avail.has_gpu());
        let display = format!("{avail}");
        assert!(display.contains("CUDA:0"));
        assert!(display.contains("RTX 4090"));
    }

    // ----- FallbackReason tests -----

    #[test]
    fn test_fallback_reason_display() {
        assert!(!format!("{}", FallbackReason::FeatureNotEnabled).is_empty());
        assert!(!format!("{}", FallbackReason::DriverNotFound).is_empty());
        assert!(!format!("{}", FallbackReason::NoCompatibleDevice).is_empty());
        assert!(!format!("{}", FallbackReason::InsufficientMemory).is_empty());
        assert!(!format!("{}", FallbackReason::DeviceNotFound).is_empty());
        assert!(!format!("{}", FallbackReason::RuntimeError).is_empty());
        assert!(!format!("{}", FallbackReason::DataTooSmall).is_empty());
        assert!(!format!("{}", FallbackReason::UserRequestedCpu).is_empty());
    }

    // ----- DeviceSelector tests -----

    #[test]
    fn test_selector_cpu_only_config() {
        let selector = DeviceSelector::new(GpuAccelConfig::cpu_only());
        // CPU 模式下 should_use_gpu 始终返回 false
        let hint = WorkloadHint::with_operation(1_000_000, 768, "test");
        assert!(!selector.should_use_gpu(&hint));
    }

    #[test]
    fn test_selector_small_data_no_gpu() {
        let selector = DeviceSelector::auto();
        // 小数据量 → should_use_gpu 返回 false
        let hint = WorkloadHint::with_operation(10, 768, "test_small");
        assert!(!selector.should_use_gpu(&hint));
    }

    #[test]
    fn test_selector_default_is_auto() {
        let selector = DeviceSelector::default();
        assert_eq!(selector.config().device_preference, DevicePreference::Auto);
    }

    #[test]
    fn test_fallback_count_increments() {
        DeviceSelector::reset_fallback_count();
        let initial = DeviceSelector::fallback_count();

        let selector = DeviceSelector::auto();
        // 小数据量会触发 DataTooSmall fallback
        let hint = WorkloadHint::with_operation(5, 768, "test_count");
        let _ = selector.should_use_gpu(&hint);

        // 计数至少不减
        assert!(DeviceSelector::fallback_count() >= initial);
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_selector_candle_device_cpu_mode() {
        let selector = DeviceSelector::new(GpuAccelConfig::cpu_only());
        let device = selector.candle_device();
        assert!(matches!(device, candle_core::Device::Cpu));
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_selector_candle_device_auto_never_panics() {
        // Auto 模式永远不 panic — 最差返回 CPU
        let selector = DeviceSelector::auto();
        let device = selector.candle_device();
        // 无论是否有 GPU，都应该返回一个有效设备
        let _ = device; // 不 panic 就是成功
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_selector_candle_device_with_fallback() {
        let selector = DeviceSelector::auto();
        let (device, _is_gpu) = selector.candle_device_with_fallback("test");
        // 无论是否有 GPU，都应返回可用设备
        let _ = device;
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_selector_invalid_cuda_device_fallback() {
        // 请求不存在的 CUDA:99 → 应 fallback 到 CPU
        let selector = DeviceSelector::new(GpuAccelConfig {
            device_preference: DevicePreference::Cuda(99),
            ..Default::default()
        });
        let device = selector.candle_device();
        // 应该 fallback 到 CPU（CUDA:99 不太可能存在）
        // 注意：如果确实有 100 块 GPU... 那也 OK
        let _ = device;
    }

    #[test]
    fn test_probe_device_availability_never_panics() {
        let avail = probe_device_availability();
        // 无论硬件环境如何，都应返回有效结果
        match &avail {
            DeviceAvailability::CudaAvailable { name, .. } => {
                assert!(!name.is_empty());
            }
            DeviceAvailability::MetalAvailable => {}
            DeviceAvailability::CpuOnly(reason) => {
                assert!(!reason.is_empty());
            }
        }
    }

    #[test]
    fn test_global_selector_is_auto() {
        let selector = global_selector();
        assert_eq!(selector.config().device_preference, DevicePreference::Auto);
    }

    // ----- 安全降级集成测试 -----

    #[test]
    fn test_safe_degradation_cpu_only_full_pipeline() {
        // 端到端：CPU-only config → should_use_gpu=false → 一切正常
        let config = GpuAccelConfig::cpu_only();
        let selector = DeviceSelector::new(config);

        // 无论数据量多大，CPU 模式下都不使用 GPU
        for &items in &[10, 1_000, 100_000, 10_000_000] {
            let hint = WorkloadHint::with_operation(items, 768, "safe_degrade_test");
            assert!(
                !selector.should_use_gpu(&hint),
                "CPU-only config should never use GPU, even for {} items",
                items
            );
        }
    }

    #[test]
    fn test_safe_degradation_auto_small_data_threshold() {
        // Auto 模式下小数据量自动降级到 CPU
        let config = GpuAccelConfig {
            device_preference: DevicePreference::Auto,
            auto_min_vectors: 5_000,
            auto_min_data_bytes: 2 * 1024 * 1024, // 2MB
            ..Default::default()
        };
        let selector = DeviceSelector::new(config);

        // 100 个 768 维向量 = 100*768*4 = 307KB < 2MB → 不用 GPU
        let hint_small = WorkloadHint::with_operation(100, 768, "small_batch");
        assert!(!selector.should_use_gpu(&hint_small));

        // 1000 个向量 < 5000 阈值 → 不用 GPU
        let hint_medium = WorkloadHint::with_operation(1_000, 768, "medium_batch");
        assert!(!selector.should_use_gpu(&hint_medium));
    }

    #[test]
    fn test_safe_degradation_custom_thresholds() {
        // 自定义阈值测试
        let config = GpuAccelConfig {
            device_preference: DevicePreference::Auto,
            auto_min_vectors: 100,
            auto_min_data_bytes: 1024, // 1KB — 极低阈值
            ..Default::default()
        };
        let selector = DeviceSelector::new(config);

        // 50 < 100 → 不用 GPU
        let hint = WorkloadHint::with_operation(50, 128, "low_threshold");
        assert!(!selector.should_use_gpu(&hint));
    }

    #[test]
    fn test_workload_hint_operations() {
        // 不同操作类型的 WorkloadHint
        let ops = ["kmeans", "brute_force_search", "embedding_batch", "pq_train"];
        for op in ops {
            let hint = WorkloadHint::with_operation(10_000, 384, op);
            assert_eq!(hint.operation, op);
            assert_eq!(hint.estimated_data_bytes(), 10_000 * 384 * 4);
        }
    }

    #[test]
    fn test_gpu_fallback_event_logging() {
        // 各类降级事件不 panic
        let events = vec![
            GpuFallbackEvent {
                reason: FallbackReason::FeatureNotEnabled,
                detail: "test".into(),
                operation: "test_op".into(),
                requested_device: DevicePreference::Auto,
            },
            GpuFallbackEvent {
                reason: FallbackReason::UserRequestedCpu,
                detail: "user choice".into(),
                operation: "test_op2".into(),
                requested_device: DevicePreference::Cpu,
            },
            GpuFallbackEvent {
                reason: FallbackReason::DataTooSmall,
                detail: "only 10 items".into(),
                operation: "kmeans".into(),
                requested_device: DevicePreference::Auto,
            },
            GpuFallbackEvent {
                reason: FallbackReason::InsufficientMemory,
                detail: "need 8GB, have 4GB".into(),
                operation: "large_index".into(),
                requested_device: DevicePreference::Cuda(0),
            },
            GpuFallbackEvent {
                reason: FallbackReason::RuntimeError,
                detail: "CUDA OOM".into(),
                operation: "matmul".into(),
                requested_device: DevicePreference::Cuda(0),
            },
        ];
        for event in &events {
            event.log(); // should never panic
        }
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_safe_degradation_compute_backend_integration() {
        use crate::query::compute_backend;

        // CPU-only selector → select_backend 应返回非 GPU 后端
        let selector = DeviceSelector::new(GpuAccelConfig::cpu_only());
        let backend = compute_backend::select_backend_with_selector(64, &selector);
        assert!(!backend.is_gpu(), "CPU-only config should not produce GPU backend");
        // 应该是 candle-cpu（dim>32）或 cpu-rust
        assert!(
            backend.name() == "candle-cpu" || backend.name() == "cpu-rust",
            "Expected candle-cpu or cpu-rust, got {}",
            backend.name()
        );
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_safe_degradation_compute_backend_small_dim() {
        use crate::query::compute_backend;

        // 小维度 → 应使用纯 Rust CPU 后端
        let selector = DeviceSelector::new(GpuAccelConfig::cpu_only());
        let backend = compute_backend::select_backend_with_selector(8, &selector);
        assert_eq!(backend.name(), "cpu-rust");
        assert!(!backend.is_gpu());
    }

    #[cfg(feature = "candle")]
    #[test]
    fn test_select_backend_default_never_panics() {
        use crate::query::compute_backend;

        // 默认 select_backend 永远不 panic
        for &dim in &[1, 3, 16, 32, 64, 128, 256, 512, 1024] {
            let backend = compute_backend::select_backend(dim);
            assert!(!backend.name().is_empty());
        }
    }

    #[test]
    fn test_device_preference_roundtrip() {
        // from_str_flexible → Display 往返一致性
        let prefs = vec![
            ("auto", DevicePreference::Auto),
            ("cpu", DevicePreference::Cpu),
            ("cuda:0", DevicePreference::Cuda(0)),
            ("cuda:2", DevicePreference::Cuda(2)),
            ("metal", DevicePreference::Metal),
        ];
        for (s, expected) in &prefs {
            let parsed = DevicePreference::from_str_flexible(s);
            assert_eq!(&parsed, expected, "Parsing '{s}' failed");
            let display = format!("{parsed}");
            let re_parsed = DevicePreference::from_str_flexible(&display);
            assert_eq!(&re_parsed, expected, "Roundtrip failed for '{s}'");
        }
    }

    #[test]
    fn test_availability_display_format() {
        // 验证 Display 格式不为空
        let avail_cpu = DeviceAvailability::CpuOnly("No CUDA".into());
        let avail_metal = DeviceAvailability::MetalAvailable;
        let avail_cuda = DeviceAvailability::CudaAvailable {
            ordinal: 0,
            name: "RTX 3090".into(),
            free_memory_bytes: 10 * 1024 * 1024 * 1024,
            total_memory_bytes: 24 * 1024 * 1024 * 1024,
        };

        for avail in &[avail_cpu, avail_metal, avail_cuda] {
            let display = format!("{avail}");
            assert!(!display.is_empty(), "Display should not be empty");
        }
    }

    #[test]
    fn test_memory_budget_respected() {
        // 设置很小的内存预算 → 即使 GPU 可用，大数据量也应 fallback
        let config = GpuAccelConfig {
            device_preference: DevicePreference::Auto,
            memory_budget: Some(1024), // 仅 1KB
            auto_min_vectors: 1, // 不限制向量数
            auto_min_data_bytes: 1, // 不限制数据量
            ..Default::default()
        };
        let selector = DeviceSelector::new(config);

        // 如果环境恰好有 GPU，巨量数据应受内存预算限制
        let hint = WorkloadHint::with_operation(1_000_000, 768, "memory_test");
        // 1M * 768 * 4 = 3GB >> 1KB budget
        // 无论 GPU 是否存在，内存预算检查都应在 GPU 可用时生效
        let result = selector.should_use_gpu(&hint);
        // 在 CPU-only 环境：false（无 GPU）
        // 在 GPU 环境：false（内存不足）
        // 两种情况都是 false
        assert!(!result, "Memory budget 1KB should reject 3GB workload");
    }
}
