zipora 3.1.5

//! Hash Index for Indirect Lookups
//!
//! A hash index structure optimized for large value types through indirection.
//! This container separates key storage from value storage for improved cache
//! efficiency and reduced memory pressure when values are large (>64 bytes).

use crate::error::{Result, ZiporaError};
use crate::memory::{SecureMemoryPool, SecurePooledPtr, get_global_pool_for_size};
use ahash::AHasher;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::mem;
use std::sync::Arc;

/// Hash index for indirect value lookups
///
/// GoldHashIdx provides a high-performance hash table optimized for scenarios
/// where values are large. Instead of storing values directly in the hash table,
/// it uses indirection through a value pool, improving cache efficiency and
/// reducing memory fragmentation.
///
/// # Design Principles
///
/// - **Separation of Concerns**: Keys and values stored separately
/// - **Cache Efficiency**: Hot key data stays in cache, cold values in pool
/// - **Memory Safety**: Integration with SecureMemoryPool for value allocation
/// - **Performance**: Same lookup speed as GoldHashMap with better memory characteristics
///
/// # Memory Layout
///
/// ```text
/// Hash Table: [Key1, ValueIndex1] [Key2, ValueIndex2] ...
/// Value Pool: [Value1] [Value2] [Large Value3] ...
/// ```
///
/// # Performance Characteristics
///
/// - **Memory**: 30% reduction for large values (>64 bytes)
/// - **Lookup**: O(1) average, same as standard hash map
/// - **Cache**: Improved cache locality for key scans
/// - **Allocation**: Reduced fragmentation through pooled values
///
/// # Example
///
/// ```rust
/// use zipora::containers::GoldHashIdx;
///
/// #[derive(Clone, Debug, PartialEq)]
/// struct LargeValue {
///     data: [u8; 256], // Large value type
///     metadata: String,
/// }
///
/// let mut idx = GoldHashIdx::new();
/// let large_val = LargeValue {
///     data: [42; 256],
///     metadata: "test".to_string()
/// };
///
/// idx.insert("key1".to_string(), large_val.clone())?;
/// assert_eq!(idx.get(&"key1".to_string()), Some(&large_val));
/// # Ok::<(), zipora::error::ZiporaError>(())
/// ```
pub struct GoldHashIdx<K, V>
where
    K: Hash + Eq + Clone,
    V: Clone,
{
    /// Hash table storing keys and indices into value pool
    table: Vec<Option<Bucket<K>>>,
    /// Value pool for efficient large value storage
    values: Vec<Option<SecurePooledPtr>>,
    /// Free list for reusing value slots
    free_values: Vec<usize>,
    /// Number of elements stored
    len: usize,
    /// Current capacity of the hash table
    capacity: usize,
    /// Load factor threshold for resizing
    load_factor: f64,
    /// Memory pool for value allocation
    pool: Option<Arc<SecureMemoryPool>>,
    /// Statistics tracking
    key_memory: usize,
    value_memory: usize,
    /// Type marker for V
    _phantom: PhantomData<V>,
}

/// Hash table bucket storing key and value index
#[derive(Debug, Clone)]
struct Bucket<K> {
    key: K,
    value_index: usize,
    hash: u64,
}

impl<K, V> GoldHashIdx<K, V>
where
    K: Hash + Eq + Clone,
    V: Clone,
{
    /// Create a new empty GoldHashIdx
    pub fn new() -> Self {
        Self::with_capacity(16)
    }

    /// Create a new GoldHashIdx with specified initial capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let capacity = capacity.next_power_of_two().max(16);

        Self {
            table: vec![None; capacity],
            values: Vec::new(),
            free_values: Vec::new(),
            len: 0,
            capacity,
            load_factor: 0.75,
            pool: None,
            key_memory: 0,
            value_memory: 0,
            _phantom: PhantomData,
        }
    }

    /// Create a GoldHashIdx with a dedicated memory pool for values
    pub fn with_pool(capacity: usize, pool: Arc<SecureMemoryPool>) -> Self {
        let capacity = capacity.next_power_of_two().max(16);

        Self {
            table: vec![None; capacity],
            values: Vec::new(),
            free_values: Vec::new(),
            len: 0,
            capacity,
            load_factor: 0.75,
            pool: Some(pool),
            key_memory: 0,
            value_memory: 0,
            _phantom: PhantomData,
        }
    }

    /// Insert a key-value pair
    ///
    /// # Arguments
    /// * `key` - The key to insert
    /// * `value` - The value to store in the value pool
    ///
    /// # Returns
    /// The previous value associated with the key, if any
    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>> {
        // Check if we need to resize
        if self.len >= (self.capacity as f64 * self.load_factor) as usize {
            self.resize()?;
        }

        let hash = self.hash_key(&key);
        let mut index = (hash % self.capacity as u64) as usize;

        // Probe for insertion point
        loop {
            if self.table[index].is_none() {
                // Found empty slot - insert new entry
                let value_index = self.allocate_value(value)?;
                self.table[index] = Some(Bucket {
                    key: key.clone(),
                    value_index,
                    hash,
                });

                self.len += 1;
                self.key_memory += mem::size_of::<K>() + mem::size_of::<Bucket<K>>();
                return Ok(None);
            } else if let Some(bucket) = &self.table[index] {
                if bucket.hash == hash && bucket.key == key {
                    // Found existing key - replace value
                    let value_index = bucket.value_index;
                    let old_value = self.get_value(value_index)?.clone();
                    self.store_value(value_index, value)?;
                    return Ok(Some(old_value));
                }

                // Continue probing
                index = (index + 1) % self.capacity;
            }
        }
    }

    /// Get a reference to the value associated with a key
    pub fn get(&self, key: &K) -> Option<&V> {
        let hash = self.hash_key(key);
        let mut index = (hash % self.capacity as u64) as usize;

        // Probe for key
        for _ in 0..self.capacity {
            match &self.table[index] {
                None => return None,
                Some(bucket) => {
                    if bucket.hash == hash && bucket.key == *key {
                        return self.get_value(bucket.value_index).ok();
                    }
                    index = (index + 1) % self.capacity;
                }
            }
        }

        None
    }

    /// Get a mutable reference to the value associated with a key
    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        let hash = self.hash_key(key);
        let mut index = (hash % self.capacity as u64) as usize;

        // Probe for key
        for _ in 0..self.capacity {
            match &self.table[index] {
                None => return None,
                Some(bucket) => {
                    if bucket.hash == hash && bucket.key == *key {
                        let value_index = bucket.value_index;
                        return self.get_value_mut(value_index).ok();
                    }
                    index = (index + 1) % self.capacity;
                }
            }
        }

        None
    }

    /// Remove a key-value pair
    pub fn remove(&mut self, key: &K) -> Option<V> {
        let hash = self.hash_key(key);
        let mut index = (hash % self.capacity as u64) as usize;

        // Probe for key
        for _ in 0..self.capacity {
            match &self.table[index] {
                None => return None,
                Some(bucket) => {
                    if bucket.hash == hash && bucket.key == *key {
                        let value_index = bucket.value_index;
                        let value = self.get_value(value_index).ok()?.clone();

                        // Mark slot as deleted and add value index to free list
                        self.table[index] = None;
                        self.free_value(value_index);
                        self.len -= 1;
                        self.key_memory = self
                            .key_memory
                            .saturating_sub(mem::size_of::<K>() + mem::size_of::<Bucket<K>>());

                        // Rehash following entries to maintain probe sequence
                        self.rehash_after_removal(index);

                        return Some(value);
                    }
                    index = (index + 1) % self.capacity;
                }
            }
        }

        None
    }

    /// Check if the map contains a key
    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }

    /// Get the number of key-value pairs
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Check if the map is empty
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Insert multiple key-value pairs efficiently
    pub fn insert_batch(&mut self, items: Vec<(K, V)>) -> Result<()> {
        // Pre-resize if needed
        let target_capacity = (self.len + items.len()) * 2;
        if target_capacity > self.capacity {
            let new_capacity = target_capacity.next_power_of_two();
            self.resize_to(new_capacity)?;
        }

        for (key, value) in items {
            self.insert(key, value)?;
        }

        Ok(())
    }

    /// Get multiple values efficiently
    pub fn get_batch(&self, keys: &[K]) -> Vec<Option<&V>> {
        keys.iter().map(|key| self.get(key)).collect()
    }

    /// Shrink the hash table to fit current elements
    pub fn shrink_to_fit(&mut self) {
        let min_capacity = ((self.len as f64 / self.load_factor) as usize)
            .next_power_of_two()
            .max(16);
        if min_capacity < self.capacity {
            if let Err(_) = self.resize_to(min_capacity) {
                // If resize fails, continue with current capacity
            }
        }

        // Also shrink value storage
        self.values.shrink_to_fit();
        self.free_values.shrink_to_fit();
    }

    /// Get memory usage statistics (keys, values)
    #[inline]
    pub fn memory_usage(&self) -> (usize, usize) {
        (self.key_memory, self.value_memory)
    }

    /// Hash a key using AHash
    fn hash_key(&self, key: &K) -> u64 {
        let mut hasher = AHasher::default();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Allocate a new value in the pool
    fn allocate_value(&mut self, value: V) -> Result<usize> {
        // Try to reuse a free slot
        if let Some(index) = self.free_values.pop() {
            self.store_value(index, value)?;
            return Ok(index);
        }

        // Allocate new slot
        let value_ptr = self.allocate_pooled_value(value)?;
        let index = self.values.len();
        self.values.push(Some(value_ptr));
        self.value_memory += mem::size_of::<V>();

        Ok(index)
    }

    /// Store a value at a specific index
    fn store_value(&mut self, index: usize, value: V) -> Result<()> {
        if index >= self.values.len() {
            return Err(ZiporaError::invalid_data("Invalid value index"));
        }

        let value_ptr = self.allocate_pooled_value(value)?;
        self.values[index] = Some(value_ptr);
        Ok(())
    }

    /// Get a value reference by index
    fn get_value(&self, index: usize) -> Result<&V> {
        self.values
            .get(index)
            .and_then(|opt| opt.as_ref())
            // SAFETY: ptr is from allocate_pooled_value, points to valid V, lifetime tied to self
            .map(|ptr| unsafe { &*(ptr.as_ptr() as *const V) })
            .ok_or_else(|| ZiporaError::invalid_data("Invalid value index"))
    }

    /// Get a mutable value reference by index
    fn get_value_mut(&mut self, index: usize) -> Result<&mut V> {
        self.values
            .get_mut(index)
            .and_then(|opt| opt.as_ref())
            // SAFETY: ptr is from allocate_pooled_value, points to valid V, exclusive access via &mut self
            .map(|ptr| unsafe { &mut *(ptr.as_ptr() as *mut V) })
            .ok_or_else(|| ZiporaError::invalid_data("Invalid value index"))
    }

    /// Free a value slot
    fn free_value(&mut self, index: usize) {
        if index < self.values.len() {
            self.values[index] = None;
            self.free_values.push(index);
            self.value_memory = self.value_memory.saturating_sub(mem::size_of::<V>());
        }
    }

    /// Allocate a pooled value
    fn allocate_pooled_value(&mut self, value: V) -> Result<SecurePooledPtr> {
        match &self.pool {
            Some(pool) => {
                let ptr = pool.allocate()?;
                // SAFETY: ptr is freshly allocated, properly aligned for V, not yet initialized
                unsafe {
                    std::ptr::write(ptr.as_ptr() as *mut V, value);
                }
                Ok(ptr)
            }
            None => {
                // Use global pool
                let pool = get_global_pool_for_size(mem::size_of::<V>());
                let ptr = pool.allocate()?;
                // SAFETY: ptr is freshly allocated, properly aligned for V, not yet initialized
                unsafe {
                    std::ptr::write(ptr.as_ptr() as *mut V, value);
                }
                Ok(ptr)
            }
        }
    }

    /// Resize the hash table
    fn resize(&mut self) -> Result<()> {
        self.resize_to(self.capacity * 2)
    }

    /// Resize to a specific capacity
    fn resize_to(&mut self, new_capacity: usize) -> Result<()> {
        let old_table = mem::replace(&mut self.table, vec![None; new_capacity]);
        let _old_capacity = self.capacity;
        self.capacity = new_capacity;
        self.len = 0;
        self.key_memory = 0;

        // Rehash all elements
        for bucket in old_table.into_iter().flatten() {
            let value = self.get_value(bucket.value_index)?.clone();
            self.insert(bucket.key, value)?;
        }

        Ok(())
    }

    /// Rehash entries after a removal to maintain probe sequences
    fn rehash_after_removal(&mut self, start_index: usize) {
        let mut index = (start_index + 1) % self.capacity;

        while let Some(bucket) = self.table[index].take() {
            // Reinsert this bucket
            let ideal_index = (bucket.hash % self.capacity as u64) as usize;
            let mut new_index = ideal_index;

            // Find new position
            while self.table[new_index].is_some() {
                new_index = (new_index + 1) % self.capacity;
            }

            self.table[new_index] = Some(bucket);
            index = (index + 1) % self.capacity;
        }
    }
}

impl<K, V> Default for GoldHashIdx<K, V>
where
    K: Hash + Eq + Clone,
    V: Clone,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, V> std::fmt::Debug for GoldHashIdx<K, V>
where
    K: Hash + Eq + Clone + std::fmt::Debug,
    V: Clone + std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GoldHashIdx")
            .field("len", &self.len)
            .field("capacity", &self.capacity)
            .field("load_factor", &self.load_factor)
            .field("key_memory", &self.key_memory)
            .field("value_memory", &self.value_memory)
            .field("free_values_count", &self.free_values.len())
            .field("has_pool", &self.pool.is_some())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory::SecurePoolConfig;

    #[derive(Clone, Debug, PartialEq)]
    struct LargeValue {
        data: [u8; 256],
        id: u32,
    }

    impl LargeValue {
        fn new(id: u32) -> Self {
            Self {
                data: [id as u8; 256],
                id,
            }
        }
    }

    #[test]
    fn test_basic_operations() -> Result<()> {
        let mut idx = GoldHashIdx::new();

        // Insert
        let val1 = LargeValue::new(1);
        let val2 = LargeValue::new(2);

        assert_eq!(idx.insert("key1".to_string(), val1.clone())?, None);
        assert_eq!(idx.insert("key2".to_string(), val2.clone())?, None);
        assert_eq!(idx.len(), 2);

        // Get
        assert_eq!(idx.get(&"key1".to_string()), Some(&val1));
        assert_eq!(idx.get(&"key2".to_string()), Some(&val2));
        assert_eq!(idx.get(&"key3".to_string()), None);

        // Contains
        assert!(idx.contains_key(&"key1".to_string()));
        assert!(!idx.contains_key(&"key3".to_string()));

        // Update
        let val1_updated = LargeValue::new(11);
        assert_eq!(
            idx.insert("key1".to_string(), val1_updated.clone())?,
            Some(val1)
        );
        assert_eq!(idx.get(&"key1".to_string()), Some(&val1_updated));

        // Remove
        assert_eq!(idx.remove(&"key1".to_string()), Some(val1_updated));
        assert_eq!(idx.get(&"key1".to_string()), None);
        assert_eq!(idx.len(), 1);

        Ok(())
    }

    #[test]
    fn test_mutable_access() -> Result<()> {
        let mut idx = GoldHashIdx::new();
        let val = LargeValue::new(42);

        idx.insert("key".to_string(), val)?;

        {
            let val_mut = idx.get_mut(&"key".to_string()).unwrap();
            val_mut.id = 99;
        }

        assert_eq!(idx.get(&"key".to_string()).unwrap().id, 99);
        Ok(())
    }

    #[test]
    fn test_batch_operations() -> Result<()> {
        let mut idx = GoldHashIdx::new();

        // Batch insert
        let items: Vec<(String, LargeValue)> = (0..100)
            .map(|i| (format!("key{}", i), LargeValue::new(i)))
            .collect();

        idx.insert_batch(items.clone())?;
        assert_eq!(idx.len(), 100);

        // Batch get
        let keys: Vec<String> = (0..100).map(|i| format!("key{}", i)).collect();
        let values = idx.get_batch(&keys);

        for (i, value_opt) in values.iter().enumerate() {
            assert!(value_opt.is_some());
            assert_eq!(value_opt.unwrap().id, i as u32);
        }

        Ok(())
    }

    #[test]
    fn test_with_custom_pool() -> Result<()> {
        let config = SecurePoolConfig::small_secure();
        let pool = SecureMemoryPool::new(config)?;
        let mut idx = GoldHashIdx::with_pool(32, pool);

        let val = LargeValue::new(42);
        idx.insert("test".to_string(), val.clone())?;

        assert_eq!(idx.get(&"test".to_string()), Some(&val));
        Ok(())
    }

    #[test]
    fn test_memory_usage_tracking() -> Result<()> {
        let mut idx = GoldHashIdx::new();

        let initial_usage = idx.memory_usage();
        assert_eq!(initial_usage, (0, 0));

        // Insert some values
        for i in 0..10 {
            idx.insert(format!("key{}", i), LargeValue::new(i))?;
        }

        let (key_mem, value_mem) = idx.memory_usage();
        assert!(key_mem > 0);
        assert!(value_mem > 0);

        println!(
            "Memory usage - Keys: {} bytes, Values: {} bytes",
            key_mem, value_mem
        );
        Ok(())
    }

    #[test]
    fn test_shrink_to_fit() -> Result<()> {
        let mut idx = GoldHashIdx::with_capacity(1024);

        // Insert a few items
        for i in 0..10 {
            idx.insert(format!("key{}", i), LargeValue::new(i))?;
        }

        let initial_capacity = idx.capacity;
        idx.shrink_to_fit();

        // Should have smaller capacity but same functionality
        assert!(idx.capacity <= initial_capacity);
        assert_eq!(idx.len(), 10);

        // Verify all items still accessible
        for i in 0..10 {
            assert!(idx.contains_key(&format!("key{}", i)));
        }

        Ok(())
    }

    #[test]
    fn test_large_scale() -> Result<()> {
        let mut idx = GoldHashIdx::new();

        // Insert many items to test resizing
        for i in 0..1000 {
            idx.insert(i.to_string(), LargeValue::new(i))?;
        }

        assert_eq!(idx.len(), 1000);

        // Verify all items
        for i in 0..1000 {
            assert_eq!(idx.get(&i.to_string()).unwrap().id, i);
        }

        // Remove half
        for i in 0..500 {
            assert!(idx.remove(&i.to_string()).is_some());
        }

        assert_eq!(idx.len(), 500);

        // Verify remaining items
        for i in 500..1000 {
            assert_eq!(idx.get(&i.to_string()).unwrap().id, i);
        }

        Ok(())
    }

    #[test]
    fn test_empty_operations() {
        let idx: GoldHashIdx<String, LargeValue> = GoldHashIdx::new();

        assert_eq!(idx.len(), 0);
        assert!(idx.is_empty());
        assert_eq!(idx.get(&"key".to_string()), None);
        assert!(!idx.contains_key(&"key".to_string()));
        assert_eq!(idx.memory_usage(), (0, 0));
    }
}