//! GPU index building configuration
//!
//! Extends the existing `VectorIndexConfig` with GPU-specific settings.
//! All GPU config fields are optional — the system auto-selects sensible defaults.

use std::fmt;

/// GPU index building configuration.
///
/// This is the primary configuration struct for GPU-accelerated index construction.
/// All fields have sensible defaults; the simplest usage is `GpuIndexConfig::auto()`.
///
/// # Design Principle: Zero-Impact when GPU Absent
///
/// When GPU is not available, `GpuIndexConfig` gracefully degrades:
/// - `GpuBuildMode::Auto` → detects no GPU → uses CPU path
/// - `GpuBuildMode::FullGpu` → detects no GPU → returns error, caller falls back
/// - `GpuBuildMode::Hybrid` → detects no GPU → runs CPU-only with same parameters
#[derive(Debug, Clone)]
pub struct GpuIndexConfig {
    /// GPU device ID (-1 = auto-select best device, 0/1/2... = specific device)
    pub device_id: i32,

    /// GPU index build mode
    pub mode: GpuBuildMode,

    /// GPU memory budget in bytes.
    /// If `None`, uses 80% of free GPU memory.
    /// If dataset exceeds this budget, automatically switches to streaming/batched build.
    pub gpu_memory_budget: Option<usize>,

    /// Enable CAGRA graph index algorithm (GPU-exclusive, NVIDIA 2023).
    /// CAGRA is 10-100x faster than HNSW on GPU for both build and search.
    /// Only used when GPU is actually available.
    pub use_cagra: bool,

    /// CAGRA graph degree (default: 64).
    /// Higher values → better recall but larger index.
    pub cagra_graph_degree: Option<usize>,

    /// CAGRA intermediate graph degree used during build (default: 128).
    /// Higher values → better graph quality but slower build.
    pub cagra_intermediate_graph_degree: Option<usize>,

    /// Maximum number of KMeans iterations for IVF training on GPU (default: 50).
    pub max_kmeans_iterations: Option<usize>,

    /// KMeans convergence threshold (default: 1e-4).
    /// Training stops when inertia change ratio falls below this.
    pub kmeans_tolerance: Option<f64>,

    /// Whether to log GPU operations for debugging.
    pub verbose: bool,
}

/// GPU build mode — controls how CPU and GPU cooperate.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GpuBuildMode {
    /// Fully GPU build: transfer all data to GPU, run everything on device.
    ///
    /// Best for datasets that fit in GPU memory.
    /// If GPU is not available, this mode will fail with an error
    /// (caller should catch and fall back to CPU).
    FullGpu,

    /// Hybrid CPU/GPU: CPU handles I/O + preprocessing, GPU handles compute.
    ///
    /// Best for large datasets that exceed GPU memory.
    /// If GPU is not available, this degrades to CPU-only with the same thread config.
    Hybrid {
        /// Number of CPU threads for data preprocessing (default: num_cpus)
        cpu_threads: usize,
        /// Number of vectors per GPU batch (default: auto-computed from GPU memory)
        gpu_batch_size: usize,
    },

    /// Auto-select based on data size and GPU memory.
    ///
    /// Decision logic:
    /// - No GPU → CPU path (no error, transparent fallback)
    /// - Data fits in GPU memory → FullGpu
    /// - Data exceeds GPU memory → Hybrid with auto-tuned batch size
    ///
    /// **This is the recommended default.**
    Auto,
}

impl Default for GpuIndexConfig {
    fn default() -> Self {
        Self {
            device_id: -1, // auto-select
            mode: GpuBuildMode::Auto,
            gpu_memory_budget: None,
            use_cagra: false,
            cagra_graph_degree: None,
            cagra_intermediate_graph_degree: None,
            max_kmeans_iterations: None,
            kmeans_tolerance: None,
            verbose: false,
        }
    }
}

impl GpuIndexConfig {
    /// Create a config that automatically selects the best GPU strategy.
    ///
    /// This is the simplest and recommended way to use GPU acceleration.
    /// If no GPU is available, the system transparently falls back to CPU.
    pub fn auto() -> Self {
        Self::default()
    }

    /// Create a config for full GPU build on a specific device.
    pub fn full_gpu(device_id: i32) -> Self {
        Self {
            device_id,
            mode: GpuBuildMode::FullGpu,
            ..Default::default()
        }
    }

    /// Create a config for hybrid CPU/GPU build.
    pub fn hybrid(cpu_threads: usize, gpu_batch_size: usize) -> Self {
        Self {
            mode: GpuBuildMode::Hybrid {
                cpu_threads,
                gpu_batch_size,
            },
            ..Default::default()
        }
    }

    /// Enable CAGRA graph index (GPU-exclusive algorithm).
    pub fn with_cagra(mut self) -> Self {
        self.use_cagra = true;
        self
    }

    /// Set CAGRA graph parameters.
    pub fn with_cagra_params(
        mut self,
        graph_degree: usize,
        intermediate_graph_degree: usize,
    ) -> Self {
        self.use_cagra = true;
        self.cagra_graph_degree = Some(graph_degree);
        self.cagra_intermediate_graph_degree = Some(intermediate_graph_degree);
        self
    }

    /// Set GPU memory budget in bytes.
    pub fn with_memory_budget(mut self, budget_bytes: usize) -> Self {
        self.gpu_memory_budget = Some(budget_bytes);
        self
    }

    /// Set GPU memory budget in megabytes (convenience).
    pub fn with_memory_budget_mb(mut self, budget_mb: usize) -> Self {
        self.gpu_memory_budget = Some(budget_mb * 1024 * 1024);
        self
    }

    /// Enable verbose logging.
    pub fn with_verbose(mut self) -> Self {
        self.verbose = true;
        self
    }

    /// Get the effective CAGRA graph degree.
    pub fn effective_cagra_graph_degree(&self) -> usize {
        self.cagra_graph_degree.unwrap_or(64)
    }

    /// Get the effective CAGRA intermediate graph degree.
    pub fn effective_cagra_intermediate_graph_degree(&self) -> usize {
        self.cagra_intermediate_graph_degree.unwrap_or(128)
    }

    /// Get the effective max KMeans iterations.
    pub fn effective_max_kmeans_iterations(&self) -> usize {
        self.max_kmeans_iterations.unwrap_or(50)
    }

    /// Get the effective KMeans tolerance.
    pub fn effective_kmeans_tolerance(&self) -> f64 {
        self.kmeans_tolerance.unwrap_or(1e-4)
    }
}

impl fmt::Display for GpuBuildMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            GpuBuildMode::FullGpu => write!(f, "FullGpu"),
            GpuBuildMode::Hybrid {
                cpu_threads,
                gpu_batch_size,
            } => write!(
                f,
                "Hybrid(cpu_threads={cpu_threads}, gpu_batch={gpu_batch_size})"
            ),
            GpuBuildMode::Auto => write!(f, "Auto"),
        }
    }
}

impl fmt::Display for GpuIndexConfig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "GpuIndexConfig {{ device={}, mode={}, cagra={}, budget={} }}",
            if self.device_id < 0 {
                "auto".to_string()
            } else {
                format!("cuda:{}", self.device_id)
            },
            self.mode,
            self.use_cagra,
            self.gpu_memory_budget
                .map(|b| format!("{}MB", b / (1024 * 1024)))
                .unwrap_or_else(|| "auto".to_string()),
        )
    }
}

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

    #[test]
    fn test_default_config() {
        let config = GpuIndexConfig::default();
        assert_eq!(config.device_id, -1);
        assert_eq!(config.mode, GpuBuildMode::Auto);
        assert!(!config.use_cagra);
        assert!(config.gpu_memory_budget.is_none());
    }

    #[test]
    fn test_auto_config() {
        let config = GpuIndexConfig::auto();
        assert_eq!(config.mode, GpuBuildMode::Auto);
    }

    #[test]
    fn test_full_gpu_config() {
        let config = GpuIndexConfig::full_gpu(0);
        assert_eq!(config.device_id, 0);
        assert_eq!(config.mode, GpuBuildMode::FullGpu);
    }

    #[test]
    fn test_hybrid_config() {
        let config = GpuIndexConfig::hybrid(8, 100_000);
        assert_eq!(
            config.mode,
            GpuBuildMode::Hybrid {
                cpu_threads: 8,
                gpu_batch_size: 100_000
            }
        );
    }

    #[test]
    fn test_cagra_config() {
        let config = GpuIndexConfig::auto().with_cagra_params(32, 64);
        assert!(config.use_cagra);
        assert_eq!(config.effective_cagra_graph_degree(), 32);
        assert_eq!(config.effective_cagra_intermediate_graph_degree(), 64);
    }

    #[test]
    fn test_memory_budget_mb() {
        let config = GpuIndexConfig::auto().with_memory_budget_mb(4096);
        assert_eq!(config.gpu_memory_budget, Some(4096 * 1024 * 1024));
    }

    #[test]
    fn test_display() {
        let config = GpuIndexConfig::auto();
        let s = format!("{config}");
        assert!(s.contains("Auto"));
        assert!(s.contains("auto"));
    }
}
