//! Global Vector Index for tracking vector locations across shards

use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use parking_lot::RwLock;

/// Global index that tracks which shard each vector is stored in
pub struct GlobalVectorIndex {
    /// ID -> Shard mapping
    id_to_shard: RwLock<HashMap<String, usize>>,
    
    /// Count of vectors per shard
    shard_counts: Vec<AtomicUsize>,
    
    /// Total number of shards
    shard_count: usize,
}

impl GlobalVectorIndex {
    /// Create a new global index
    pub fn new(shard_count: usize) -> Self {
        let shard_counts = (0..shard_count)
            .map(|_| AtomicUsize::new(0))
            .collect();
        
        Self {
            id_to_shard: RwLock::new(HashMap::new()),
            shard_counts,
            shard_count,
        }
    }
    
    /// Register a vector's location
    pub fn register(&self, id: &str, shard_id: usize) {
        let mut map = self.id_to_shard.write();
        let is_new = !map.contains_key(id);
        map.insert(id.to_string(), shard_id);
        
        if is_new && shard_id < self.shard_count {
            self.shard_counts[shard_id].fetch_add(1, Ordering::Relaxed);
        }
    }
    
    /// Register multiple vectors
    pub fn register_batch(&self, items: &[(String, usize)]) {
        let mut map = self.id_to_shard.write();
        for (id, shard_id) in items {
            let is_new = !map.contains_key(id);
            map.insert(id.clone(), *shard_id);
            
            if is_new && *shard_id < self.shard_count {
                self.shard_counts[*shard_id].fetch_add(1, Ordering::Relaxed);
            }
        }
    }
    
    /// Unregister a vector
    pub fn unregister(&self, id: &str) -> Option<usize> {
        let mut map = self.id_to_shard.write();
        if let Some(shard_id) = map.remove(id) {
            if shard_id < self.shard_count {
                self.shard_counts[shard_id].fetch_sub(1, Ordering::Relaxed);
            }
            Some(shard_id)
        } else {
            None
        }
    }
    
    /// Locate which shard a vector is in
    pub fn locate(&self, id: &str) -> Option<usize> {
        let map = self.id_to_shard.read();
        map.get(id).copied()
    }
    
    /// Check if a vector exists
    pub fn contains(&self, id: &str) -> bool {
        let map = self.id_to_shard.read();
        map.contains_key(id)
    }
    
    /// Get the total count of vectors
    pub fn total_count(&self) -> usize {
        self.shard_counts
            .iter()
            .map(|c| c.load(Ordering::Relaxed))
            .sum()
    }
    
    /// Get count per shard
    pub fn shard_counts(&self) -> Vec<usize> {
        self.shard_counts
            .iter()
            .map(|c| c.load(Ordering::Relaxed))
            .collect()
    }
    
    /// Get the least loaded shard (for round-robin or load balancing)
    pub fn least_loaded_shard(&self) -> usize {
        self.shard_counts
            .iter()
            .enumerate()
            .min_by_key(|(_, count)| count.load(Ordering::Relaxed))
            .map(|(id, _)| id)
            .unwrap_or(0)
    }
    
    /// Get shard statistics
    pub fn stats(&self) -> Vec<(usize, usize)> {
        self.shard_counts
            .iter()
            .enumerate()
            .map(|(id, count)| (id, count.load(Ordering::Relaxed)))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_register_and_locate() {
        let index = GlobalVectorIndex::new(4);
        
        index.register("vec_1", 0);
        index.register("vec_2", 1);
        index.register("vec_3", 2);
        
        assert_eq!(index.locate("vec_1"), Some(0));
        assert_eq!(index.locate("vec_2"), Some(1));
        assert_eq!(index.locate("vec_3"), Some(2));
        assert_eq!(index.locate("vec_4"), None);
    }
    
    #[test]
    fn test_unregister() {
        let index = GlobalVectorIndex::new(4);
        
        index.register("vec_1", 0);
        assert!(index.contains("vec_1"));
        assert_eq!(index.total_count(), 1);
        
        let shard = index.unregister("vec_1");
        assert_eq!(shard, Some(0));
        assert!(!index.contains("vec_1"));
        assert_eq!(index.total_count(), 0);
    }
    
    #[test]
    fn test_shard_counts() {
        let index = GlobalVectorIndex::new(4);
        
        for i in 0..10 {
            index.register(&format!("vec_{}", i), i % 4);
        }
        
        let counts = index.shard_counts();
        assert_eq!(counts[0], 3); // 0, 4, 8
        assert_eq!(counts[1], 3); // 1, 5, 9
        assert_eq!(counts[2], 2); // 2, 6
        assert_eq!(counts[3], 2); // 3, 7
        assert_eq!(index.total_count(), 10);
    }
    
    #[test]
    fn test_least_loaded_shard() {
        let index = GlobalVectorIndex::new(4);
        
        // Shard 0: 3, Shard 1: 1, Shard 2: 2, Shard 3: 0
        index.register("a", 0);
        index.register("b", 0);
        index.register("c", 0);
        index.register("d", 1);
        index.register("e", 2);
        index.register("f", 2);
        
        assert_eq!(index.least_loaded_shard(), 3);
    }
    
    #[test]
    fn test_register_batch() {
        let index = GlobalVectorIndex::new(4);
        
        let items = vec![
            ("vec_1".to_string(), 0),
            ("vec_2".to_string(), 1),
            ("vec_3".to_string(), 2),
        ];
        
        index.register_batch(&items);
        
        assert_eq!(index.total_count(), 3);
        assert_eq!(index.locate("vec_1"), Some(0));
        assert_eq!(index.locate("vec_2"), Some(1));
        assert_eq!(index.locate("vec_3"), Some(2));
    }
}
