//! Write Shard for parallel write processing
//!
//! Each shard has its own write queue, batch aggregator, and WAL writer,
//! allowing independent parallel processing.

use std::sync::Arc;
use std::path::PathBuf;
use parking_lot::Mutex;

use super::write_queue::{WriteQueue, WriteQueueConfig, WriteQueueError, WriteResult};
use super::batch_aggregator::{BatchConfig, BatchWorker, BatchFlushHandler};
use super::metrics::WriteMetrics;
use crate::types::MemoryNode;

/// Configuration for a write shard
#[derive(Debug, Clone)]
pub struct WriteShardConfig {
    /// Shard ID
    pub id: usize,
    
    /// Write queue configuration
    pub queue_config: WriteQueueConfig,
    
    /// Batch aggregator configuration
    pub batch_config: BatchConfig,
    
    /// Path for shard data (WAL, etc.)
    pub path: PathBuf,
}

impl WriteShardConfig {
    /// Create config for a shard
    pub fn new(id: usize, base_path: impl Into<PathBuf>) -> Self {
        let base_path = base_path.into();
        Self {
            id,
            queue_config: WriteQueueConfig::default(),
            batch_config: BatchConfig::default(),
            path: base_path.join(format!("shard_{}", id)),
        }
    }
    
    /// Set queue configuration
    pub fn with_queue_config(mut self, config: WriteQueueConfig) -> Self {
        self.queue_config = config;
        self
    }
    
    /// Set batch configuration
    pub fn with_batch_config(mut self, config: BatchConfig) -> Self {
        self.batch_config = config;
        self
    }
}

/// Write Shard
///
/// A single shard that processes writes independently.
pub struct WriteShard {
    /// Shard ID
    id: usize,
    
    /// Write queue
    queue: WriteQueue,
    
    /// Configuration
    config: WriteShardConfig,
    
    /// Metrics reference
    metrics: Arc<WriteMetrics>,
    
    /// Worker handle (if started) - wrapped in Mutex for interior mutability
    worker_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
}

impl WriteShard {
    /// Create a new write shard (without starting worker)
    pub fn new(config: WriteShardConfig, metrics: Arc<WriteMetrics>) -> Self {
        let queue = WriteQueue::new(config.queue_config.clone(), metrics.clone());
        
        Self {
            id: config.id,
            queue,
            config,
            metrics,
            worker_handle: Mutex::new(None),
        }
    }
    
    /// Start the shard worker with a flush handler
    /// 
    /// Uses interior mutability so this can be called on Arc<WriteShard>
    pub fn start(&self, handler: Arc<dyn BatchFlushHandler>) {
        let mut handle_guard = self.worker_handle.lock();
        
        if handle_guard.is_some() {
            return; // Already started
        }
        
        if let Some(receiver) = self.queue.take_receiver() {
            let worker = BatchWorker::new(
                receiver,
                self.config.batch_config.clone(),
                handler,
                self.metrics.clone(),
                self.id,
            );
            
            let handle = tokio::spawn(worker.run());
            *handle_guard = Some(handle);
            
            tracing::info!("Started write shard {} worker", self.id);
        }
    }
    
    /// Get shard ID
    pub fn id(&self) -> usize {
        self.id
    }
    
    /// Get current queue length
    pub fn queue_len(&self) -> usize {
        self.queue.len()
    }
    
    /// Get queue capacity
    pub fn queue_capacity(&self) -> usize {
        self.queue.capacity()
    }
    
    /// Get queue fill ratio
    pub fn queue_fill_ratio(&self) -> f64 {
        self.queue.fill_ratio()
    }
    
    /// Check if queue is empty
    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }
    
    /// Write a document (async, returns immediately)
    pub async fn write(&self, node: MemoryNode) -> Result<WriteResult, WriteShardError> {
        let sequence = self.queue.enqueue(node.clone()).await
            .map_err(|e| WriteShardError::QueueError(e))?;
        
        Ok(WriteResult::success(node.id, sequence, self.id))
    }
    
    /// Write a document and wait for completion
    pub async fn write_sync(&self, node: MemoryNode) -> Result<WriteResult, WriteShardError> {
        self.queue.enqueue_sync(node).await
            .map_err(|e| WriteShardError::QueueError(e))
    }
    
    /// Try to write without blocking
    pub fn try_write(&self, node: MemoryNode) -> Result<WriteResult, WriteShardError> {
        let id = node.id.clone();
        let sequence = self.queue.try_enqueue(node)
            .map_err(|e| WriteShardError::QueueError(e))?;
        
        Ok(WriteResult::success(id, sequence, self.id))
    }
    
    /// Flush pending writes
    pub async fn flush(&self) -> Result<(), WriteShardError> {
        // Currently, flush is handled by the batch worker's periodic timer
        // In a full implementation, we'd signal the worker to flush immediately
        Ok(())
    }
    
    /// Shutdown the shard
    pub async fn shutdown(&self) -> Result<(), WriteShardError> {
        // Close the queue (signals worker to drain and exit)
        self.queue.close();
        
        // Wait for worker to finish if we have a handle
        // Note: We can't actually await here since we don't own the handle
        // This would need to be redesigned for proper shutdown
        
        tracing::info!("Shutdown requested for shard {}", self.id);
        Ok(())
    }
    
    /// Get configuration
    pub fn config(&self) -> &WriteShardConfig {
        &self.config
    }
    
    /// Get path
    pub fn path(&self) -> &PathBuf {
        &self.config.path
    }
}

/// Errors from write shard
#[derive(Debug, thiserror::Error)]
pub enum WriteShardError {
    #[error("Queue error: {0}")]
    QueueError(#[from] WriteQueueError),
    
    #[error("Worker not started")]
    WorkerNotStarted,
    
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::batch_aggregator::InMemoryBatchHandler;
    use tempfile::tempdir;
    
    #[tokio::test]
    async fn test_write_shard_basic() {
        let dir = tempdir().unwrap();
        let config = WriteShardConfig::new(0, dir.path());
        let metrics = Arc::new(WriteMetrics::new());
        let mut shard = WriteShard::new(config, metrics);
        
        // Start worker
        let handler = Arc::new(InMemoryBatchHandler::new());
        shard.start(handler.clone());
        
        // Write a document
        let node = MemoryNode::new("test content");
        let result = shard.write(node).await.unwrap();
        
        assert!(result.success);
        assert_eq!(result.shard_id, 0);
        
        // Allow time for batch to flush
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        
        // Verify document was written
        assert_eq!(handler.len(), 1);
    }
    
    #[tokio::test]
    async fn test_write_shard_multiple() {
        let dir = tempdir().unwrap();
        let config = WriteShardConfig::new(0, dir.path())
            .with_batch_config(BatchConfig {
                max_batch_size: 5,
                max_wait_ms: 50,
                ..Default::default()
            });
        let metrics = Arc::new(WriteMetrics::new());
        let mut shard = WriteShard::new(config, metrics.clone());
        
        let handler = Arc::new(InMemoryBatchHandler::new());
        shard.start(handler.clone());
        
        // Write multiple documents
        for i in 0..10 {
            let node = MemoryNode::new(&format!("content {}", i));
            shard.write(node).await.unwrap();
        }
        
        // Allow time for batches to flush
        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        
        // Verify all documents were written
        assert_eq!(handler.len(), 10);
        
        // Check metrics
        let snapshot = metrics.snapshot();
        assert_eq!(snapshot.batch_documents, 10);
    }
}
