//! Concurrent Document Store for high-throughput write operations
//!
//! This is the main entry point for the high-concurrency write system.
//! It combines all components: write queues, batch aggregators, backpressure,
//! sharding, and WAL for a complete solution.

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

use tokio::sync::RwLock;

use crate::types::MemoryNode;
use crate::storage::DocumentStore;
use crate::storage::concurrent::{
    WriteMetrics, MetricsSnapshot,
    WriteQueueConfig,
    BatchConfig, BatchFlushHandler, BatchFlushError, WriteBatch,
    BackpressureController, BackpressureConfig,
    ShardRouter, ShardRouterConfig, ShardStrategy,
    WriteShard, WriteShardConfig,
    WriteResult,
};
use crate::storage::wal::{
    WALWriter, WALWriterConfig, SyncMode,
    WALEntry, WALOperation,
    WALCompactor, CompactorConfig,
    WALReader,
};

/// Configuration for the concurrent document store
#[derive(Debug, Clone)]
pub struct ConcurrentStoreConfig {
    /// Base path for all data
    pub path: PathBuf,
    
    /// Number of shards (0 = auto based on CPU count)
    pub shard_count: usize,
    
    /// Shard routing strategy
    pub shard_strategy: ShardStrategy,
    
    /// Write queue configuration
    pub queue_config: WriteQueueConfig,
    
    /// Batch configuration
    pub batch_config: BatchConfig,
    
    /// Backpressure configuration
    pub backpressure_config: BackpressureConfig,
    
    /// WAL configuration
    pub wal_sync_mode: SyncMode,
    
    /// WAL max file size
    pub wal_max_file_size: usize,
    
    /// Compaction interval (seconds)
    pub compaction_interval_secs: u64,
    
    /// Whether to enable WAL (can be disabled for testing)
    pub enable_wal: bool,
    
    /// Whether to auto-start workers
    pub auto_start: bool,
}

impl Default for ConcurrentStoreConfig {
    fn default() -> Self {
        Self {
            path: PathBuf::from("data"),
            shard_count: 0, // Auto
            shard_strategy: ShardStrategy::Hash,
            queue_config: WriteQueueConfig::default(),
            batch_config: BatchConfig::default(),
            backpressure_config: BackpressureConfig::default(),
            wal_sync_mode: SyncMode::Batch,
            wal_max_file_size: 64 * 1024 * 1024, // 64MB
            compaction_interval_secs: 60,
            enable_wal: true,
            auto_start: true,
        }
    }
}

impl ConcurrentStoreConfig {
    /// Create config with specified path
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            ..Default::default()
        }
    }
    
    /// Set shard count
    pub fn with_shards(mut self, count: usize) -> Self {
        self.shard_count = count;
        self
    }
    
    /// Set shard strategy
    pub fn with_strategy(mut self, strategy: ShardStrategy) -> Self {
        self.shard_strategy = strategy;
        self
    }
    
    /// Configure for high throughput
    pub fn high_throughput(mut self) -> Self {
        self.queue_config = WriteQueueConfig::high_throughput();
        self.batch_config = BatchConfig::high_throughput();
        self.wal_sync_mode = SyncMode::Batch;
        self
    }
    
    /// Configure for low latency
    pub fn low_latency(mut self) -> Self {
        self.queue_config = WriteQueueConfig::low_latency();
        self.batch_config = BatchConfig::low_latency();
        self.wal_sync_mode = SyncMode::Every;
        self
    }
    
    /// Disable WAL (for testing only)
    pub fn without_wal(mut self) -> Self {
        self.enable_wal = false;
        self
    }
}

/// High-concurrency document store
///
/// Provides high-throughput write operations with:
/// - Sharded write queues for parallel processing
/// - Batch aggregation for efficient I/O
/// - Backpressure control for overload protection
/// - Write-ahead logging for durability
/// - Background compaction
pub struct ConcurrentDocumentStore {
    /// Configuration
    config: ConcurrentStoreConfig,
    
    /// Shard router
    router: Arc<ShardRouter>,
    
    /// Shards
    shards: Vec<Arc<WriteShard>>,
    
    /// Document store (for reads and compaction)
    doc_store: Arc<DocumentStore>,
    
    /// WAL writers (one per shard)
    wal_writers: Vec<Arc<RwLock<WALWriter>>>,
    
    /// Metrics
    metrics: Arc<WriteMetrics>,
    
    /// Compactor handle
    compactor_handle: Option<tokio::task::JoinHandle<()>>,
    
    /// Shutdown flag
    shutdown: Arc<std::sync::atomic::AtomicBool>,
}

impl ConcurrentDocumentStore {
    /// Create a new concurrent document store
    pub async fn new(config: ConcurrentStoreConfig) -> std::io::Result<Self> {
        // Create directories
        std::fs::create_dir_all(&config.path)?;
        let wal_path = config.path.join("wal");
        std::fs::create_dir_all(&wal_path)?;
        
        // Determine shard count
        let shard_count = if config.shard_count == 0 {
            num_cpus::get() * 2
        } else {
            config.shard_count
        };
        
        // Create metrics
        let metrics = Arc::new(WriteMetrics::new());
        
        // Create document store
        let docs_path = config.path.join("documents.jsonl");
        let doc_store = Arc::new(DocumentStore::open_or_create(&docs_path)?);
        
        // Create WAL writers and shards
        let mut wal_writers = Vec::with_capacity(shard_count);
        let mut shards = Vec::with_capacity(shard_count);
        
        for i in 0..shard_count {
            // Create WAL writer for this shard
            let wal_config = WALWriterConfig {
                path: wal_path.clone(),
                max_file_size: config.wal_max_file_size,
                sync_mode: config.wal_sync_mode,
                retention_count: 10,
                buffer_size: 64 * 1024,
            };
            
            let wal_writer = if config.enable_wal {
                WALWriter::new(wal_config, i, metrics.clone())?
            } else {
                // Create a dummy WAL writer that doesn't actually write
                WALWriter::new(wal_config, i, metrics.clone())?
            };
            
            wal_writers.push(Arc::new(RwLock::new(wal_writer)));
            
            // Create shard
            let shard_config = WriteShardConfig::new(i, &config.path)
                .with_queue_config(config.queue_config.clone())
                .with_batch_config(config.batch_config.clone());
            
            let shard = WriteShard::new(shard_config, metrics.clone());
            shards.push(Arc::new(shard));
        }
        
        // Create backpressure controller
        let backpressure = BackpressureController::new(
            config.backpressure_config.clone(),
            metrics.clone(),
        );
        
        // Create router config
        let router_config = ShardRouterConfig {
            shard_count,
            strategy: config.shard_strategy,
            auto_balance: true,
        };
        
        // Create router
        let router = Arc::new(ShardRouter::new(
            shards.iter().cloned().collect(),
            router_config,
            backpressure,
            metrics.clone(),
        ));
        
        let shutdown = Arc::new(std::sync::atomic::AtomicBool::new(false));
        
        let mut store = Self {
            config,
            router,
            shards,
            doc_store,
            wal_writers,
            metrics,
            compactor_handle: None,
            shutdown,
        };
        
        // Auto-start if configured
        if store.config.auto_start {
            store.start().await?;
        }
        
        Ok(store)
    }
    
    /// Start the store (workers and compactor)
    pub async fn start(&mut self) -> std::io::Result<()> {
        // Start shard workers
        for (i, shard) in self.shards.iter().enumerate() {
            let handler = Arc::new(ShardBatchHandler {
                doc_store: self.doc_store.clone(),
                wal_writer: self.wal_writers[i].clone(),
                enable_wal: self.config.enable_wal,
            });
            // Start the worker - WriteShard now uses interior mutability
            shard.start(handler);
        }
        
        // Start compactor
        if self.config.enable_wal {
            let compactor_config = CompactorConfig {
                wal_path: self.config.path.join("wal"),
                interval_secs: self.config.compaction_interval_secs,
                min_wal_files: 3,
                min_wal_bytes: 10 * 1024 * 1024,
                retention_count: 2,
                auto_compact: true,
            };
            
            let compactor = WALCompactor::new(compactor_config, self.metrics.clone());
            let doc_store = self.doc_store.clone();
            
            let handle = tokio::spawn(async move {
                compactor.run_loop(doc_store).await;
            });
            
            self.compactor_handle = Some(handle);
        }
        
        tracing::info!(
            "ConcurrentDocumentStore started: shards={}, wal={}",
            self.shards.len(),
            self.config.enable_wal
        );
        
        Ok(())
    }
    
    /// Write a document (async, returns immediately after queuing)
    pub async fn write(&self, node: MemoryNode) -> Result<WriteResult, ConcurrentStoreError> {
        // Write to WAL first if enabled
        if self.config.enable_wal {
            let shard_id = self.router.route(&node);
            let wal_writer = &self.wal_writers[shard_id];
            
            let mut writer = wal_writer.write().await;
            let seq = writer.next_sequence();
            let entry = WALEntry::insert(seq, "default".to_string(), node.clone());
            
            writer.append(&entry)
                .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
        }
        
        // Route to shard
        self.router.write(node).await
            .map_err(|e| ConcurrentStoreError::RouterError(e.to_string()))
    }
    
    /// Write a document and wait for it to be persisted
    pub async fn write_sync(&self, node: MemoryNode) -> Result<WriteResult, ConcurrentStoreError> {
        // Write to WAL first if enabled
        if self.config.enable_wal {
            let shard_id = self.router.route(&node);
            let wal_writer = &self.wal_writers[shard_id];
            
            let mut writer = wal_writer.write().await;
            let seq = writer.next_sequence();
            let entry = WALEntry::insert(seq, "default".to_string(), node.clone());
            
            writer.append(&entry)
                .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
            
            writer.sync()
                .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
        }
        
        self.router.write_sync(node).await
            .map_err(|e| ConcurrentStoreError::RouterError(e.to_string()))
    }
    
    /// Write multiple documents
    pub async fn write_batch(&self, nodes: Vec<MemoryNode>) -> Result<Vec<WriteResult>, ConcurrentStoreError> {
        let mut results = Vec::with_capacity(nodes.len());
        
        for node in nodes {
            let result = self.write(node).await?;
            results.push(result);
        }
        
        Ok(results)
    }
    
    /// Get a document by ID (reads from underlying store)
    pub fn get(&self, id: &str) -> Option<MemoryNode> {
        self.doc_store.get(id)
    }
    
    /// Check if a document exists
    pub fn contains(&self, id: &str) -> bool {
        self.doc_store.contains(id)
    }
    
    /// Get document count
    pub fn len(&self) -> usize {
        self.doc_store.len()
    }
    
    /// Check if store is empty
    pub fn is_empty(&self) -> bool {
        self.doc_store.is_empty()
    }
    
    /// Get all document IDs
    pub fn ids(&self) -> Vec<String> {
        self.doc_store.ids()
    }
    
    /// Get metrics snapshot
    pub fn metrics(&self) -> MetricsSnapshot {
        self.metrics.snapshot()
    }
    
    /// Get current queue fill ratio
    pub fn queue_fill_ratio(&self) -> f64 {
        self.router.get_total_fill_ratio()
    }
    
    /// Flush all pending writes
    pub async fn flush(&self) -> Result<(), ConcurrentStoreError> {
        // Flush all WAL writers
        for wal_writer in &self.wal_writers {
            let mut writer = wal_writer.write().await;
            writer.sync()
                .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
        }
        
        // Flush router/shards
        self.router.flush().await
            .map_err(|e| ConcurrentStoreError::RouterError(e.to_string()))?;
        
        // Flush document store
        self.doc_store.flush()
            .map_err(|e| ConcurrentStoreError::StoreError(e.to_string()))?;
        
        Ok(())
    }
    
    /// Shutdown the store gracefully
    pub async fn shutdown(&self) -> Result<(), ConcurrentStoreError> {
        self.shutdown.store(true, std::sync::atomic::Ordering::SeqCst);
        
        // Flush everything
        self.flush().await?;
        
        // Shutdown router/shards
        self.router.shutdown().await
            .map_err(|e| ConcurrentStoreError::RouterError(e.to_string()))?;
        
        // Close WAL writers
        for wal_writer in &self.wal_writers {
            let mut writer = wal_writer.write().await;
            writer.close()
                .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
        }
        
        tracing::info!("ConcurrentDocumentStore shutdown complete");
        
        Ok(())
    }
    
    /// Recover from WAL after crash
    pub fn recover(&self) -> Result<RecoveryResult, ConcurrentStoreError> {
        let wal_path = self.config.path.join("wal");
        let reader = WALReader::new(&wal_path);
        
        let entries = reader.read_all()
            .map_err(|e| ConcurrentStoreError::WalError(e.to_string()))?;
        
        let mut result = RecoveryResult::default();
        result.entries_found = entries.len();
        
        for entry in entries {
            if !entry.verify_checksum() {
                result.entries_skipped += 1;
                continue;
            }
            
            match entry.operation {
                WALOperation::Insert | WALOperation::Update => {
                    if let Some(node) = entry.data {
                        if self.doc_store.put(node).is_ok() {
                            result.entries_replayed += 1;
                        } else {
                            result.entries_failed += 1;
                        }
                    }
                }
                WALOperation::Delete => {
                    self.doc_store.delete(&entry.doc_id);
                    result.entries_replayed += 1;
                }
            }
            
            result.last_sequence = entry.sequence;
        }
        
        // Flush recovered data
        self.doc_store.flush()
            .map_err(|e| ConcurrentStoreError::StoreError(e.to_string()))?;
        
        Ok(result)
    }
    
    /// Get underlying document store (for direct reads)
    pub fn doc_store(&self) -> &DocumentStore {
        &self.doc_store
    }
    
    /// Get shard count
    pub fn shard_count(&self) -> usize {
        self.shards.len()
    }
}

/// Batch handler for shards
struct ShardBatchHandler {
    doc_store: Arc<DocumentStore>,
    wal_writer: Arc<RwLock<WALWriter>>,
    enable_wal: bool,
}

#[async_trait::async_trait]
impl BatchFlushHandler for ShardBatchHandler {
    async fn handle_batch(&self, batch: WriteBatch, shard_id: usize) -> Result<Vec<WriteResult>, BatchFlushError> {
        let mut results = Vec::with_capacity(batch.len());
        
        // Write to WAL if enabled
        if self.enable_wal {
            let mut writer = self.wal_writer.write().await;
            for request in &batch.requests {
                let seq = writer.next_sequence();
                let entry = WALEntry::insert(seq, "default".to_string(), request.node.clone());
                writer.append(&entry)
                    .map_err(|e| BatchFlushError::WalError(e.to_string()))?;
            }
            writer.sync()
                .map_err(|e| BatchFlushError::WalError(e.to_string()))?;
        }
        
        // Write to document store
        for request in batch.requests {
            let id = request.node.id.clone();
            let sequence = request.sequence;
            
            match self.doc_store.put(request.node.clone()) {
                Ok(()) => {
                    let result = WriteResult::success(id, sequence, shard_id);
                    request.notify_completion(result.clone());
                    results.push(result);
                }
                Err(e) => {
                    let result = WriteResult::failure(id, sequence, e.to_string());
                    request.notify_completion(result.clone());
                    results.push(result);
                }
            }
        }
        
        Ok(results)
    }
}

/// Errors from concurrent store
#[derive(Debug, thiserror::Error)]
pub enum ConcurrentStoreError {
    #[error("WAL error: {0}")]
    WalError(String),
    
    #[error("Router error: {0}")]
    RouterError(String),
    
    #[error("Store error: {0}")]
    StoreError(String),
    
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
}

/// Result of WAL recovery
#[derive(Debug, Default)]
pub struct RecoveryResult {
    /// Number of entries found in WAL
    pub entries_found: usize,
    
    /// Number of entries successfully replayed
    pub entries_replayed: usize,
    
    /// Number of entries skipped (checksum errors)
    pub entries_skipped: usize,
    
    /// Number of entries that failed to replay
    pub entries_failed: usize,
    
    /// Last sequence number seen
    pub last_sequence: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    
    #[tokio::test]
    async fn test_concurrent_store_basic() {
        let dir = tempdir().unwrap();
        let config = ConcurrentStoreConfig::new(dir.path())
            .with_shards(2)
            .without_wal(); // Disable WAL for faster testing
        
        let store = ConcurrentDocumentStore::new(config).await.unwrap();
        
        // Write a document
        let node = MemoryNode::new("test content");
        let id = node.id.clone();
        
        let result = store.write(node).await.unwrap();
        assert!(result.success);
        
        // Allow time for batch processing
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        
        // Check metrics
        let metrics = store.metrics();
        assert!(metrics.total_requests >= 1);
    }
    
    #[tokio::test]
    async fn test_concurrent_store_high_throughput() {
        let dir = tempdir().unwrap();
        let config = ConcurrentStoreConfig::new(dir.path())
            .with_shards(4)
            .high_throughput()
            .without_wal();
        
        let store = ConcurrentDocumentStore::new(config).await.unwrap();
        
        // Write many documents
        for i in 0..100 {
            let node = MemoryNode::new(&format!("content {}", i));
            store.write(node).await.unwrap();
        }
        
        // Allow time for batch processing
        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        
        let metrics = store.metrics();
        assert_eq!(metrics.total_requests, 100);
    }
}
