//! Vector Shard Router - Routes operations to appropriate shards

use std::sync::Arc;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use rayon::prelude::*;

use crate::storage::concurrent::{
    WriteMetrics, MetricsSnapshot, 
    BackpressureController, BackpressureAction,
    ShardStrategy,
};

use super::config::ConcurrentHNSWConfig;
use super::types::{VectorWriteRequest, VectorWriteResult, VectorShardError, VectorSearchResult};
use super::shard::VectorShard;
use super::global_index::GlobalVectorIndex;

/// Vector Shard Router - distributes operations across shards
pub struct VectorShardRouter {
    /// Shards
    shards: Vec<Arc<VectorShard>>,
    
    /// Number of shards
    shard_count: usize,
    
    /// Routing strategy
    strategy: ShardStrategy,
    
    /// Global index
    global_index: Arc<GlobalVectorIndex>,
    
    /// Backpressure controller
    backpressure: BackpressureController,
    
    /// Metrics
    metrics: Arc<WriteMetrics>,
    
    /// HNSW dimension (for validation)
    dimension: usize,
}

impl VectorShardRouter {
    /// Create a new router
    pub fn new(config: ConcurrentHNSWConfig) -> Self {
        let shard_count = config.effective_shard_count();
        let metrics = Arc::new(WriteMetrics::new());
        let global_index = Arc::new(GlobalVectorIndex::new(shard_count));
        
        // Create shards
        let shards: Vec<Arc<VectorShard>> = (0..shard_count)
            .map(|i| {
                let mut hnsw_config = config.hnsw.clone();
                // Divide max_elements across shards
                hnsw_config.max_elements = (config.hnsw.max_elements / shard_count).max(1000);
                
                Arc::new(VectorShard::new(
                    i,
                    hnsw_config,
                    &config,
                    Arc::clone(&metrics),
                ))
            })
            .collect();
        
        let backpressure = BackpressureController::new(config.backpressure_config, Arc::clone(&metrics));
        
        Self {
            shards,
            shard_count,
            strategy: config.shard_strategy,
            global_index,
            backpressure,
            metrics,
            dimension: config.hnsw.dimension,
        }
    }
    
    /// Get shard count
    pub fn shard_count(&self) -> usize {
        self.shard_count
    }
    
    /// Route an ID to a shard
    fn route(&self, id: &str) -> usize {
        match self.strategy {
            ShardStrategy::Hash => {
                let mut hasher = DefaultHasher::new();
                id.hash(&mut hasher);
                (hasher.finish() as usize) % self.shard_count
            }
            ShardStrategy::RoundRobin => {
                self.global_index.least_loaded_shard()
            }
            ShardStrategy::Collection => {
                // For vector index, treat Collection same as Hash
                let mut hasher = DefaultHasher::new();
                id.hash(&mut hasher);
                (hasher.finish() as usize) % self.shard_count
            }
        }
    }
    
    /// Get shard by ID (or route if not found)
    fn get_shard(&self, id: &str) -> &Arc<VectorShard> {
        let shard_id = self.global_index
            .locate(id)
            .unwrap_or_else(|| self.route(id));
        &self.shards[shard_id]
    }
    
    // ========================================================================
    // Write Operations
    // ========================================================================
    
    /// Insert a vector (async, queued)
    pub async fn insert(&self, id: String, vector: Vec<f32>) -> Result<VectorWriteResult, VectorShardError> {
        // Validate dimension
        if vector.len() != self.dimension {
            return Err(VectorShardError::DimensionMismatch {
                expected: self.dimension,
                got: vector.len(),
            });
        }
        
        // Check backpressure
        let fill_ratio = self.queue_fill_ratio();
        match self.backpressure.evaluate(fill_ratio) {
            BackpressureAction::Accept => {}
            BackpressureAction::Slowdown(duration) => {
                tokio::time::sleep(duration).await;
            }
            BackpressureAction::Reject => {
                return Err(VectorShardError::QueueFull);
            }
        }
        
        // Route to shard
        let shard_id = self.route(&id);
        let shard = &self.shards[shard_id];
        
        // Enqueue
        let request = VectorWriteRequest::insert(id.clone(), vector, 0);
        shard.enqueue(request)?;
        
        // Register in global index
        self.global_index.register(&id, shard_id);
        
        // Trigger batch processing (non-blocking)
        let shard_clone = Arc::clone(shard);
        tokio::spawn(async move {
            shard_clone.process_batch();
        });
        
        Ok(VectorWriteResult::success(0, shard_id, 0))
    }
    
    /// Insert and wait for completion
    pub async fn insert_sync(&self, id: String, vector: Vec<f32>) -> Result<VectorWriteResult, VectorShardError> {
        // Validate dimension
        if vector.len() != self.dimension {
            return Err(VectorShardError::DimensionMismatch {
                expected: self.dimension,
                got: vector.len(),
            });
        }
        
        // Route to shard
        let shard_id = self.route(&id);
        let shard = &self.shards[shard_id];
        
        // Direct insert (bypasses queue for sync)
        let start = std::time::Instant::now();
        let internal_id = shard.insert_direct(id.clone(), vector)?;
        let latency = start.elapsed().as_micros() as u64;
        
        // Register in global index
        self.global_index.register(&id, shard_id);
        
        self.metrics.record_write_time(latency);
        self.metrics.record_success();
        
        Ok(VectorWriteResult::success(internal_id, shard_id, latency))
    }
    
    /// Batch insert
    pub async fn insert_batch(&self, items: Vec<(String, Vec<f32>)>) -> Result<Vec<VectorWriteResult>, VectorShardError> {
        // Validate dimensions
        for (id, vector) in &items {
            if vector.len() != self.dimension {
                return Err(VectorShardError::DimensionMismatch {
                    expected: self.dimension,
                    got: vector.len(),
                });
            }
            
            // Check for duplicates
            if self.global_index.contains(id) {
                return Err(VectorShardError::DuplicateId(id.clone()));
            }
        }
        
        // Group items by shard
        let mut shard_items: Vec<Vec<(String, Vec<f32>)>> = vec![Vec::new(); self.shard_count];
        
        for (id, vector) in items {
            let shard_id = self.route(&id);
            shard_items[shard_id].push((id, vector));
        }
        
        // Insert into each shard
        let mut results = Vec::new();
        
        for (shard_id, items) in shard_items.into_iter().enumerate() {
            if items.is_empty() {
                continue;
            }
            
            let shard = &self.shards[shard_id];
            let start = std::time::Instant::now();
            
            // Register all IDs first
            let ids: Vec<(String, usize)> = items
                .iter()
                .map(|(id, _)| (id.clone(), shard_id))
                .collect();
            self.global_index.register_batch(&ids);
            
            // Direct batch insert via public API
            match shard.insert_batch_direct(items.clone()) {
                Ok(internal_ids) => {
                    let latency = start.elapsed().as_micros() as u64;
                    for internal_id in internal_ids {
                        results.push(VectorWriteResult::success(internal_id, shard_id, latency));
                    }
                }
                Err(e) => {
                    let latency = start.elapsed().as_micros() as u64;
                    for _ in 0..items.len() {
                        results.push(VectorWriteResult::failure(shard_id, e.to_string(), latency));
                    }
                }
            }
        }
        
        Ok(results)
    }
    
    /// Delete a vector
    pub async fn delete(&self, id: &str) -> Result<VectorWriteResult, VectorShardError> {
        // Find shard
        let shard_id = self.global_index
            .locate(id)
            .ok_or_else(|| VectorShardError::NotFound(id.to_string()))?;
        
        let shard = &self.shards[shard_id];
        let start = std::time::Instant::now();
        
        shard.delete_direct(id)?;
        
        // Unregister from global index
        self.global_index.unregister(id);
        
        let latency = start.elapsed().as_micros() as u64;
        Ok(VectorWriteResult::delete_success(shard_id, latency))
    }
    
    // ========================================================================
    // Read Operations
    // ========================================================================
    
    /// Search across all shards
    pub fn search(&self, query: &[f32], k: usize) -> Vec<VectorSearchResult> {
        // Parallel search across all shards
        let mut all_results: Vec<VectorSearchResult> = self.shards
            .par_iter()
            .flat_map(|shard| shard.search(query, k))
            .collect();
        
        // Sort by distance and take top-k
        all_results.sort_by(|a, b| {
            a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal)
        });
        all_results.truncate(k);
        
        all_results
    }
    
    /// Search with custom ef
    pub fn search_with_ef(&self, query: &[f32], k: usize, ef: usize) -> Vec<VectorSearchResult> {
        let mut all_results: Vec<VectorSearchResult> = self.shards
            .par_iter()
            .flat_map(|shard| shard.search_with_ef(query, k, ef))
            .collect();
        
        all_results.sort_by(|a, b| {
            a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal)
        });
        all_results.truncate(k);
        
        all_results
    }
    
    /// Get a vector by ID
    pub fn get(&self, id: &str) -> Option<Arc<Vec<f32>>> {
        let shard_id = self.global_index.locate(id)?;
        self.shards[shard_id].get(id)
    }
    
    /// Check if vector exists
    pub fn contains(&self, id: &str) -> bool {
        self.global_index.contains(id)
    }
    
    /// Get total vector count
    pub fn len(&self) -> usize {
        self.global_index.total_count()
    }
    
    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    
    // ========================================================================
    // Management
    // ========================================================================
    
    /// Get metrics snapshot
    pub fn metrics(&self) -> MetricsSnapshot {
        self.metrics.snapshot()
    }
    
    /// Get queue fill ratio (max across all shards)
    pub fn queue_fill_ratio(&self) -> f64 {
        self.shards
            .iter()
            .map(|s| s.queue_fill_ratio())
            .fold(0.0, f64::max)
    }
    
    /// Flush all pending writes
    pub async fn flush(&self) -> Result<(), VectorShardError> {
        for shard in &self.shards {
            shard.flush();
        }
        Ok(())
    }
    
    /// Check backpressure status
    pub fn backpressure_status(&self) -> BackpressureAction {
        self.backpressure.evaluate(self.queue_fill_ratio())
    }
    
    /// Get shard statistics
    pub fn shard_stats(&self) -> Vec<(usize, usize)> {
        self.global_index.stats()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::HNSWConfig;
    
    fn create_test_router(shard_count: usize) -> VectorShardRouter {
        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)
            .with_shards(shard_count);
        
        VectorShardRouter::new(config)
    }
    
    #[test]
    fn test_routing_consistency() {
        let router = create_test_router(4);
        
        // Same ID should always route to same shard
        let shard1 = router.route("test_id");
        let shard2 = router.route("test_id");
        assert_eq!(shard1, shard2);
    }
    
    #[test]
    fn test_routing_distribution() {
        let router = create_test_router(4);
        
        let mut counts = [0usize; 4];
        for i in 0..1000 {
            let shard = router.route(&format!("id_{}", i));
            counts[shard] += 1;
        }
        
        // Each shard should have roughly 250 items (allow 50% variance)
        for count in &counts {
            assert!(*count > 100, "Shard count too low: {}", count);
            assert!(*count < 400, "Shard count too high: {}", count);
        }
    }
    
    #[tokio::test]
    async fn test_insert_sync() {
        let router = create_test_router(4);
        
        let result = router.insert_sync(
            "vec_1".to_string(),
            vec![1.0, 0.0, 0.0, 0.0],
        ).await.unwrap();
        
        assert!(result.success);
        assert!(router.contains("vec_1"));
        assert_eq!(router.len(), 1);
    }
    
    #[tokio::test]
    async fn test_insert_batch() {
        let router = create_test_router(4);
        
        let items: Vec<(String, Vec<f32>)> = (0..100)
            .map(|i| (format!("vec_{}", i), vec![i as f32, 0.0, 0.0, 0.0]))
            .collect();
        
        let results = router.insert_batch(items).await.unwrap();
        
        assert_eq!(results.len(), 100);
        assert!(results.iter().all(|r| r.success));
        assert_eq!(router.len(), 100);
        
        // Check distribution across shards
        let stats = router.shard_stats();
        let total: usize = stats.iter().map(|(_, c)| c).sum();
        assert_eq!(total, 100);
    }
    
    #[tokio::test]
    async fn test_search_across_shards() {
        let router = create_test_router(4);
        
        // Insert vectors that will be distributed across shards
        let items: Vec<(String, Vec<f32>)> = (0..100)
            .map(|i| (format!("vec_{}", i), vec![i as f32, 0.0, 0.0, 0.0]))
            .collect();
        
        router.insert_batch(items).await.unwrap();
        
        // Search should find results from all shards
        let query = vec![50.0, 0.0, 0.0, 0.0];
        let results = router.search(&query, 10);
        
        assert_eq!(results.len(), 10);
        // Results should be sorted by distance
        for i in 1..results.len() {
            assert!(results[i].distance >= results[i-1].distance);
        }
    }
    
    #[tokio::test]
    async fn test_delete() {
        let router = create_test_router(4);
        
        router.insert_sync("vec_1".to_string(), vec![1.0, 0.0, 0.0, 0.0]).await.unwrap();
        assert!(router.contains("vec_1"));
        
        router.delete("vec_1").await.unwrap();
        assert!(!router.contains("vec_1"));
    }
    
    #[tokio::test]
    async fn test_dimension_validation() {
        let router = create_test_router(4);
        
        // Wrong dimension should fail
        let result = router.insert_sync(
            "vec_1".to_string(),
            vec![1.0, 0.0, 0.0], // 3D instead of 4D
        ).await;
        
        assert!(matches!(result, Err(VectorShardError::DimensionMismatch { .. })));
    }
}
