//! End-to-end hybrid CPU/GPU index building pipeline
//!
//! Orchestrates the full index build flow:
//!
//! ```text
//! [CPU] Data loading → [CPU] Preprocessing → [GPU] KMeans training
//! → [GPU] Quantization → [GPU/CPU] Graph build → [CPU] Serialize to Lance
//! ```
//!
//! ## Design Principles
//!
//! 1. **Best-effort GPU**: Uses GPU when available, falls back transparently.
//! 2. **Streaming**: Processes data in batches to handle datasets >> GPU memory.
//! 3. **Observable**: Reports progress through `PipelineProgress` callbacks.
//! 4. **Composable**: Each stage can be used independently.

use std::sync::Arc;
use std::time::Instant;

use super::cagra::{CagraBuildStats, CagraConfig, CagraIndex};
use super::config::{GpuBuildMode, GpuIndexConfig};
use super::device::{GpuAvailability, GpuDeviceInfo};
use super::error::{GpuError, GpuResult};
use super::kmeans::{GpuKMeansConfig, GpuKMeansTrainer, KMeansResult};
use super::memory::GpuMemoryPool;

/// Overall pipeline stage for progress reporting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PipelineStage {
    /// Probing GPU device
    DeviceProbe,
    /// Loading and preprocessing data
    DataPrep,
    /// Training KMeans centroids (IVF)
    KMeansTraining,
    /// Training quantization codebooks (PQ/SQ)
    QuantizationTraining,
    /// Building graph index (CAGRA/HNSW)
    GraphBuild,
    /// Assigning vectors to partitions
    Assignment,
    /// Serializing index to storage
    Serialization,
    /// Pipeline completed
    Completed,
}

impl std::fmt::Display for PipelineStage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::DeviceProbe => write!(f, "Device Probe"),
            Self::DataPrep => write!(f, "Data Preprocessing"),
            Self::KMeansTraining => write!(f, "KMeans Training"),
            Self::QuantizationTraining => write!(f, "Quantization Training"),
            Self::GraphBuild => write!(f, "Graph Build"),
            Self::Assignment => write!(f, "Partition Assignment"),
            Self::Serialization => write!(f, "Serialization"),
            Self::Completed => write!(f, "Completed"),
        }
    }
}

/// Pipeline execution result.
#[derive(Debug, Clone)]
pub struct PipelineResult {
    /// Which device was ultimately used
    pub device_used: DeviceUsed,
    /// Total build time in milliseconds
    pub total_time_ms: u64,
    /// Time spent per stage (stage name → ms)
    pub stage_times: Vec<(String, u64)>,
    /// GPU device info (if used)
    pub gpu_device: Option<GpuDeviceInfo>,
    /// Number of vectors processed
    pub num_vectors: usize,
    /// Vector dimension
    pub dimension: usize,
    /// Index type built
    pub index_type: String,
    /// Whether any stage fell back to CPU
    pub had_cpu_fallback: bool,
    /// KMeans: number of centroids trained
    pub num_centroids: Option<usize>,
    /// CAGRA build stats (if applicable)
    pub cagra_stats: Option<CagraBuildStats>,
}

/// Which device was used for the compute-heavy stages.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeviceUsed {
    /// Pure CPU
    Cpu,
    /// GPU with CPU fallback for some stages
    GpuHybrid,
    /// Full GPU
    Gpu,
}

impl std::fmt::Display for DeviceUsed {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Cpu => write!(f, "CPU"),
            Self::GpuHybrid => write!(f, "GPU+CPU Hybrid"),
            Self::Gpu => write!(f, "GPU"),
        }
    }
}

/// Progress callback type.
pub type ProgressCallback = Box<dyn Fn(PipelineStage, f64) + Send + Sync>;

/// The main hybrid CPU/GPU index building pipeline.
pub struct IndexBuildPipeline {
    /// GPU index configuration
    gpu_config: GpuIndexConfig,

    /// GPU memory pool (if GPU is available)
    memory_pool: Option<GpuMemoryPool>,

    /// GPU availability information
    gpu_availability: GpuAvailability,

    /// Whether GPU is actually usable for this build
    use_gpu: bool,

    /// Progress callback
    progress_cb: Option<ProgressCallback>,
}

impl IndexBuildPipeline {
    /// Create a new index build pipeline.
    ///
    /// Probes GPU device and determines the build strategy.
    pub fn new(gpu_config: GpuIndexConfig) -> Self {
        let gpu_availability = super::gpu_available();

        // Determine if GPU should be used based on availability and config mode
        let use_gpu = match gpu_config.mode {
            GpuBuildMode::FullGpu => gpu_availability.is_available(),
            GpuBuildMode::Auto => gpu_availability.is_available(),
            GpuBuildMode::Hybrid { .. } => gpu_availability.is_available(),
        };

        let memory_pool = if use_gpu {
            gpu_availability.best_device().map(|info| {
                GpuMemoryPool::new(info.clone(), gpu_config.gpu_memory_budget)
            })
        } else {
            None
        };

        Self {
            gpu_config,
            memory_pool,
            gpu_availability,
            use_gpu,
            progress_cb: None,
        }
    }

    /// Set a progress callback.
    pub fn with_progress(mut self, cb: ProgressCallback) -> Self {
        self.progress_cb = Some(cb);
        self
    }

    /// Check if GPU will be used for this build.
    pub fn will_use_gpu(&self) -> bool {
        self.use_gpu
    }

    /// Get the GPU availability status.
    pub fn gpu_status(&self) -> &GpuAvailability {
        &self.gpu_availability
    }

    /// Report progress.
    fn report(&self, stage: PipelineStage, progress: f64) {
        if let Some(ref cb) = self.progress_cb {
            cb(stage, progress);
        }
    }

    // ========================================================================
    // IVF Index Building
    // ========================================================================

    /// Build an IVF-based index (IVF_FLAT, IVF_SQ, IVF_PQ).
    ///
    /// Steps:
    /// 1. Train KMeans centroids on GPU/CPU
    /// 2. Train quantization codebook (if applicable) on GPU/CPU
    /// 3. Assign vectors to partitions (streaming, GPU/CPU)
    ///
    /// Returns the trained centroids and assignment.
    pub fn build_ivf(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        n_partitions: usize,
    ) -> GpuResult<IvfBuildResult> {
        let start = Instant::now();
        let mut stage_times = Vec::new();

        // Stage 1: KMeans training
        self.report(PipelineStage::KMeansTraining, 0.0);
        let kmeans_start = Instant::now();

        let kmeans_config = GpuKMeansConfig {
            n_clusters: n_partitions,
            max_iterations: self.gpu_config.effective_max_kmeans_iterations(),
            tolerance: self.gpu_config.effective_kmeans_tolerance(),
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(kmeans_config);
        let kmeans_result = trainer.train(vectors, n_vectors, dimension, self.memory_pool.as_ref())?;

        stage_times.push((
            "KMeans Training".into(),
            kmeans_start.elapsed().as_millis() as u64,
        ));
        self.report(PipelineStage::KMeansTraining, 100.0);

        // Stage 2: Assign vectors to nearest centroid
        // (Already done by KMeans — we have assignments in the result)
        self.report(PipelineStage::Assignment, 0.0);
        let assign_start = Instant::now();

        // If the KMeans result already has assignments, use them directly.
        // Otherwise, compute assignments from centroids.
        let assignments = if kmeans_result.assignments.len() == n_vectors {
            kmeans_result.assignments.clone()
        } else {
            self.assign_vectors(vectors, n_vectors, dimension, &kmeans_result.centroids, n_partitions)?
        };

        stage_times.push((
            "Assignment".into(),
            assign_start.elapsed().as_millis() as u64,
        ));
        self.report(PipelineStage::Assignment, 100.0);

        let used_gpu_for_kmeans = kmeans_result.used_gpu;
        let had_fallback = self.use_gpu && !used_gpu_for_kmeans;

        let device_used = if !self.use_gpu {
            DeviceUsed::Cpu
        } else if had_fallback {
            DeviceUsed::GpuHybrid
        } else {
            DeviceUsed::Gpu
        };

        Ok(IvfBuildResult {
            centroids: kmeans_result.centroids,
            assignments,
            n_partitions,
            pipeline_result: PipelineResult {
                device_used,
                total_time_ms: start.elapsed().as_millis() as u64,
                stage_times,
                gpu_device: self.gpu_device_info(),
                num_vectors: n_vectors,
                dimension,
                index_type: "IVF".into(),
                had_cpu_fallback: had_fallback,
                num_centroids: Some(n_partitions),
                cagra_stats: None,
            },
        })
    }

    /// Assign vectors to their nearest centroid.
    fn assign_vectors(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        centroids: &[f32],
        n_centroids: usize,
    ) -> GpuResult<Vec<u32>> {
        use rayon::prelude::*;

        let mut assignments = vec![0u32; n_vectors];

        // Process in parallel on CPU (simple and fast enough)
        assignments
            .par_iter_mut()
            .enumerate()
            .for_each(|(i, assign)| {
                let v = &vectors[i * dimension..(i + 1) * dimension];
                let mut best_dist = f32::MAX;
                let mut best_idx = 0u32;

                for c in 0..n_centroids {
                    let centroid = &centroids[c * dimension..(c + 1) * dimension];
                    let dist: f32 =
                        v.iter().zip(centroid).map(|(a, b)| (a - b) * (a - b)).sum();
                    if dist < best_dist {
                        best_dist = dist;
                        best_idx = c as u32;
                    }
                }

                *assign = best_idx;
            });

        Ok(assignments)
    }

    // ========================================================================
    // CAGRA Index Building
    // ========================================================================

    /// Build a CAGRA graph index.
    ///
    /// Tries GPU first, falls back to CPU k-NN graph.
    pub fn build_cagra(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        cagra_config: Option<CagraConfig>,
    ) -> GpuResult<CagraBuildResult> {
        let start = Instant::now();

        self.report(PipelineStage::GraphBuild, 0.0);

        let config = cagra_config.unwrap_or_else(|| CagraConfig {
            graph_degree: self.gpu_config.effective_cagra_graph_degree(),
            intermediate_graph_degree: self
                .gpu_config
                .effective_cagra_intermediate_graph_degree(),
            ..Default::default()
        });

        let index = CagraIndex::build(
            vectors,
            n_vectors,
            dimension,
            &config,
            self.memory_pool.as_ref(),
        )?;

        let stats = index.build_stats();
        self.report(PipelineStage::GraphBuild, 100.0);

        let device_used = if index.built_on_gpu() {
            DeviceUsed::Gpu
        } else {
            DeviceUsed::Cpu
        };

        Ok(CagraBuildResult {
            index,
            pipeline_result: PipelineResult {
                device_used,
                total_time_ms: start.elapsed().as_millis() as u64,
                stage_times: vec![(
                    "Graph Build".into(),
                    start.elapsed().as_millis() as u64,
                )],
                gpu_device: self.gpu_device_info(),
                num_vectors: n_vectors,
                dimension,
                index_type: "CAGRA".into(),
                had_cpu_fallback: !stats.used_gpu,
                num_centroids: None,
                cagra_stats: Some(stats),
            },
        })
    }

    // ========================================================================
    // Helpers
    // ========================================================================

    fn gpu_device_info(&self) -> Option<GpuDeviceInfo> {
        self.gpu_availability.best_device().cloned()
    }
}

/// Result of IVF index building.
pub struct IvfBuildResult {
    /// Trained centroids [n_partitions × dimension]
    pub centroids: Vec<f32>,
    /// Partition assignment for each vector
    pub assignments: Vec<u32>,
    /// Number of partitions
    pub n_partitions: usize,
    /// Pipeline execution result
    pub pipeline_result: PipelineResult,
}

/// Result of CAGRA index building.
pub struct CagraBuildResult {
    /// Built CAGRA graph index
    pub index: CagraIndex,
    /// Pipeline execution result
    pub pipeline_result: PipelineResult,
}

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

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

    fn generate_data(n: usize, dim: usize) -> Vec<f32> {
        (0..n * dim)
            .map(|i| ((i * 7 + 13) % 1000) as f32 / 1000.0)
            .collect()
    }

    fn test_gpu_config() -> GpuIndexConfig {
        GpuIndexConfig {
            mode: GpuBuildMode::Auto,
            ..Default::default()
        }
    }

    #[test]
    fn test_pipeline_ivf_cpu() {
        let n = 500;
        let dim = 8;
        let data = generate_data(n, dim);

        let pipeline = IndexBuildPipeline::new(test_gpu_config());
        let result = pipeline.build_ivf(&data, n, dim, 4).unwrap();

        assert_eq!(result.assignments.len(), n);
        assert_eq!(result.n_partitions, 4);
        assert_eq!(result.centroids.len(), 4 * dim);

        // All assignments should be valid partition IDs
        for &a in &result.assignments {
            assert!(a < 4, "Invalid assignment: {}", a);
        }

        let pr = &result.pipeline_result;
        assert_eq!(pr.num_vectors, n);
        assert_eq!(pr.dimension, dim);
        assert_eq!(pr.num_centroids, Some(4));
    }

    #[test]
    fn test_pipeline_cagra_cpu() {
        let n = 200;
        let dim = 4;
        let data = generate_data(n, dim);

        let pipeline = IndexBuildPipeline::new(test_gpu_config());

        let config = CagraConfig {
            graph_degree: 8,
            ..Default::default()
        };

        let result = pipeline.build_cagra(&data, n, dim, Some(config)).unwrap();

        assert_eq!(result.index.len(), n);
        assert_eq!(result.index.graph_degree(), 8);

        let pr = &result.pipeline_result;
        assert_eq!(pr.index_type, "CAGRA");
        assert!(pr.cagra_stats.is_some());
    }

    #[test]
    fn test_pipeline_progress_callback() {
        use std::sync::Mutex;

        let stages = Arc::new(Mutex::new(Vec::new()));
        let stages_clone = stages.clone();

        let pipeline = IndexBuildPipeline::new(test_gpu_config()).with_progress(Box::new(
            move |stage, progress| {
                stages_clone.lock().unwrap().push((stage, progress));
            },
        ));

        let data = generate_data(100, 4);
        pipeline.build_ivf(&data, 100, 4, 4).unwrap();

        let recorded = stages.lock().unwrap();
        assert!(!recorded.is_empty(), "No progress events recorded");

        // Should have KMeansTraining and Assignment stages
        let stage_names: Vec<_> = recorded.iter().map(|(s, _)| *s).collect();
        assert!(stage_names.contains(&PipelineStage::KMeansTraining));
        assert!(stage_names.contains(&PipelineStage::Assignment));
    }

    #[test]
    fn test_pipeline_will_use_gpu() {
        let pipeline = IndexBuildPipeline::new(test_gpu_config());
        // In CI/test environments, GPU is typically not available
        // The pipeline should still work (CPU fallback)
        let _ = pipeline.will_use_gpu();
    }

    #[test]
    fn test_pipeline_result_display() {
        assert_eq!(format!("{}", DeviceUsed::Cpu), "CPU");
        assert_eq!(format!("{}", DeviceUsed::Gpu), "GPU");
        assert_eq!(format!("{}", DeviceUsed::GpuHybrid), "GPU+CPU Hybrid");
    }

    #[test]
    fn test_pipeline_stage_display() {
        assert_eq!(
            format!("{}", PipelineStage::KMeansTraining),
            "KMeans Training"
        );
        assert_eq!(format!("{}", PipelineStage::Completed), "Completed");
    }

    #[test]
    fn test_ivf_partition_distribution() {
        // Use clustered data to verify partitions make sense
        let dim = 4;
        let mut data = Vec::new();
        // Cluster 0: near origin
        for i in 0..50 {
            for d in 0..dim {
                data.push((i * 3 + d) as f32 / 1000.0);
            }
        }
        // Cluster 1: far from origin
        for i in 0..50 {
            for d in 0..dim {
                data.push(100.0 + (i * 3 + d) as f32 / 1000.0);
            }
        }

        let pipeline = IndexBuildPipeline::new(test_gpu_config());
        let result = pipeline.build_ivf(&data, 100, dim, 2).unwrap();

        // Vectors in the same cluster should mostly get the same assignment
        let cluster0_assign = result.assignments[0];
        let cluster1_assign = result.assignments[50];
        assert_ne!(
            cluster0_assign, cluster1_assign,
            "Two distinct clusters should get different partition IDs"
        );

        // Most vectors in cluster 0 should share the same assignment
        let same_as_first: usize = result.assignments[..50]
            .iter()
            .filter(|&&a| a == cluster0_assign)
            .count();
        assert!(
            same_as_first >= 40,
            "At least 80% of cluster 0 should be in the same partition, got {}%",
            same_as_first * 2
        );
    }
}
