//! Batch Aggregator for efficient write batching
//!
//! Collects write requests and flushes them in batches based on
//! configurable size, time, and byte thresholds.

use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::interval;

use crate::types::MemoryNode;
use super::write_queue::{WriteRequest, WriteBatch, WriteResult};
use super::metrics::{WriteMetrics, Timer};

/// Configuration for batch aggregation
#[derive(Debug, Clone)]
pub struct BatchConfig {
    /// Maximum number of documents per batch
    pub max_batch_size: usize,
    
    /// Maximum time to wait before flushing (ms)
    pub max_wait_ms: u64,
    
    /// Maximum total bytes per batch
    pub max_batch_bytes: usize,
    
    /// Whether to flush immediately when batch is full
    pub flush_on_full: bool,
}

impl Default for BatchConfig {
    fn default() -> Self {
        Self {
            max_batch_size: 100,
            max_wait_ms: 50,
            max_batch_bytes: 1024 * 1024, // 1MB
            flush_on_full: true,
        }
    }
}

impl BatchConfig {
    /// Create config for high throughput
    pub fn high_throughput() -> Self {
        Self {
            max_batch_size: 500,
            max_wait_ms: 100,
            max_batch_bytes: 10 * 1024 * 1024, // 10MB
            flush_on_full: true,
        }
    }
    
    /// Create config for low latency
    pub fn low_latency() -> Self {
        Self {
            max_batch_size: 20,
            max_wait_ms: 10,
            max_batch_bytes: 256 * 1024, // 256KB
            flush_on_full: true,
        }
    }
}

/// Result of checking if batch should be flushed
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FlushReason {
    /// Batch reached maximum size
    SizeLimit,
    /// Batch reached maximum bytes
    ByteLimit,
    /// Maximum wait time exceeded
    TimeLimit,
    /// Manual flush requested
    Manual,
    /// Shutdown requested
    Shutdown,
}

/// Batch Aggregator that collects and flushes write requests
pub struct BatchAggregator {
    /// Current batch being accumulated
    batch: WriteBatch,
    
    /// Configuration
    config: BatchConfig,
    
    /// Metrics reference
    metrics: Arc<WriteMetrics>,
    
    /// Shard ID this aggregator belongs to
    shard_id: usize,
}

impl BatchAggregator {
    /// Create a new batch aggregator
    pub fn new(config: BatchConfig, metrics: Arc<WriteMetrics>, shard_id: usize) -> Self {
        Self {
            batch: WriteBatch::with_capacity(config.max_batch_size),
            config,
            metrics,
            shard_id,
        }
    }
    
    /// Add a request to the current batch
    pub fn add(&mut self, request: WriteRequest) {
        self.batch.push(request);
    }
    
    /// Check if batch should be flushed
    pub fn should_flush(&self) -> Option<FlushReason> {
        // Check size limit
        if self.batch.len() >= self.config.max_batch_size {
            return Some(FlushReason::SizeLimit);
        }
        
        // Check byte limit
        if self.batch.total_bytes >= self.config.max_batch_bytes {
            return Some(FlushReason::ByteLimit);
        }
        
        // Check time limit
        if !self.batch.is_empty() && 
           self.batch.age() >= Duration::from_millis(self.config.max_wait_ms) {
            return Some(FlushReason::TimeLimit);
        }
        
        None
    }
    
    /// Get the current batch for flushing
    pub fn flush(&mut self) -> Option<WriteBatch> {
        if self.batch.is_empty() {
            return None;
        }
        
        // Take current batch and replace with new empty one
        let batch = std::mem::replace(
            &mut self.batch,
            WriteBatch::with_capacity(self.config.max_batch_size)
        );
        
        Some(batch)
    }
    
    /// Get current batch size
    pub fn len(&self) -> usize {
        self.batch.len()
    }
    
    /// Check if current batch is empty
    pub fn is_empty(&self) -> bool {
        self.batch.is_empty()
    }
    
    /// Get current batch bytes
    pub fn bytes(&self) -> usize {
        self.batch.total_bytes
    }
    
    /// Get batch age
    pub fn age(&self) -> Duration {
        self.batch.age()
    }
    
    /// Get configuration
    pub fn config(&self) -> &BatchConfig {
        &self.config
    }
    
    /// Get shard ID
    pub fn shard_id(&self) -> usize {
        self.shard_id
    }
}

/// Async trait for batch flush handlers
#[async_trait::async_trait]
pub trait BatchFlushHandler: Send + Sync {
    /// Handle a batch of write requests
    async fn handle_batch(&self, batch: WriteBatch, shard_id: usize) -> Result<Vec<WriteResult>, BatchFlushError>;
}

/// Errors during batch flush
#[derive(Debug, thiserror::Error)]
pub enum BatchFlushError {
    #[error("WAL write failed: {0}")]
    WalError(String),
    
    #[error("Document store error: {0}")]
    StoreError(String),
    
    #[error("Batch processing interrupted")]
    Interrupted,
    
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
}

/// Worker that processes batches from a queue
pub struct BatchWorker {
    /// Receiver for write requests
    receiver: mpsc::Receiver<WriteRequest>,
    
    /// Batch aggregator
    aggregator: BatchAggregator,
    
    /// Flush handler
    handler: Arc<dyn BatchFlushHandler>,
    
    /// Metrics reference
    metrics: Arc<WriteMetrics>,
    
    /// Shutdown flag
    shutdown: bool,
}

impl BatchWorker {
    /// Create a new batch worker
    pub fn new(
        receiver: mpsc::Receiver<WriteRequest>,
        config: BatchConfig,
        handler: Arc<dyn BatchFlushHandler>,
        metrics: Arc<WriteMetrics>,
        shard_id: usize,
    ) -> Self {
        Self {
            receiver,
            aggregator: BatchAggregator::new(config.clone(), metrics.clone(), shard_id),
            handler,
            metrics,
            shutdown: false,
        }
    }
    
    /// Run the worker loop
    pub async fn run(mut self) {
        let flush_interval = Duration::from_millis(self.aggregator.config.max_wait_ms);
        let mut interval_timer = interval(flush_interval);
        
        loop {
            tokio::select! {
                // Receive new requests
                request = self.receiver.recv() => {
                    match request {
                        Some(req) => {
                            // Record queue time
                            let queue_time = req.queue_time();
                            self.metrics.record_queue_time(queue_time.as_micros() as u64);
                            
                            // Add to batch
                            self.aggregator.add(req);
                            
                            // Check if we should flush
                            if let Some(reason) = self.aggregator.should_flush() {
                                self.flush_batch(reason).await;
                            }
                        }
                        None => {
                            // Channel closed, flush remaining and exit
                            self.shutdown = true;
                            if !self.aggregator.is_empty() {
                                self.flush_batch(FlushReason::Shutdown).await;
                            }
                            break;
                        }
                    }
                }
                
                // Periodic flush
                _ = interval_timer.tick() => {
                    if let Some(reason) = self.aggregator.should_flush() {
                        self.flush_batch(reason).await;
                    }
                }
            }
            
            if self.shutdown {
                break;
            }
        }
        
        tracing::info!("Batch worker {} shutting down", self.aggregator.shard_id);
    }
    
    /// Flush the current batch
    async fn flush_batch(&mut self, reason: FlushReason) {
        if let Some(batch) = self.aggregator.flush() {
            let batch_size = batch.len();
            let batch_bytes = batch.total_bytes;
            let timer = Timer::start();
            
            tracing::debug!(
                "Flushing batch: shard={}, size={}, bytes={}, reason={:?}",
                self.aggregator.shard_id,
                batch_size,
                batch_bytes,
                reason
            );
            
            match self.handler.handle_batch(batch, self.aggregator.shard_id).await {
                Ok(results) => {
                    // Record metrics
                    self.metrics.record_batch_flush(batch_size, batch_bytes);
                    self.metrics.record_write_time(timer.elapsed_us());
                    
                    // Mark successes
                    for result in &results {
                        if result.success {
                            self.metrics.record_success();
                        } else {
                            self.metrics.record_failure();
                        }
                    }
                    
                    tracing::debug!(
                        "Batch flushed: shard={}, success={}, time_ms={}",
                        self.aggregator.shard_id,
                        results.iter().filter(|r| r.success).count(),
                        timer.elapsed_ms()
                    );
                }
                Err(e) => {
                    tracing::error!(
                        "Batch flush failed: shard={}, error={}",
                        self.aggregator.shard_id,
                        e
                    );
                    
                    // Record all as failures
                    for _ in 0..batch_size {
                        self.metrics.record_failure();
                    }
                }
            }
        }
    }
}

/// Simple in-memory batch handler for testing
pub struct InMemoryBatchHandler {
    /// Storage for documents
    store: Arc<parking_lot::RwLock<Vec<MemoryNode>>>,
}

impl InMemoryBatchHandler {
    pub fn new() -> Self {
        Self {
            store: Arc::new(parking_lot::RwLock::new(Vec::new())),
        }
    }
    
    pub fn documents(&self) -> Vec<MemoryNode> {
        self.store.read().clone()
    }
    
    pub fn len(&self) -> usize {
        self.store.read().len()
    }
}

impl Default for InMemoryBatchHandler {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl BatchFlushHandler for InMemoryBatchHandler {
    async fn handle_batch(&self, batch: WriteBatch, shard_id: usize) -> Result<Vec<WriteResult>, BatchFlushError> {
        let mut results = Vec::with_capacity(batch.len());
        let mut store = self.store.write();
        
        for request in batch.requests {
            let id = request.node.id.clone();
            let sequence = request.sequence;
            
            store.push(request.node.clone());
            
            let result = WriteResult::success(id, sequence, shard_id);
            
            // Notify completion if requested
            request.notify_completion(result.clone());
            
            results.push(result);
        }
        
        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_batch_aggregator_size_limit() {
        let config = BatchConfig {
            max_batch_size: 5,
            max_wait_ms: 1000,
            max_batch_bytes: 1024 * 1024,
            flush_on_full: true,
        };
        let metrics = Arc::new(WriteMetrics::new());
        let mut aggregator = BatchAggregator::new(config, metrics, 0);
        
        // Add 4 requests - should not trigger flush
        for i in 0..4 {
            let node = MemoryNode::new(&format!("content {}", i));
            let request = WriteRequest::new(node, i);
            aggregator.add(request);
        }
        assert!(aggregator.should_flush().is_none());
        
        // Add 5th request - should trigger flush
        let node = MemoryNode::new("content 4");
        let request = WriteRequest::new(node, 4);
        aggregator.add(request);
        assert_eq!(aggregator.should_flush(), Some(FlushReason::SizeLimit));
    }
    
    #[test]
    fn test_batch_aggregator_flush() {
        let config = BatchConfig::default();
        let metrics = Arc::new(WriteMetrics::new());
        let mut aggregator = BatchAggregator::new(config, metrics, 0);
        
        // Add some requests
        for i in 0..3 {
            let node = MemoryNode::new(&format!("content {}", i));
            let request = WriteRequest::new(node, i);
            aggregator.add(request);
        }
        
        assert_eq!(aggregator.len(), 3);
        
        // Flush
        let batch = aggregator.flush().unwrap();
        assert_eq!(batch.len(), 3);
        assert!(aggregator.is_empty());
    }
    
    #[tokio::test]
    async fn test_batch_worker() {
        let (tx, rx) = mpsc::channel(100);
        let config = BatchConfig {
            max_batch_size: 5,
            max_wait_ms: 100,
            ..Default::default()
        };
        let metrics = Arc::new(WriteMetrics::new());
        let handler = Arc::new(InMemoryBatchHandler::new());
        
        let worker = BatchWorker::new(
            rx,
            config,
            handler.clone(),
            metrics.clone(),
            0,
        );
        
        // Spawn worker
        let worker_handle = tokio::spawn(worker.run());
        
        // Send some requests
        for i in 0..10 {
            let node = MemoryNode::new(&format!("content {}", i));
            let request = WriteRequest::new(node, i);
            tx.send(request).await.unwrap();
        }
        
        // Close channel to trigger shutdown
        drop(tx);
        
        // Wait for worker to finish
        worker_handle.await.unwrap();
        
        // Check handler received all documents
        assert_eq!(handler.len(), 10);
        
        // Check metrics
        let snapshot = metrics.snapshot();
        assert_eq!(snapshot.batch_documents, 10);
        assert!(snapshot.batches_flushed >= 1);
    }
}
