//! Storage traits for slab-like containers with stable keys.
//!
//! Storage provides insert/remove/get operations where keys remain
//! valid until explicitly removed. This enables node-based data structures
//! (lists, skip lists, heaps) to use keys instead of pointers.
//!
//! # Trait Hierarchy
//!
//! ```text
//! Storage<T>           - base trait: get, remove, len
//!     │
//!     ├── BoundedStorage<T>   - fixed capacity, try_insert -> Result
//!     │
//!     └── UnboundedStorage<T> - growable, insert -> Key (infallible)
//! ```
//!
//! # Key Generation Models
//!
//! **Storage-assigned keys** (slab-like): Storage generates keys on insert.
//! - `BoxedStorage`, `slab::Slab`, `nexus_slab::Slab`
//!
//! **Value-provided keys** (map-like): Value contains its own key via [`Keyed`] trait.
//! - `HashMap<K, V>` where `V: Keyed<Key = K>`
//!
//! # Implementations
//!
//! | Type | Traits | Key Model |
//! |------|--------|-----------|
//! | `BoxedStorage<T>` | `BoundedStorage` | Storage-assigned |
//! | `slab::Slab<T>` | `UnboundedStorage` | Storage-assigned |
//! | `nexus_slab::Slab<T>` | `BoundedStorage` | Storage-assigned |
//! | `HashMap<K, V>` | `UnboundedStorage` | Value-provided (`V: Keyed`) |

use crate::Key;

use core::mem::MaybeUninit;
use core::ptr::NonNull;
use std::alloc::{Layout, alloc, dealloc, handle_alloc_error};
use std::collections::HashMap;
use std::hash::{BuildHasher, Hash};
use std::marker::PhantomData;

// =============================================================================
// Traits
// =============================================================================

/// Base storage trait with stable keys.
///
/// Provides common operations for slab-like containers. Keys remain
/// valid until explicitly removed, enabling node-based data structures
/// to use keys instead of pointers.
///
/// See [`BoundedStorage`] for fixed-capacity storage and
/// [`UnboundedStorage`] for growable storage.
pub trait Storage<T> {
    /// Key type for this storage.
    type Key: Key;

    /// Removes and returns the value at `key`, if present.
    fn remove(&mut self, key: Self::Key) -> Option<T>;

    /// Returns a reference to the value at `key`, if present.
    fn get(&self, key: Self::Key) -> Option<&T>;

    /// Returns a mutable reference to the value at `key`, if present.
    fn get_mut(&mut self, key: Self::Key) -> Option<&mut T>;

    /// Returns the number of occupied slots.
    fn len(&self) -> usize;

    /// Returns `true` if no slots are occupied.
    #[inline]
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns a reference without bounds checking.
    ///
    /// # Safety
    ///
    /// `key` must be valid and occupied.
    unsafe fn get_unchecked(&self, key: Self::Key) -> &T;

    /// Returns a mutable reference without bounds checking.
    ///
    /// # Safety
    ///
    /// `key` must be valid and occupied.
    unsafe fn get_unchecked_mut(&mut self, key: Self::Key) -> &mut T;

    /// Removes an element without bounds checking.
    ///
    /// # Safety
    ///
    /// The key must be valid and occupied.
    unsafe fn remove_unchecked(&mut self, key: Self::Key) -> T;
}

/// Fixed-capacity storage where insertion can fail.
///
/// Use this for pre-allocated, bounded storage where capacity is known
/// upfront. Insertion returns `Result<Key, Full<T>>`.
///
/// # Implementations
///
/// - [`BoxedStorage<T>`] - runtime capacity, heap allocated
/// - `nexus_slab::Slab<T>` - page-aligned, mlockable (feature `nexus-slab`)
pub trait BoundedStorage<T>: Storage<T> {
    /// Attempts to insert a value, returning its stable key.
    ///
    /// Returns `Err(Full(value))` if storage is at capacity.
    fn try_insert(&mut self, value: T) -> Result<Self::Key, Full<T>>;

    /// Returns the total capacity (number of slots).
    fn capacity(&self) -> usize;

    /// Returns `true` if all slots are occupied.
    #[inline]
    fn is_full(&self) -> bool {
        self.len() >= self.capacity()
    }
}

/// Growable storage where insertion always succeeds.
///
/// Use this when you want storage that grows as needed. Insertion
/// is infallible (may allocate).
///
/// # Implementations
///
/// - `slab::Slab<T>` - growable, heap allocated (feature `slab`)
/// - `HashMap<K, V>` where `V: Keyed<Key = K>` (feature `std`)
pub trait UnboundedStorage<T>: Storage<T> {
    /// Inserts a value, returning its stable key.
    ///
    /// This operation is infallible but may allocate.
    fn insert(&mut self, value: T) -> Self::Key;
}

/// Trait for values that contain their own key.
///
/// Used with `HashMap` storage where the key is extracted from the value
/// rather than generated by the storage.
///
/// # Example
///
/// ```
/// use nexus_collections::Keyed;
///
/// #[derive(Clone, PartialEq, Eq, Hash)]
/// struct OrderId(String);
///
/// struct Order {
///     id: OrderId,
///     price: u64,
///     qty: u64,
/// }
///
/// impl Keyed for Order {
///     type Key = OrderId;
///
///     fn key(&self) -> Self::Key {
///         self.id.clone()
///     }
/// }
/// ```
pub trait Keyed {
    /// The key type for this value.
    type Key;

    /// Returns the key for this value.
    fn key(&self) -> Self::Key;
}

// =============================================================================
// Error Type
// =============================================================================

/// Error returned when fixed-capacity storage is full.
///
/// Contains the value that could not be inserted, allowing recovery.
/// Modeled after `std::sync::mpsc::SendError`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Full<T>(pub T);

impl<T> Full<T> {
    /// Returns the value that could not be inserted.
    pub fn into_inner(self) -> T {
        self.0
    }
}

impl<T> core::fmt::Display for Full<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "storage is full")
    }
}

impl<T: core::fmt::Debug> std::error::Error for Full<T> {}

// =============================================================================
// BoxedStorage - runtime capacity, single allocation, bitmap occupancy
// =============================================================================

/// Fixed-capacity storage with runtime-determined size.
///
/// Uses a single heap allocation containing:
/// - Entry array (`MaybeUninit<T>`)
/// - Occupancy bitmap (`u64` words)
/// - Free stack (keys)
///
/// Capacity is rounded up to the next power of 2 for bitmap efficiency.
///
/// # Example
///
/// ```
/// use nexus_collections::{BoxedStorage, BoundedStorage, Storage};
///
/// let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(1000);
/// assert!(storage.capacity() >= 1000); // Rounded to 1024
///
/// let key = storage.try_insert(42).unwrap();
/// assert_eq!(storage.get(key), Some(&42));
/// ```
pub struct BoxedStorage<T, K: Key = u32> {
    /// Single allocation containing entries, bitmap, and free stack.
    ptr: NonNull<u8>,
    /// Capacity (always power of 2).
    capacity: usize,
    /// Number of free slots.
    free_len: usize,
    /// Cached layout for deallocation.
    layout: Layout,
    /// Offset to bitmap from ptr.
    bitmap_offset: usize,
    /// Offset to free stack from ptr.
    free_stack_offset: usize,
    _marker: PhantomData<(T, K)>,
}

impl<T, K: Key> BoxedStorage<T, K> {
    /// Creates storage with at least `min_capacity` slots.
    ///
    /// Actual capacity is rounded up to the next power of 2.
    ///
    /// # Panics
    ///
    /// Panics if `min_capacity` is 0 or exceeds the key type's maximum.
    pub fn with_capacity(min_capacity: usize) -> Self {
        assert!(min_capacity > 0, "capacity must be > 0");

        // Round up to power of 2 for bitmap efficiency
        let capacity = min_capacity.next_power_of_two();

        assert!(
            capacity <= K::NONE.as_usize(),
            "capacity exceeds key type maximum"
        );

        // Calculate layout
        // Layout: [entries][padding][bitmap][padding][free_stack]
        let entries_layout = Layout::array::<MaybeUninit<T>>(capacity).unwrap();
        let bitmap_words = bitmap_words(capacity);
        let bitmap_layout = Layout::array::<u64>(bitmap_words).unwrap();
        let free_stack_layout = Layout::array::<K>(capacity).unwrap();

        let (layout, bitmap_offset) = entries_layout.extend(bitmap_layout).unwrap();
        let (layout, free_stack_offset) = layout.extend(free_stack_layout).unwrap();
        let layout = layout.pad_to_align();

        // Allocate
        let ptr = unsafe { alloc(layout) };
        if ptr.is_null() {
            handle_alloc_error(layout);
        }
        let ptr = unsafe { NonNull::new_unchecked(ptr) };

        // Initialize bitmap to all zeros (all slots free, but we track via free_stack)
        unsafe {
            let bitmap_ptr = ptr.as_ptr().add(bitmap_offset) as *mut u64;
            core::ptr::write_bytes(bitmap_ptr, 0, bitmap_words);
        }

        // Initialize free stack
        unsafe {
            let free_stack_ptr = ptr.as_ptr().add(free_stack_offset) as *mut K;
            for i in 0..capacity {
                free_stack_ptr.add(i).write(K::from_usize(i));
            }
        }

        Self {
            ptr,
            capacity,
            free_len: capacity,
            layout,
            bitmap_offset,
            free_stack_offset,
            _marker: PhantomData,
        }
    }

    /// Removes all elements from storage.
    ///
    /// This drops all stored values and makes all slots available for reuse.
    ///
    /// # Warning
    ///
    /// If any data structures (List, Heap, etc.) still reference keys in
    /// this storage, they will have dangling references. Only call this when
    /// you know nothing else references the storage, or after clearing those
    /// data structures first.
    ///
    /// For owned wrappers like `OwnedList` and `OwnedHeap`, this is handled
    /// automatically.
    pub fn clear(&mut self) {
        // Drop all occupied values
        for i in 0..self.capacity {
            if self.is_occupied(i) {
                // Safety: slot is occupied
                unsafe {
                    let ptr = self.entries_ptr().add(i);
                    std::ptr::drop_in_place((*ptr).as_mut_ptr());
                }
            }
        }

        // Reset bitmap to all zeros (all vacant)
        unsafe {
            std::ptr::write_bytes(self.bitmap_ptr(), 0, bitmap_words(self.capacity));
        }

        // Rebuild free stack
        let free_stack = self.free_stack_ptr();
        for i in 0..self.capacity {
            unsafe {
                *free_stack.add(i) = K::from_usize(i);
            }
        }
        self.free_len = self.capacity;
    }

    #[inline]
    fn entries_ptr(&self) -> *mut MaybeUninit<T> {
        self.ptr.as_ptr() as *mut MaybeUninit<T>
    }

    #[inline]
    fn bitmap_ptr(&self) -> *mut u64 {
        unsafe { self.ptr.as_ptr().add(self.bitmap_offset) as *mut u64 }
    }

    #[inline]
    fn free_stack_ptr(&self) -> *mut K {
        unsafe { self.ptr.as_ptr().add(self.free_stack_offset) as *mut K }
    }

    #[inline]
    fn is_occupied(&self, idx: usize) -> bool {
        let word = idx / 64;
        let bit = idx % 64;
        unsafe {
            let bitmap = self.bitmap_ptr();
            (*bitmap.add(word) & (1 << bit)) != 0
        }
    }

    #[inline]
    fn set_occupied(&mut self, idx: usize) {
        let word = idx / 64;
        let bit = idx % 64;
        unsafe {
            let bitmap = self.bitmap_ptr();
            *bitmap.add(word) |= 1 << bit;
        }
    }

    #[inline]
    fn set_vacant(&mut self, idx: usize) {
        let word = idx / 64;
        let bit = idx % 64;
        unsafe {
            let bitmap = self.bitmap_ptr();
            *bitmap.add(word) &= !(1 << bit);
        }
    }
}

impl<T, K: Key> Storage<T> for BoxedStorage<T, K> {
    type Key = K;

    #[inline]
    fn remove(&mut self, key: Self::Key) -> Option<T> {
        let i = key.as_usize();
        if i >= self.capacity || !self.is_occupied(i) {
            return None;
        }

        self.set_vacant(i);
        let value = unsafe { self.entries_ptr().add(i).read().assume_init() };

        unsafe {
            self.free_stack_ptr().add(self.free_len).write(key);
        }
        self.free_len += 1;

        Some(value)
    }

    #[inline]
    fn get(&self, key: Self::Key) -> Option<&T> {
        let i = key.as_usize();
        if i >= self.capacity || !self.is_occupied(i) {
            return None;
        }

        Some(unsafe { (*self.entries_ptr().add(i)).assume_init_ref() })
    }

    #[inline]
    fn get_mut(&mut self, key: Self::Key) -> Option<&mut T> {
        let i = key.as_usize();
        if i >= self.capacity || !self.is_occupied(i) {
            return None;
        }

        Some(unsafe { (*self.entries_ptr().add(i)).assume_init_mut() })
    }

    #[inline]
    fn len(&self) -> usize {
        self.capacity - self.free_len
    }

    #[inline]
    unsafe fn get_unchecked(&self, key: Self::Key) -> &T {
        unsafe { (*self.entries_ptr().add(key.as_usize())).assume_init_ref() }
    }

    #[inline]
    unsafe fn get_unchecked_mut(&mut self, key: Self::Key) -> &mut T {
        unsafe { (*self.entries_ptr().add(key.as_usize())).assume_init_mut() }
    }

    #[inline]
    unsafe fn remove_unchecked(&mut self, key: Self::Key) -> T {
        let i = key.as_usize();

        self.set_vacant(i);
        let value = unsafe { self.entries_ptr().add(i).read().assume_init() };

        unsafe {
            self.free_stack_ptr().add(self.free_len).write(key);
        }
        self.free_len += 1;

        value
    }
}

impl<T, K: Key> BoundedStorage<T> for BoxedStorage<T, K> {
    #[inline]
    fn try_insert(&mut self, value: T) -> Result<Self::Key, Full<T>> {
        if self.free_len == 0 {
            return Err(Full(value));
        }

        self.free_len -= 1;
        let key = unsafe { *self.free_stack_ptr().add(self.free_len) };
        let i = key.as_usize();

        unsafe {
            self.entries_ptr().add(i).write(MaybeUninit::new(value));
        }
        self.set_occupied(i);

        Ok(key)
    }

    #[inline]
    fn capacity(&self) -> usize {
        self.capacity
    }
}

impl<T, K: Key> Drop for BoxedStorage<T, K> {
    fn drop(&mut self) {
        // Drop all occupied entries
        for i in 0..self.capacity {
            if self.is_occupied(i) {
                unsafe {
                    self.entries_ptr().add(i).read().assume_init_drop();
                }
            }
        }

        // Deallocate
        unsafe {
            dealloc(self.ptr.as_ptr(), self.layout);
        }
    }
}

// Safety: BoxedStorage owns its data, safe to send if T is Send
unsafe impl<T: Send, K: Key> Send for BoxedStorage<T, K> {}

// =============================================================================
// HashMap implementation (UnboundedStorage for Keyed values)
// =============================================================================

impl<K, V, S> Storage<V> for HashMap<K, V, S>
where
    K: Key + Hash + Eq,
    S: BuildHasher,
{
    type Key = K;

    #[inline]
    fn remove(&mut self, key: Self::Key) -> Option<V> {
        self.remove(&key)
    }

    #[inline]
    fn get(&self, key: Self::Key) -> Option<&V> {
        self.get(&key)
    }

    #[inline]
    fn get_mut(&mut self, key: Self::Key) -> Option<&mut V> {
        self.get_mut(&key)
    }

    #[inline]
    fn len(&self) -> usize {
        self.len()
    }

    #[inline]
    unsafe fn get_unchecked(&self, key: Self::Key) -> &V {
        unsafe { self.get(&key).unwrap_unchecked() }
    }

    #[inline]
    unsafe fn get_unchecked_mut(&mut self, key: Self::Key) -> &mut V {
        unsafe { self.get_mut(&key).unwrap_unchecked() }
    }

    #[inline]
    unsafe fn remove_unchecked(&mut self, key: Self::Key) -> V {
        unsafe { self.remove(&key).unwrap_unchecked() }
    }
}

impl<K, V, S> UnboundedStorage<V> for HashMap<K, V, S>
where
    K: Key + Hash + Eq + Clone,
    V: Keyed<Key = K>,
    S: BuildHasher + Default,
{
    #[inline]
    fn insert(&mut self, value: V) -> Self::Key {
        let key = value.key();
        self.insert(key.clone(), value);
        key
    }
}

// =============================================================================
// slab::Slab implementation (UnboundedStorage)
// =============================================================================

#[cfg(feature = "slab")]
impl<T> Storage<T> for slab::Slab<T> {
    type Key = usize;

    #[inline]
    fn remove(&mut self, key: Self::Key) -> Option<T> {
        self.try_remove(key)
    }

    #[inline]
    fn get(&self, key: Self::Key) -> Option<&T> {
        self.get(key)
    }

    #[inline]
    fn get_mut(&mut self, key: Self::Key) -> Option<&mut T> {
        self.get_mut(key)
    }

    #[inline]
    fn len(&self) -> usize {
        self.len()
    }

    #[inline]
    unsafe fn get_unchecked(&self, key: Self::Key) -> &T {
        unsafe { self.get(key).unwrap_unchecked() }
    }

    #[inline]
    unsafe fn get_unchecked_mut(&mut self, key: Self::Key) -> &mut T {
        unsafe { self.get_mut(key).unwrap_unchecked() }
    }

    #[inline]
    unsafe fn remove_unchecked(&mut self, key: Self::Key) -> T {
        self.remove(key)
    }
}

#[cfg(feature = "slab")]
impl<T> UnboundedStorage<T> for slab::Slab<T> {
    #[inline]
    fn insert(&mut self, value: T) -> Self::Key {
        self.insert(value)
    }
}

// =============================================================================
// nexus_slab implementation
// =============================================================================

#[cfg(feature = "nexus-slab")]
use nexus_slab::{DynamicSlab, FixedSlab, Key as NexusKey};

// Base Storage impl - works for both modes via const generic
#[cfg(feature = "nexus-slab")]
impl<T, const MODE: bool> Storage<T> for nexus_slab::Slab<T, MODE> {
    type Key = NexusKey;

    #[inline]
    fn remove(&mut self, key: Self::Key) -> Option<T> {
        if self.contains(key) {
            Some(nexus_slab::Slab::remove(self, key))
        } else {
            None
        }
    }

    #[inline]
    fn get(&self, key: Self::Key) -> Option<&T> {
        nexus_slab::Slab::get(self, key)
    }

    #[inline]
    fn get_mut(&mut self, key: Self::Key) -> Option<&mut T> {
        nexus_slab::Slab::get_mut(self, key)
    }

    #[inline]
    fn len(&self) -> usize {
        nexus_slab::Slab::len(self)
    }

    #[inline]
    unsafe fn get_unchecked(&self, key: Self::Key) -> &T {
        unsafe { nexus_slab::Slab::get_unchecked(self, key) }
    }

    #[inline]
    unsafe fn get_unchecked_mut(&mut self, key: Self::Key) -> &mut T {
        unsafe { nexus_slab::Slab::get_unchecked_mut(self, key) }
    }

    #[inline]
    unsafe fn remove_unchecked(&mut self, key: Self::Key) -> T {
        // 0.5.0 remove() panics on invalid key, so this is safe if caller upholds invariant
        nexus_slab::Slab::remove(self, key)
    }
}

// FixedSlab: bounded, fallible insert
#[cfg(feature = "nexus-slab")]
impl<T> BoundedStorage<T> for FixedSlab<T> {
    #[inline]
    fn try_insert(&mut self, value: T) -> Result<Self::Key, Full<T>> {
        self.try_insert(value).map_err(|e| Full(e.0))
    }

    #[inline]
    fn capacity(&self) -> usize {
        nexus_slab::Slab::capacity(self)
    }
}

// DynamicSlab: unbounded, infallible insert
#[cfg(feature = "nexus-slab")]
impl<T> UnboundedStorage<T> for DynamicSlab<T> {
    #[inline]
    fn insert(&mut self, value: T) -> Self::Key {
        nexus_slab::Slab::insert(self, value)
    }
}

// =============================================================================
// Helper functions
// =============================================================================

#[inline]
const fn bitmap_words(capacity: usize) -> usize {
    (capacity + 63) / 64
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn new_is_empty() {
        let storage: BoxedStorage<u64> = BoxedStorage::with_capacity(16);
        assert!(storage.is_empty());
        assert!(!storage.is_full());
        assert_eq!(storage.len(), 0);
        assert_eq!(storage.capacity(), 16);
    }

    #[test]
    fn capacity_rounds_to_power_of_two() {
        let storage: BoxedStorage<u64> = BoxedStorage::with_capacity(100);
        assert_eq!(storage.capacity(), 128);

        let storage: BoxedStorage<u64> = BoxedStorage::with_capacity(1000);
        assert_eq!(storage.capacity(), 1024);
    }

    #[test]
    fn insert_get_remove() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(16);

        let key = storage.try_insert(42).unwrap();
        assert_eq!(storage.len(), 1);
        assert_eq!(storage.get(key), Some(&42));

        let removed = storage.remove(key);
        assert_eq!(removed, Some(42));
        assert_eq!(storage.get(key), None);
        assert_eq!(storage.len(), 0);
    }

    #[test]
    fn get_mut() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(16);

        let key = storage.try_insert(10).unwrap();
        *storage.get_mut(key).unwrap() = 20;

        assert_eq!(storage.get(key), Some(&20));
    }

    #[test]
    fn fill_to_capacity() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(4);

        let k0 = storage.try_insert(0).unwrap();
        let k1 = storage.try_insert(1).unwrap();
        let k2 = storage.try_insert(2).unwrap();
        let k3 = storage.try_insert(3).unwrap();

        assert!(storage.is_full());

        let err = storage.try_insert(4);
        assert!(err.is_err());
        assert_eq!(err.unwrap_err().into_inner(), 4);

        assert_eq!(storage.get(k0), Some(&0));
        assert_eq!(storage.get(k1), Some(&1));
        assert_eq!(storage.get(k2), Some(&2));
        assert_eq!(storage.get(k3), Some(&3));
    }

    #[test]
    fn slot_reuse() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(4);

        let k0 = storage.try_insert(0).unwrap();
        let _k1 = storage.try_insert(1).unwrap();

        storage.remove(k0);

        // Next insert reuses k0's slot (LIFO)
        let k2 = storage.try_insert(2).unwrap();
        assert_eq!(k2, k0);
    }

    #[test]
    fn remove_nonexistent() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(16);

        let key = storage.try_insert(42).unwrap();
        storage.remove(key);

        // Double remove returns None
        assert_eq!(storage.remove(key), None);
    }

    #[test]
    fn clear_storage() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(16);

        storage.try_insert(1).unwrap();
        storage.try_insert(2).unwrap();
        storage.try_insert(3).unwrap();

        assert_eq!(storage.len(), 3);

        storage.clear();

        assert_eq!(storage.len(), 0);
        assert!(storage.is_empty());
        assert!(!storage.is_full());
    }

    #[test]
    fn drop_cleans_up() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        static DROP_COUNT: AtomicUsize = AtomicUsize::new(0);

        #[derive(Debug)]
        struct DropCounter;
        impl Drop for DropCounter {
            fn drop(&mut self) {
                DROP_COUNT.fetch_add(1, Ordering::SeqCst);
            }
        }

        DROP_COUNT.store(0, Ordering::SeqCst);

        {
            let mut storage: BoxedStorage<DropCounter> = BoxedStorage::with_capacity(8);
            storage.try_insert(DropCounter).unwrap();
            storage.try_insert(DropCounter).unwrap();
            storage.try_insert(DropCounter).unwrap();
        }

        assert_eq!(DROP_COUNT.load(Ordering::SeqCst), 3);
    }

    #[test]
    fn large_capacity() {
        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(4096);
        assert_eq!(storage.capacity(), 4096);

        // Fill it
        let mut keys = Vec::with_capacity(4096);
        for i in 0..4096 {
            keys.push(storage.try_insert(i as u64).unwrap());
        }
        assert!(storage.is_full());

        // Verify all values
        for (i, key) in keys.iter().enumerate() {
            assert_eq!(storage.get(*key), Some(&(i as u64)));
        }
    }

    #[test]
    fn u16_key() {
        let mut storage: BoxedStorage<u64, u16> = BoxedStorage::with_capacity(100);

        let key = storage.try_insert(42).unwrap();
        assert_eq!(storage.get(key), Some(&42));
    }

    // =========================================================================
    // HashMap tests
    // =========================================================================

    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    struct OrderId(u64);

    // Implement Key trait for OrderId
    impl crate::Key for OrderId {
        const NONE: Self = OrderId(u64::MAX);

        fn from_usize(val: usize) -> Self {
            OrderId(val as u64)
        }

        fn as_usize(&self) -> usize {
            self.0 as usize
        }

        fn is_none(&self) -> bool {
            self.0 == u64::MAX
        }
    }

    #[derive(Debug)]
    struct Order {
        id: OrderId,
        price: u64,
    }

    impl Keyed for Order {
        type Key = OrderId;

        fn key(&self) -> Self::Key {
            self.id.clone()
        }
    }

    #[test]
    fn hashmap_storage_basic() {
        let mut storage: HashMap<OrderId, Order> = HashMap::new();

        let order = Order {
            id: OrderId(1),
            price: 100,
        };

        let key = UnboundedStorage::insert(&mut storage, order);
        assert_eq!(key, OrderId(1));

        let retrieved = Storage::get(&storage, OrderId(1));
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().price, 100);
    }

    #[test]
    fn hashmap_storage_remove() {
        let mut storage: HashMap<OrderId, Order> = HashMap::new();

        UnboundedStorage::insert(
            &mut storage,
            Order {
                id: OrderId(1),
                price: 100,
            },
        );
        UnboundedStorage::insert(
            &mut storage,
            Order {
                id: OrderId(2),
                price: 200,
            },
        );

        assert_eq!(Storage::len(&storage), 2);

        let removed = Storage::remove(&mut storage, OrderId(1));
        assert!(removed.is_some());
        assert_eq!(removed.unwrap().price, 100);

        assert_eq!(Storage::len(&storage), 1);
        assert!(Storage::get(&storage, OrderId(1)).is_none());
    }

    #[test]
    fn hashmap_storage_get_mut() {
        let mut storage: HashMap<OrderId, Order> = HashMap::new();

        UnboundedStorage::insert(
            &mut storage,
            Order {
                id: OrderId(1),
                price: 100,
            },
        );

        if let Some(order) = Storage::get_mut(&mut storage, OrderId(1)) {
            order.price = 150;
        }

        assert_eq!(Storage::get(&storage, OrderId(1)).unwrap().price, 150);
    }

    // =========================================================================
    // slab::Slab tests
    // =========================================================================

    #[cfg(feature = "slab")]
    mod slab_tests {
        use super::*;

        #[test]
        fn insert_get_remove() {
            let mut storage = slab::Slab::new();

            let key = UnboundedStorage::insert(&mut storage, 42);
            assert_eq!(Storage::get(&storage, key), Some(&42));

            let removed = Storage::remove(&mut storage, key);
            assert_eq!(removed, Some(42));
            assert_eq!(Storage::get(&storage, key), None);
        }

        #[test]
        fn slot_reuse() {
            let mut storage = slab::Slab::new();

            let key1 = UnboundedStorage::insert(&mut storage, 1);
            Storage::remove(&mut storage, key1);

            let key2 = UnboundedStorage::insert(&mut storage, 2);
            assert_eq!(key1, key2); // Slot reused
        }
    }

    // =========================================================================
    // nexus_slab::Slab tests
    // =========================================================================

    #[cfg(feature = "nexus-slab")]
    mod nexus_slab_tests {
        use super::*;

        #[test]
        fn insert_get_remove() {
            let mut storage: nexus_slab::DynamicSlab<u64> = nexus_slab::Slab::with_capacity(16).unwrap();

            let key = UnboundedStorage::insert(&mut storage, 42);
            assert_eq!(Storage::get(&storage, key), Some(&42));

            let removed = Storage::remove(&mut storage, key);
            assert_eq!(removed, Some(42));
            assert_eq!(Storage::get(&storage, key), None);
        }
    }

    // =========================================================================
    // Benchmarks
    // =========================================================================

    #[test]
    #[ignore]
    fn bench_boxed_storage() {
        use std::time::Instant;

        const CAPACITY: usize = 4096;
        const ITERATIONS: usize = 100_000;

        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(CAPACITY);

        // Warmup
        for i in 0..CAPACITY {
            storage.try_insert(i as u64).unwrap();
        }
        for i in 0..CAPACITY {
            storage.remove(u32::from_usize(i));
        }

        // Collect timings
        let mut insert_ns = Vec::with_capacity(ITERATIONS);
        let mut get_ns = Vec::with_capacity(ITERATIONS);
        let mut remove_ns = Vec::with_capacity(ITERATIONS);

        for i in 0..ITERATIONS {
            // Insert
            let start = Instant::now();
            let key = storage.try_insert(i as u64).unwrap();
            insert_ns.push(start.elapsed().as_nanos() as u64);

            // Get
            let start = Instant::now();
            let _ = std::hint::black_box(storage.get(key));
            get_ns.push(start.elapsed().as_nanos() as u64);

            // Remove
            let start = Instant::now();
            let _ = std::hint::black_box(storage.remove(key));
            remove_ns.push(start.elapsed().as_nanos() as u64);
        }

        // Sort for percentiles
        insert_ns.sort_unstable();
        get_ns.sort_unstable();
        remove_ns.sort_unstable();

        fn percentile(sorted: &[u64], p: f64) -> u64 {
            let idx = ((p / 100.0) * sorted.len() as f64) as usize;
            sorted[idx.min(sorted.len() - 1)]
        }

        fn print_stats(name: &str, sorted: &[u64]) {
            println!(
                "{:8} | p50: {:4} ns | p90: {:4} ns | p99: {:4} ns | p999: {:5} ns",
                name,
                percentile(sorted, 50.0),
                percentile(sorted, 90.0),
                percentile(sorted, 99.0),
                percentile(sorted, 99.9),
            );
        }

        println!(
            "\nBoxedStorage<u64> ({} iterations, capacity {})",
            ITERATIONS, CAPACITY
        );
        println!("---------------------------------------------------------");
        print_stats("insert", &insert_ns);
        print_stats("get", &get_ns);
        print_stats("remove", &remove_ns);
        println!();
    }

    #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
    #[test]
    #[ignore]
    fn bench_boxed_storage_tsc() {
        const CAPACITY: usize = 4096;
        const ITERATIONS: usize = 100_000;

        #[inline]
        fn rdtsc() -> u64 {
            unsafe {
                core::arch::x86_64::_mm_lfence();
                core::arch::x86_64::_rdtsc()
            }
        }

        let mut storage: BoxedStorage<u64> = BoxedStorage::with_capacity(CAPACITY);

        // Warmup
        for i in 0..CAPACITY {
            storage.try_insert(i as u64).unwrap();
        }
        for i in 0..CAPACITY {
            storage.remove(u32::from_usize(i));
        }

        // Collect timings
        let mut insert_cycles = Vec::with_capacity(ITERATIONS);
        let mut get_cycles = Vec::with_capacity(ITERATIONS);
        let mut remove_cycles = Vec::with_capacity(ITERATIONS);

        for i in 0..ITERATIONS {
            // Insert
            let start = rdtsc();
            let key = storage.try_insert(i as u64).unwrap();
            let end = rdtsc();
            insert_cycles.push(end - start);

            // Get
            let start = rdtsc();
            let _ = std::hint::black_box(storage.get(key));
            let end = rdtsc();
            get_cycles.push(end - start);

            // Remove
            let start = rdtsc();
            let _ = std::hint::black_box(storage.remove(key));
            let end = rdtsc();
            remove_cycles.push(end - start);
        }

        // Sort for percentiles
        insert_cycles.sort_unstable();
        get_cycles.sort_unstable();
        remove_cycles.sort_unstable();

        fn percentile(sorted: &[u64], p: f64) -> u64 {
            let idx = ((p / 100.0) * sorted.len() as f64) as usize;
            sorted[idx.min(sorted.len() - 1)]
        }

        fn print_stats(name: &str, sorted: &[u64]) {
            println!(
                "{:8} | p50: {:5} cycles | p90: {:5} cycles | p99: {:5} cycles | p999: {:6} cycles",
                name,
                percentile(sorted, 50.0),
                percentile(sorted, 90.0),
                percentile(sorted, 99.0),
                percentile(sorted, 99.9),
            );
        }

        println!(
            "\nBoxedStorage<u64> ({} iterations, capacity {})",
            ITERATIONS, CAPACITY
        );
        println!("------------------------------------------------------------------------");
        print_stats("insert", &insert_cycles);
        print_stats("get", &get_cycles);
        print_stats("remove", &remove_cycles);
        println!();
    }
}
