//! Vector Shard - A single shard containing an HNSW index with write queue

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Instant;
use std::collections::VecDeque;

use parking_lot::{Mutex, RwLock};
use tokio::sync::oneshot;

use crate::types::HNSWConfig;
use crate::vector::hnsw::HNSWIndex;
use crate::storage::concurrent::{WriteMetrics, BatchConfig};

use super::types::{VectorWriteOp, VectorWriteRequest, VectorWriteResult, VectorShardError, VectorSearchResult};
use super::config::ConcurrentHNSWConfig;

/// A single vector shard containing an HNSW index
pub struct VectorShard {
    /// Shard ID
    shard_id: usize,
    
    /// The HNSW index for this shard
    index: RwLock<HNSWIndex>,
    
    /// Write queue
    write_queue: Mutex<VecDeque<VectorWriteRequest>>,
    
    /// Queue capacity
    queue_capacity: usize,
    
    /// Batch configuration
    batch_config: BatchConfig,
    
    /// Sequence counter
    sequence: AtomicU64,
    
    /// Metrics
    metrics: Arc<WriteMetrics>,
    
    /// Pending count
    pending_count: AtomicUsize,
}

impl VectorShard {
    /// Create a new vector shard
    pub fn new(
        shard_id: usize,
        hnsw_config: HNSWConfig,
        config: &ConcurrentHNSWConfig,
        metrics: Arc<WriteMetrics>,
    ) -> Self {
        Self {
            shard_id,
            index: RwLock::new(HNSWIndex::new(hnsw_config)),
            write_queue: Mutex::new(VecDeque::with_capacity(config.queue_config.capacity)),
            queue_capacity: config.queue_config.capacity,
            batch_config: config.batch_config.clone(),
            sequence: AtomicU64::new(0),
            metrics,
            pending_count: AtomicUsize::new(0),
        }
    }
    
    /// Get shard ID
    pub fn shard_id(&self) -> usize {
        self.shard_id
    }
    
    /// Get number of vectors in this shard
    pub fn len(&self) -> usize {
        self.index.read().len()
    }
    
    /// Check if shard is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    
    /// Get queue fill ratio
    pub fn queue_fill_ratio(&self) -> f64 {
        let pending = self.pending_count.load(Ordering::Relaxed);
        pending as f64 / self.queue_capacity as f64
    }
    
    /// Enqueue a write request (async)
    pub fn enqueue(&self, mut request: VectorWriteRequest) -> Result<u64, VectorShardError> {
        let sequence = self.sequence.fetch_add(1, Ordering::SeqCst);
        request.sequence = sequence;
        
        let mut queue = self.write_queue.lock();
        
        if queue.len() >= self.queue_capacity {
            return Err(VectorShardError::QueueFull);
        }
        
        queue.push_back(request);
        self.pending_count.fetch_add(1, Ordering::Relaxed);
        self.metrics.record_request();
        
        Ok(sequence)
    }
    
    /// Enqueue and wait for completion
    pub async fn enqueue_sync(&self, request: VectorWriteRequest) -> Result<VectorWriteResult, VectorShardError> {
        let (tx, rx) = oneshot::channel();
        let request = request.with_completion(tx);
        
        self.enqueue(request)?;
        
        // Process immediately
        self.process_batch();
        
        rx.await.map_err(|_| VectorShardError::WriteTimeout)
    }
    
    /// Process pending writes in batch
    pub fn process_batch(&self) -> usize {
        let start = Instant::now();
        
        // Drain queue up to batch size
        let batch: Vec<VectorWriteRequest> = {
            let mut queue = self.write_queue.lock();
            let batch_size = self.batch_config.max_batch_size.min(queue.len());
            
            if batch_size == 0 {
                return 0;
            }
            
            queue.drain(..batch_size).collect()
        };
        
        let batch_len = batch.len();
        self.pending_count.fetch_sub(batch_len, Ordering::Relaxed);
        
        // Separate inserts for batch processing
        let mut inserts: Vec<(String, Vec<f32>, Option<oneshot::Sender<VectorWriteResult>>)> = Vec::new();
        let mut deletes: Vec<(String, Option<oneshot::Sender<VectorWriteResult>>)> = Vec::new();
        
        for request in batch {
            match request.op {
                VectorWriteOp::Insert { id, vector } => {
                    inserts.push((id, vector, request.completion));
                }
                VectorWriteOp::Delete { id } => {
                    deletes.push((id, request.completion));
                }
                VectorWriteOp::BatchInsert { items } => {
                    for (id, vector) in items {
                        inserts.push((id, vector, None));
                    }
                    // Notify completion for batch
                    if let Some(tx) = request.completion {
                        let _ = tx.send(VectorWriteResult::success(0, self.shard_id, 0));
                    }
                }
            }
        }
        
        // Process inserts in batch
        if !inserts.is_empty() {
            let items: Vec<(String, Vec<f32>)> = inserts
                .iter()
                .map(|(id, v, _)| (id.clone(), v.clone()))
                .collect();
            
            let result = {
                let index = self.index.write();
                index.insert_batch(items)
            };
            
            match result {
                Ok(internal_ids) => {
                    let latency = start.elapsed().as_micros() as u64;
                    
                    for (i, (_, _, completion)) in inserts.into_iter().enumerate() {
                        if let Some(tx) = completion {
                            let internal_id = internal_ids.get(i).copied().unwrap_or(0);
                            let _ = tx.send(VectorWriteResult::success(
                                internal_id,
                                self.shard_id,
                                latency,
                            ));
                        }
                    }
                    
                    self.metrics.record_batch_flush(internal_ids.len(), latency as usize);
                }
                Err(e) => {
                    let latency = start.elapsed().as_micros() as u64;
                    let error_msg = e.to_string();
                    
                    for (_, _, completion) in inserts {
                        if let Some(tx) = completion {
                            let _ = tx.send(VectorWriteResult::failure(
                                self.shard_id,
                                error_msg.clone(),
                                latency,
                            ));
                        }
                    }
                    
                    self.metrics.record_failure();
                }
            }
        }
        
        // Process deletes one by one
        for (id, completion) in deletes {
            let result = {
                let index = self.index.write();
                index.delete(&id)
            };
            
            let latency = start.elapsed().as_micros() as u64;
            
            if let Some(tx) = completion {
                match result {
                    Ok(_) => {
                        let _ = tx.send(VectorWriteResult::delete_success(self.shard_id, latency));
                    }
                    Err(e) => {
                        let _ = tx.send(VectorWriteResult::failure(
                            self.shard_id,
                            e.to_string(),
                            latency,
                        ));
                    }
                }
            }
        }
        
        batch_len
    }
    
    /// Flush all pending writes
    pub fn flush(&self) -> usize {
        let mut total = 0;
        loop {
            let processed = self.process_batch();
            if processed == 0 {
                break;
            }
            total += processed;
        }
        total
    }
    
    // ========================================================================
    // Direct Read Operations (no queue)
    // ========================================================================
    
    /// Search for k nearest neighbors
    pub fn search(&self, query: &[f32], k: usize) -> Vec<VectorSearchResult> {
        let index = self.index.read();
        match index.search(query, k) {
            Ok(results) => {
                // Convert internal results to our result type with external IDs
                let reverse_map = index.reverse_id_map.read();
                results.into_iter().map(|r| {
                    let external_id = reverse_map.get(r.id)
                        .cloned()
                        .unwrap_or_else(|| format!("unknown_{}", r.id));
                    VectorSearchResult {
                        id: external_id,
                        distance: r.distance,
                        shard_id: self.shard_id,
                    }
                }).collect()
            }
            Err(_) => Vec::new(),
        }
    }
    
    /// Search with custom ef parameter (uses default search for now)
    pub fn search_with_ef(&self, query: &[f32], k: usize, _ef: usize) -> Vec<VectorSearchResult> {
        // Note: Current HNSWIndex doesn't expose ef parameter in search
        // Using standard search as fallback
        self.search(query, k)
    }
    
    /// Get a vector by external ID
    pub fn get(&self, id: &str) -> Option<Arc<Vec<f32>>> {
        let index = self.index.read();
        index.get_arc(id)
    }
    
    /// Check if vector exists
    pub fn contains(&self, id: &str) -> bool {
        let index = self.index.read();
        index.contains(id)
    }
    
    // ========================================================================
    // Direct Write Operations (for sync API)
    // ========================================================================
    
    /// Insert directly (bypasses queue)
    pub fn insert_direct(&self, id: String, vector: Vec<f32>) -> Result<usize, VectorShardError> {
        let index = self.index.write();
        index.insert(id, vector).map_err(|e| e.into())
    }
    
    /// Batch insert directly (bypasses queue)
    pub fn insert_batch_direct(&self, items: Vec<(String, Vec<f32>)>) -> Result<Vec<usize>, VectorShardError> {
        let index = self.index.write();
        index.insert_batch(items).map_err(|e| e.into())
    }
    
    /// Delete directly (bypasses queue)
    pub fn delete_direct(&self, id: &str) -> Result<(), VectorShardError> {
        let index = self.index.write();
        index.delete(id).map_err(|e| e.into())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    fn create_test_shard() -> VectorShard {
        let hnsw_config = HNSWConfig {
            dimension: 4,
            m: 16,
            m_max0: 32,
            ef_construction: 200,
            max_elements: 10000,
            ..Default::default()
        };
        
        let config = ConcurrentHNSWConfig::with_hnsw(hnsw_config.clone());
        let metrics = Arc::new(WriteMetrics::new());
        
        VectorShard::new(0, hnsw_config, &config, metrics)
    }
    
    #[test]
    fn test_direct_insert_and_search() {
        let shard = create_test_shard();
        
        // Insert some vectors
        shard.insert_direct("vec_1".to_string(), vec![1.0, 0.0, 0.0, 0.0]).unwrap();
        shard.insert_direct("vec_2".to_string(), vec![0.0, 1.0, 0.0, 0.0]).unwrap();
        shard.insert_direct("vec_3".to_string(), vec![0.0, 0.0, 1.0, 0.0]).unwrap();
        
        assert_eq!(shard.len(), 3);
        
        // Search
        let query = vec![1.0, 0.1, 0.0, 0.0];
        let results = shard.search(&query, 2);
        
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "vec_1"); // Closest
    }
    
    #[test]
    fn test_queue_insert() {
        let shard = create_test_shard();
        
        // Enqueue inserts
        let req1 = VectorWriteRequest::insert("vec_1".to_string(), vec![1.0, 0.0, 0.0, 0.0], 0);
        let req2 = VectorWriteRequest::insert("vec_2".to_string(), vec![0.0, 1.0, 0.0, 0.0], 0);
        
        shard.enqueue(req1).unwrap();
        shard.enqueue(req2).unwrap();
        
        assert!(shard.queue_fill_ratio() > 0.0);
        
        // Process batch
        let processed = shard.process_batch();
        assert_eq!(processed, 2);
        assert_eq!(shard.len(), 2);
    }
    
    #[test]
    fn test_flush() {
        let shard = create_test_shard();
        
        // Enqueue many inserts
        for i in 0..10 {
            let req = VectorWriteRequest::insert(
                format!("vec_{}", i),
                vec![i as f32, 0.0, 0.0, 0.0],
                0,
            );
            shard.enqueue(req).unwrap();
        }
        
        // Flush all
        let flushed = shard.flush();
        assert_eq!(flushed, 10);
        assert_eq!(shard.len(), 10);
    }
    
    #[test]
    fn test_delete() {
        let shard = create_test_shard();
        
        shard.insert_direct("vec_1".to_string(), vec![1.0, 0.0, 0.0, 0.0]).unwrap();
        assert!(shard.contains("vec_1"));
        
        shard.delete_direct("vec_1").unwrap();
        // Note: HNSW uses soft delete, so contains may still return true
        // but the vector won't appear in search results
    }
}
