//! Streaming batch builder for GPU index construction
//!
//! Handles datasets that exceed GPU memory by splitting them into batches
//! and processing them with overlapping transfer/compute (double buffering).
//!
//! ## Double Buffering Strategy
//!
//! ```text
//! Time →
//! GPU Compute:  [Batch 0 compute] [Batch 1 compute] [Batch 2 compute] ...
//! GPU Transfer:          [Batch 1 upload]  [Batch 2 upload]  [Batch 3 upload]
//! CPU Prepare:  [Batch 1 prep] [Batch 2 prep] [Batch 3 prep]
//! ```
//!
//! While GPU processes batch N, CPU prepares batch N+1 and transfers it.
//! This hides transfer latency behind compute time.
//!
//! ## Auto-degradation
//!
//! If GPU encounters OOM mid-batch, the builder:
//! 1. Halves the batch size
//! 2. Retries the failed batch
//! 3. If still OOM after 3 retries → falls back to full CPU mode

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use super::error::{GpuError, GpuResult};
use super::memory::{GpuMemoryPool, MemoryReservation};

/// Configuration for the streaming batch builder.
#[derive(Debug, Clone)]
pub struct BatchBuilderConfig {
    /// Initial batch size (number of vectors). 0 = auto-compute from GPU memory.
    pub initial_batch_size: usize,

    /// Minimum batch size before giving up and falling back to CPU.
    pub min_batch_size: usize,

    /// Maximum number of OOM retries (with halved batch size) before CPU fallback.
    pub max_oom_retries: usize,

    /// Whether to enable double buffering (overlap transfer and compute).
    pub double_buffering: bool,

    /// Prefetch distance: how many batches ahead to prepare on CPU.
    pub prefetch_distance: usize,

    /// Whether to log progress for each batch.
    pub verbose: bool,
}

impl Default for BatchBuilderConfig {
    fn default() -> Self {
        Self {
            initial_batch_size: 0,
            min_batch_size: 1024,
            max_oom_retries: 3,
            double_buffering: true,
            prefetch_distance: 1,
            verbose: false,
        }
    }
}

/// Progress tracking for batch processing.
#[derive(Debug, Clone)]
pub struct BatchProgress {
    /// Total number of vectors to process
    pub total_vectors: usize,
    /// Number of vectors processed so far
    pub processed_vectors: usize,
    /// Current batch index
    pub current_batch: usize,
    /// Total number of batches
    pub total_batches: usize,
    /// Current effective batch size
    pub batch_size: usize,
    /// Number of OOM retries so far
    pub oom_retries: usize,
    /// Whether processing has fallen back to CPU
    pub fallen_back_to_cpu: bool,
    /// Elapsed time in milliseconds
    pub elapsed_ms: u64,
    /// Estimated remaining time in milliseconds
    pub estimated_remaining_ms: u64,
}

impl BatchProgress {
    /// Progress percentage (0.0 - 100.0)
    pub fn percent(&self) -> f64 {
        if self.total_vectors == 0 {
            return 100.0;
        }
        (self.processed_vectors as f64 / self.total_vectors as f64) * 100.0
    }
}

impl std::fmt::Display for BatchProgress {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Batch {}/{} ({:.1}%): {}/{} vectors, batch_size={}, {}",
            self.current_batch,
            self.total_batches,
            self.percent(),
            self.processed_vectors,
            self.total_vectors,
            self.batch_size,
            if self.fallen_back_to_cpu {
                "CPU fallback"
            } else {
                "GPU"
            },
        )
    }
}

/// A prepared batch of vectors ready for GPU processing.
#[derive(Debug)]
pub struct VectorBatch {
    /// Batch index (0-based)
    pub batch_index: usize,

    /// Start index in the original vector array
    pub start_index: usize,

    /// Number of vectors in this batch
    pub num_vectors: usize,

    /// Vector dimension
    pub dimension: usize,

    /// Flat vector data [num_vectors × dimension]
    pub data: Vec<f32>,
}

impl VectorBatch {
    /// Size in bytes
    pub fn size_bytes(&self) -> usize {
        self.data.len() * std::mem::size_of::<f32>()
    }
}

/// Streaming batch builder that splits large datasets for GPU processing.
///
/// ## Usage Pattern
///
/// ```ignore
/// let builder = StreamingBatchBuilder::new(config, memory_pool);
/// let batches = builder.plan(vectors, n_vectors, dimension);
///
/// for batch in batches {
///     // Process on GPU or CPU
///     let result = process_batch(&batch)?;
///     builder.report_batch_complete(batch.batch_index);
/// }
/// ```
pub struct StreamingBatchBuilder {
    config: BatchBuilderConfig,
    /// Current effective batch size (may decrease on OOM)
    effective_batch_size: AtomicUsize,
    /// Number of OOM retries performed
    oom_retry_count: AtomicUsize,
    /// Whether we've fallen back to CPU
    cpu_fallback: AtomicBool,
    /// Total vectors being processed
    total_vectors: AtomicUsize,
    /// Vectors processed so far
    processed_vectors: AtomicUsize,
    /// Start time for progress estimation
    start_time: std::time::Instant,
}

impl StreamingBatchBuilder {
    /// Create a new streaming batch builder.
    pub fn new(config: BatchBuilderConfig) -> Self {
        Self {
            effective_batch_size: AtomicUsize::new(config.initial_batch_size),
            oom_retry_count: AtomicUsize::new(0),
            cpu_fallback: AtomicBool::new(false),
            total_vectors: AtomicUsize::new(0),
            processed_vectors: AtomicUsize::new(0),
            start_time: std::time::Instant::now(),
            config,
        }
    }

    /// Compute the optimal batch size for the given data and GPU memory.
    ///
    /// Returns the number of vectors per batch that fits in the memory budget.
    pub fn compute_batch_size(
        &self,
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> usize {
        if let Some(explicit) = Some(self.config.initial_batch_size).filter(|&s| s > 0) {
            return explicit.min(n_vectors);
        }

        // Auto-compute based on GPU memory
        if let Some(pool) = memory_pool {
            let batch_size = pool.optimal_batch_size(dimension);
            return batch_size.min(n_vectors).max(self.config.min_batch_size);
        }

        // No memory pool — use a conservative default
        let bytes_per_vector = dimension * std::mem::size_of::<f32>();
        let target_bytes = 256 * 1024 * 1024; // 256MB default
        let computed = if bytes_per_vector > 0 {
            target_bytes / bytes_per_vector
        } else {
            n_vectors
        };
        computed.min(n_vectors).max(self.config.min_batch_size)
    }

    /// Plan batches for the given dataset.
    ///
    /// Returns a vector of `VectorBatch` structs, each containing a slice of the data.
    pub fn plan_batches(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> Vec<VectorBatch> {
        let batch_size = self.compute_batch_size(n_vectors, dimension, memory_pool);
        self.effective_batch_size.store(batch_size, Ordering::Relaxed);
        self.total_vectors.store(n_vectors, Ordering::Relaxed);
        self.processed_vectors.store(0, Ordering::Relaxed);

        let num_batches = (n_vectors + batch_size - 1) / batch_size;
        let mut batches = Vec::with_capacity(num_batches);

        for b in 0..num_batches {
            let start = b * batch_size;
            let end = (start + batch_size).min(n_vectors);
            let count = end - start;

            let data_start = start * dimension;
            let data_end = end * dimension;
            let data = vectors[data_start..data_end].to_vec();

            batches.push(VectorBatch {
                batch_index: b,
                start_index: start,
                num_vectors: count,
                dimension,
                data,
            });
        }

        if self.config.verbose {
            tracing::info!(
                "Planned {} batches: batch_size={}, total_vectors={}, dimension={}",
                num_batches,
                batch_size,
                n_vectors,
                dimension
            );
        }

        batches
    }

    /// Report that a batch has been processed successfully.
    pub fn report_batch_complete(&self, _batch_index: usize, num_vectors: usize) {
        self.processed_vectors
            .fetch_add(num_vectors, Ordering::Relaxed);
    }

    /// Handle an OOM error during batch processing.
    ///
    /// Returns:
    /// - `Ok(new_batch_size)` if we should retry with a smaller batch
    /// - `Err(GpuError)` if we've exhausted retries and should fall back to CPU
    pub fn handle_oom(&self) -> GpuResult<usize> {
        let retries = self.oom_retry_count.fetch_add(1, Ordering::Relaxed);

        if retries >= self.config.max_oom_retries {
            self.cpu_fallback.store(true, Ordering::Relaxed);
            tracing::warn!(
                "GPU OOM after {} retries, falling back to CPU",
                retries + 1
            );
            return Err(GpuError::OutOfMemory {
                requested_bytes: 0,
                available_bytes: 0,
                device_name: "GPU".into(),
            });
        }

        let current = self.effective_batch_size.load(Ordering::Relaxed);
        let new_size = (current / 2).max(self.config.min_batch_size);
        self.effective_batch_size.store(new_size, Ordering::Relaxed);

        tracing::warn!(
            "GPU OOM (retry {}/{}): reducing batch size {} → {}",
            retries + 1,
            self.config.max_oom_retries,
            current,
            new_size,
        );

        Ok(new_size)
    }

    /// Check if we've fallen back to CPU mode.
    pub fn is_cpu_fallback(&self) -> bool {
        self.cpu_fallback.load(Ordering::Relaxed)
    }

    /// Get the current effective batch size.
    pub fn effective_batch_size(&self) -> usize {
        self.effective_batch_size.load(Ordering::Relaxed)
    }

    /// Get current progress.
    pub fn progress(&self) -> BatchProgress {
        let total = self.total_vectors.load(Ordering::Relaxed);
        let processed = self.processed_vectors.load(Ordering::Relaxed);
        let batch_size = self.effective_batch_size.load(Ordering::Relaxed);
        let elapsed = self.start_time.elapsed().as_millis() as u64;

        let total_batches = if batch_size > 0 {
            (total + batch_size - 1) / batch_size
        } else {
            0
        };

        let current_batch = if batch_size > 0 {
            processed / batch_size
        } else {
            0
        };

        let estimated_remaining = if processed > 0 {
            let rate = processed as f64 / elapsed as f64;
            let remaining = total.saturating_sub(processed);
            (remaining as f64 / rate) as u64
        } else {
            0
        };

        BatchProgress {
            total_vectors: total,
            processed_vectors: processed,
            current_batch,
            total_batches,
            batch_size,
            oom_retries: self.oom_retry_count.load(Ordering::Relaxed),
            fallen_back_to_cpu: self.is_cpu_fallback(),
            elapsed_ms: elapsed,
            estimated_remaining_ms: estimated_remaining,
        }
    }

    /// Replan batches after an OOM (with smaller batch size).
    ///
    /// Only replans the remaining unprocessed batches.
    pub fn replan_remaining(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        already_processed: usize,
    ) -> Vec<VectorBatch> {
        let batch_size = self.effective_batch_size.load(Ordering::Relaxed);
        let remaining = n_vectors.saturating_sub(already_processed);

        if remaining == 0 || batch_size == 0 {
            return Vec::new();
        }

        let num_batches = (remaining + batch_size - 1) / batch_size;
        let mut batches = Vec::with_capacity(num_batches);

        for b in 0..num_batches {
            let offset = already_processed + b * batch_size;
            let end = (offset + batch_size).min(n_vectors);
            let count = end - offset;

            let data_start = offset * dimension;
            let data_end = end * dimension;
            let data = vectors[data_start..data_end].to_vec();

            batches.push(VectorBatch {
                batch_index: b,
                start_index: offset,
                num_vectors: count,
                dimension,
                data,
            });
        }

        batches
    }
}

// ============================================================================
// Double Buffer Manager
// ============================================================================

/// Manages two GPU buffer slots for overlapping transfer and compute.
///
/// While slot A is being processed on GPU, slot B receives the next batch.
/// After A finishes, they swap roles.
pub struct DoubleBuffer {
    /// Current active slot (0 or 1)
    active_slot: usize,
    /// Size of each slot in bytes
    slot_size_bytes: usize,
    /// Whether double buffering is enabled (falls back to single buffer if false)
    enabled: bool,
}

impl DoubleBuffer {
    /// Create a new double buffer manager.
    pub fn new(slot_size_bytes: usize, enabled: bool) -> Self {
        Self {
            active_slot: 0,
            slot_size_bytes,
            enabled,
        }
    }

    /// Get the index of the active compute slot.
    pub fn compute_slot(&self) -> usize {
        self.active_slot
    }

    /// Get the index of the transfer slot (the other one).
    pub fn transfer_slot(&self) -> usize {
        if self.enabled {
            1 - self.active_slot
        } else {
            self.active_slot // Same slot when disabled
        }
    }

    /// Swap active and transfer slots.
    pub fn swap(&mut self) {
        if self.enabled {
            self.active_slot = 1 - self.active_slot;
        }
    }

    /// Total memory required for both slots.
    pub fn total_memory_bytes(&self) -> usize {
        if self.enabled {
            self.slot_size_bytes * 2
        } else {
            self.slot_size_bytes
        }
    }

    /// Whether double buffering is active.
    pub fn is_enabled(&self) -> bool {
        self.enabled
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::gpu::device::GpuDeviceInfo;

    fn test_device_info() -> GpuDeviceInfo {
        GpuDeviceInfo {
            device_id: 0,
            name: "Test GPU".into(),
            compute_capability: (8, 0),
            total_memory_bytes: 8 * 1024 * 1024 * 1024,
            free_memory_bytes: 6 * 1024 * 1024 * 1024,
            sm_count: 80,
            is_compatible: true,
        }
    }

    #[test]
    fn test_batch_planning() {
        let config = BatchBuilderConfig {
            initial_batch_size: 100,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);

        let vectors: Vec<f32> = (0..1000 * 4).map(|i| i as f32).collect();
        let batches = builder.plan_batches(&vectors, 1000, 4, None);

        assert_eq!(batches.len(), 10); // 1000 / 100 = 10 batches
        assert_eq!(batches[0].num_vectors, 100);
        assert_eq!(batches[0].start_index, 0);
        assert_eq!(batches[0].data.len(), 100 * 4);
        assert_eq!(batches[9].num_vectors, 100);
        assert_eq!(batches[9].start_index, 900);
    }

    #[test]
    fn test_batch_planning_uneven() {
        let config = BatchBuilderConfig {
            initial_batch_size: 300,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);

        let vectors: Vec<f32> = (0..1000 * 2).map(|i| i as f32).collect();
        let batches = builder.plan_batches(&vectors, 1000, 2, None);

        assert_eq!(batches.len(), 4); // ceil(1000/300) = 4
        assert_eq!(batches[0].num_vectors, 300);
        assert_eq!(batches[1].num_vectors, 300);
        assert_eq!(batches[2].num_vectors, 300);
        assert_eq!(batches[3].num_vectors, 100); // Last batch is smaller
    }

    #[test]
    fn test_auto_batch_size_with_pool() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(512 * 1024 * 1024)); // 512MB

        let config = BatchBuilderConfig::default();
        let builder = StreamingBatchBuilder::new(config);

        let batch_size = builder.compute_batch_size(1_000_000, 768, Some(&pool));
        // Should be reasonable — not the full million, not tiny
        assert!(batch_size >= 1024);
        assert!(batch_size <= 1_000_000);
    }

    #[test]
    fn test_auto_batch_size_no_pool() {
        let config = BatchBuilderConfig::default();
        let builder = StreamingBatchBuilder::new(config);

        let batch_size = builder.compute_batch_size(1_000_000, 768, None);
        // Should use 256MB default: 256MB / (768*4) ≈ 87381
        assert!(batch_size >= 1024);
        assert!(batch_size <= 1_000_000);
    }

    #[test]
    fn test_oom_retry() {
        let config = BatchBuilderConfig {
            initial_batch_size: 10000,
            min_batch_size: 1000,
            max_oom_retries: 3,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);
        builder.effective_batch_size.store(10000, Ordering::Relaxed);

        // First OOM → halve to 5000
        let new_size = builder.handle_oom().unwrap();
        assert_eq!(new_size, 5000);

        // Second OOM → halve to 2500
        let new_size = builder.handle_oom().unwrap();
        assert_eq!(new_size, 2500);

        // Third OOM → halve to 1250
        let new_size = builder.handle_oom().unwrap();
        assert_eq!(new_size, 1250);

        // Fourth OOM → error (max retries exceeded)
        assert!(builder.handle_oom().is_err());
        assert!(builder.is_cpu_fallback());
    }

    #[test]
    fn test_oom_min_batch_size() {
        let config = BatchBuilderConfig {
            initial_batch_size: 2000,
            min_batch_size: 1500,
            max_oom_retries: 5,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);
        builder.effective_batch_size.store(2000, Ordering::Relaxed);

        // OOM → halve to max(1000, min_batch=1500) = 1500
        let new_size = builder.handle_oom().unwrap();
        assert_eq!(new_size, 1500);
    }

    #[test]
    fn test_progress_tracking() {
        let config = BatchBuilderConfig {
            initial_batch_size: 100,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);

        let vectors: Vec<f32> = (0..500 * 4).map(|i| i as f32).collect();
        let batches = builder.plan_batches(&vectors, 500, 4, None);

        assert_eq!(batches.len(), 5);

        let progress = builder.progress();
        assert_eq!(progress.total_vectors, 500);
        assert_eq!(progress.processed_vectors, 0);
        assert!(!progress.fallen_back_to_cpu);

        builder.report_batch_complete(0, 100);
        builder.report_batch_complete(1, 100);

        let progress = builder.progress();
        assert_eq!(progress.processed_vectors, 200);
        assert!((progress.percent() - 40.0).abs() < 0.1);
    }

    #[test]
    fn test_replan_remaining() {
        let config = BatchBuilderConfig {
            initial_batch_size: 100,
            min_batch_size: 50,
            ..Default::default()
        };
        let builder = StreamingBatchBuilder::new(config);
        builder.effective_batch_size.store(50, Ordering::Relaxed);

        let vectors: Vec<f32> = (0..500 * 4).map(|i| i as f32).collect();
        let batches = builder.replan_remaining(&vectors, 500, 4, 200);

        // 300 remaining / 50 = 6 batches
        assert_eq!(batches.len(), 6);
        assert_eq!(batches[0].start_index, 200);
        assert_eq!(batches[0].num_vectors, 50);
        assert_eq!(batches[5].start_index, 450);
        assert_eq!(batches[5].num_vectors, 50);
    }

    #[test]
    fn test_double_buffer() {
        let db = DoubleBuffer::new(1024 * 1024, true);
        assert_eq!(db.compute_slot(), 0);
        assert_eq!(db.transfer_slot(), 1);
        assert_eq!(db.total_memory_bytes(), 2 * 1024 * 1024);
        assert!(db.is_enabled());
    }

    #[test]
    fn test_double_buffer_swap() {
        let mut db = DoubleBuffer::new(1024 * 1024, true);
        assert_eq!(db.compute_slot(), 0);
        assert_eq!(db.transfer_slot(), 1);

        db.swap();
        assert_eq!(db.compute_slot(), 1);
        assert_eq!(db.transfer_slot(), 0);

        db.swap();
        assert_eq!(db.compute_slot(), 0);
        assert_eq!(db.transfer_slot(), 1);
    }

    #[test]
    fn test_double_buffer_disabled() {
        let mut db = DoubleBuffer::new(1024 * 1024, false);
        assert_eq!(db.compute_slot(), 0);
        assert_eq!(db.transfer_slot(), 0); // Same slot when disabled
        assert_eq!(db.total_memory_bytes(), 1024 * 1024); // Single buffer
        assert!(!db.is_enabled());

        db.swap(); // No-op when disabled
        assert_eq!(db.compute_slot(), 0);
    }

    #[test]
    fn test_batch_progress_display() {
        let progress = BatchProgress {
            total_vectors: 10000,
            processed_vectors: 5000,
            current_batch: 5,
            total_batches: 10,
            batch_size: 1000,
            oom_retries: 0,
            fallen_back_to_cpu: false,
            elapsed_ms: 500,
            estimated_remaining_ms: 500,
        };
        let s = format!("{progress}");
        assert!(s.contains("50.0%"));
        assert!(s.contains("5/10"));
        assert!(s.contains("GPU"));
    }

    #[test]
    fn test_vector_batch_size() {
        let batch = VectorBatch {
            batch_index: 0,
            start_index: 0,
            num_vectors: 1000,
            dimension: 768,
            data: vec![0.0f32; 1000 * 768],
        };
        assert_eq!(batch.size_bytes(), 1000 * 768 * 4); // f32 = 4 bytes
    }
}
