armdb 0.1.13

use crate::Key;
use crate::codec::Codec;
use crate::compaction::{CompactionIndex, compact_shard};
use crate::config::Config;
use crate::disk_loc::DiskLoc;
use crate::engine::Engine;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, TypedWriteHook};
use crate::recovery::recover_typed_tree;
use crate::shard::ShardInner;
use crate::skiplist::node::{SkipNode, TypedData, TypedNode, random_height};
use crate::skiplist::{InsertResult, SkipList};
use crate::sync::MutexGuard;
use std::mem::size_of;
use std::ops::Bound;

/// Guard-protected reference to a typed value inside a `TypedTree`.
///
/// Dereferences to `T`. The `seize` guard prevents the referenced data
/// from being reclaimed while this handle exists. No `Clone` required on `T`.
pub struct TypedRef<'a, T> {
    _guard: seize::LocalGuard<'a>,
    data: *const TypedData<T>,
    _marker: std::marker::PhantomData<&'a T>,
}

// SAFETY: TypedRef holds a seize guard that prevents reclamation of the data.
// The pointer is valid for the lifetime of the guard.
unsafe impl<T: Send + Sync> Send for TypedRef<'_, T> {}
unsafe impl<T: Send + Sync> Sync for TypedRef<'_, T> {}

impl<'a, T> TypedRef<'a, T> {
    pub(crate) fn new(guard: seize::LocalGuard<'a>, data: *const TypedData<T>) -> Self {
        Self {
            _guard: guard,
            data,
            _marker: std::marker::PhantomData,
        }
    }
}

impl<T> std::ops::Deref for TypedRef<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        // SAFETY: data pointer is valid while guard is alive
        unsafe { &(*self.data).value }
    }
}

/// A tree with fixed-size keys and typed values `T`. Values are encoded via
/// a [`Codec`] for disk persistence but stored as `T` in memory — reads never
/// touch disk and return [`TypedRef<T>`](TypedRef) (guard-protected reference).
///
/// Each `TypedTree` owns its storage engine — one tree = one database directory.
///
/// # Clone-free design
///
/// Unlike `ConstTree` which copies `[u8; V]` values (they are `Copy`), `TypedTree`
/// returns [`TypedRef<T>`](TypedRef) — a guard-protected reference. The `seize` guard
/// inside `TypedRef` prevents reclamation of the old data while the reference exists.
/// This means `put()`, `delete()`, and `update()` can return the old value **without
/// requiring `T: Clone`**.
///
/// The only method that requires `T: Clone` is [`compact()`](Self::compact) — compaction
/// must copy the value into a new `TypedData` with an updated disk location.
///
/// # Write hooks
///
/// Uses [`TypedWriteHook<K, T>`] instead of `WriteHook<K>`. The hook receives `&T`
/// directly (not encoded bytes). `on_write` fires on `put`/`insert`/`delete`/`cas`/`update`.
/// Does **not** fire inside `atomic()`. Old value is always provided (it lives in
/// memory) — `NEEDS_OLD_VALUE` is ignored.
///
/// `on_init` fires once per live entry during `migrate()` or `replay_init()`
/// (enable via `NEEDS_INIT = true`).
///
/// # Usage
///
/// ```ignore
/// let tree = TypedTree::<16, MyValue, RapiraCodec>::open(
///     "data/users",
///     Config::default(),
///     RapiraCodec,
/// )?;
/// tree.put(&key, value)?;
/// if let Some(r) = tree.get(&key) {
///     println!("{:?}", &*r);  // TypedRef<MyValue> derefs to &MyValue
/// }
/// tree.close()?;
/// ```
///
/// # Iteration
///
/// `iter()`, `range()`, and `prefix_iter()` all return [`TypedIter`] which
/// implements `Iterator + DoubleEndedIterator` with `Item = (K, &T)`.
/// Lock-free, zero disk I/O.
///
/// ```ignore
/// for (key, value) in tree.iter() { }
/// let latest = tree.prefix_iter(&user_id).take(20).collect::<Vec<_>>();
/// let oldest = tree.iter().rev().take(5);  // DoubleEndedIterator
/// ```
///
/// # Features
///
/// Requires the `typed-tree` feature. Codec implementations:
/// - `rapira-codec` — [`RapiraCodec`](crate::RapiraCodec) for `T: rapira::Rapira`
/// - `bitcode-codec` — [`BitcodeCodec`](crate::BitcodeCodec) for `T: bitcode::Encode + Decode`
pub struct TypedTree<K: Key, T: Send + Sync, C: Codec<T>, H: TypedWriteHook<K, T> = NoHook> {
    index: SkipList<TypedNode<K, T>>,
    engine: Engine,
    codec: C,
    compaction_threshold: f64,
    shard_prefix_bits: usize,
    reversed: bool,
    hook: H,
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Sync> TypedTree<K, T, C> {
    /// Open or create a `TypedTree` at the given path.
    /// Recovers the index from existing data files on disk.
    pub fn open(path: impl AsRef<std::path::Path>, config: Config, codec: C) -> DbResult<Self> {
        Self::open_inner(path, config, codec, NoHook)
    }
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>> TypedTree<K, T, C, H> {
    /// Open or create a `TypedTree` with a write hook for secondary index maintenance.
    pub fn open_hooked(
        path: impl AsRef<std::path::Path>,
        config: Config,
        codec: C,
        hook: H,
    ) -> DbResult<Self> {
        Self::open_inner(path, config, codec, hook)
    }

    fn open_inner(
        path: impl AsRef<std::path::Path>,
        config: Config,
        codec: C,
        hook: H,
    ) -> DbResult<Self> {
        let compaction_threshold = config.compaction_threshold;
        let shard_prefix_bits = config.shard_prefix_bits;
        let reversed = config.reversed;
        let engine = Engine::open(path, config)?;

        let tree = Self {
            index: SkipList::new(reversed),
            engine,
            codec,
            compaction_threshold,
            shard_prefix_bits,
            reversed,
            hook,
        };

        // Recover index from disk
        let shard_dirs = tree.engine.shard_dirs();
        let shard_dir_refs = Engine::shard_dir_refs(&shard_dirs);
        let shard_ids = tree.engine.shard_ids();

        let hints = tree.engine.hints();
        let max_gsn = recover_typed_tree::<K, T, C>(
            &shard_dir_refs,
            &shard_ids,
            tree.index(),
            &tree.codec,
            hints,
            #[cfg(feature = "encryption")]
            tree.engine.cipher(),
        )?;

        tree.engine
            .gsn()
            .fetch_max(max_gsn + 1, std::sync::atomic::Ordering::Relaxed);
        if hints {
            for shard in tree.engine.shards().iter() {
                shard.set_key_len(size_of::<K>());
            }
        }
        tracing::info!(
            key_size = size_of::<K>(),
            entries = tree.len(),
            "typed_tree recovered"
        );

        Ok(tree)
    }

    /// Graceful shutdown: write hint files (if enabled), flush write buffers + fsync.
    pub fn close(self) -> DbResult<()> {
        if self.engine.hints() {
            self.sync_hints()?;
        }
        self.engine.flush()
    }

    /// Flush all shard write buffers to disk (without fsync).
    pub fn flush_buffers(&self) -> DbResult<()> {
        self.engine.flush_buffers()
    }

    /// Get the database configuration.
    pub fn config(&self) -> &Config {
        self.engine.config()
    }
}

impl<K: Key, T: Clone + Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>> CompactionIndex<K>
    for TypedTree<K, T, C, H>
{
    fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
        let guard = self.index.collector().enter();
        if let Some(node) = self.index.get(key.as_bytes(), &guard) {
            let current_data = node.load_data();
            if current_data.disk == old_loc {
                let new_data = Box::into_raw(Box::new(TypedData {
                    disk: new_loc,
                    value: current_data.value.clone(),
                }));
                let old_data_ptr = node.swap_data(new_data);
                unsafe {
                    self.index
                        .collector()
                        .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
                }
                return true;
            }
        }
        false
    }

    fn contains_key(&self, key: &K) -> bool {
        self.contains(key)
    }
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>> TypedTree<K, T, C, H> {
    /// Trigger a compaction pass across all shards.
    pub fn compact(&self) -> DbResult<usize>
    where
        T: Clone,
    {
        let mut total_compacted = 0;
        for shard in self.engine.shards().iter() {
            total_compacted += compact_shard(shard, self, self.compaction_threshold)?;
        }
        Ok(total_compacted)
    }

    /// Get a guard-protected reference to a value by key. Lock-free, zero disk I/O.
    pub fn get(&self, key: &K) -> Option<TypedRef<'_, T>> {
        metrics::counter!("armdb.ops", "op" => "get", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.get");
        let guard = self.index.collector().enter();
        let node = self.index.get(key.as_bytes(), &guard)?;
        // SAFETY: load_data returns a pointer valid while guard is alive.
        // We store the pointer (not a reference) so the borrow of guard is released,
        // allowing us to move guard into TypedRef.
        let data = node.load_data() as *const TypedData<T>;
        Some(TypedRef {
            _guard: guard,
            data,
            _marker: std::marker::PhantomData,
        })
    }

    /// Get a guard-protected reference to a value by key, returning `Err(KeyNotFound)` if absent.
    pub fn get_or_err(&self, key: &K) -> DbResult<TypedRef<'_, T>> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

    /// Insert or update a key-value pair. Returns a `TypedRef` to the old value
    /// if the key existed (valid while the guard lives).
    pub fn put(&self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "put", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.put");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let guard = self.index.collector().enter();
        self.put_locked(shard_id, &mut inner, guard, key, value)
    }

    /// Insert a key-value pair only if the key does not exist.
    /// Returns `Err(KeyExists)` if the key is already present.
    pub fn insert(&self, key: &K, value: T) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "insert", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.insert");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let guard = self.index.collector().enter();
        self.insert_locked(shard_id, &mut inner, &guard, key, value)?;
        Ok(())
    }

    /// Delete a key. Returns a `TypedRef` to the old value if the key existed.
    pub fn delete(&self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "delete", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.delete");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let guard = self.index.collector().enter();
        self.delete_locked(shard_id, &mut inner, guard, key)
    }

    /// Atomically execute multiple operations on a single shard.
    /// All keys must route to the same shard as `shard_key`.
    /// The closure must be short — shard lock is held for its duration.
    pub fn atomic<R>(
        &self,
        shard_key: &K,
        f: impl FnOnce(&mut TypedShard<'_, K, T, C, H>) -> DbResult<R>,
    ) -> DbResult<R> {
        let shard_id = self.shard_for(shard_key);
        let inner = self.engine.shards()[shard_id].lock();
        let guard = self.index.collector().enter();
        let mut shard = TypedShard {
            tree: self,
            inner,
            shard_id,
            guard,
        };
        f(&mut shard)
    }

    fn put_locked<'g>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: seize::LocalGuard<'g>,
        key: &K,
        value: T,
    ) -> DbResult<Option<TypedRef<'g, T>>> {
        self.put_locked_inner::<true>(shard_id, inner, guard, key, value)
    }

    fn put_locked_inner<'g, const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: seize::LocalGuard<'g>,
        key: &K,
        value: T,
    ) -> DbResult<Option<TypedRef<'g, T>>> {
        let mut buf = Vec::new();
        self.codec.encode_to(&value, &mut buf);
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

        // Fast path: key exists — atomic swap, no node allocation
        if let Some(existing) = self.index.get(key.as_bytes(), &guard) {
            // SAFETY: load_data returns a pointer valid while guard is alive.
            let old_data = existing.load_data() as *const TypedData<T>;
            let new_data = Box::new(TypedData {
                disk: disk_loc,
                value,
            });

            if HOOKS {
                self.hook.on_write(
                    key,
                    Some(unsafe { &(*old_data).value }),
                    Some(&new_data.value),
                );
            }

            let new_data_ptr = Box::into_raw(new_data);
            let old_data_ptr = existing.swap_data(new_data_ptr);
            let old_disk = unsafe { (*old_data_ptr).disk };
            inner.add_dead_bytes(
                old_disk.file_id as u32,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );
            unsafe {
                self.index
                    .collector()
                    .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
            return Ok(Some(TypedRef {
                _guard: guard,
                data: old_data,
                _marker: std::marker::PhantomData,
            }));
        }

        // Slow path: new key — allocate node + insert
        let height = random_height();
        let node_ptr = TypedNode::alloc(*key, value, disk_loc, height);

        match self.index.insert(node_ptr, &guard) {
            InsertResult::Inserted => {
                if HOOKS {
                    let data = unsafe { &*node_ptr }.load_data();
                    self.hook.on_write(key, None, Some(&data.value));
                }
                Ok(None)
            }
            InsertResult::Exists(existing) => {
                // Race: another shard inserted same key between get and insert
                let old_data = existing.load_data() as *const TypedData<T>;
                let new_data_ptr = unsafe {
                    let ptr = (*node_ptr).data.load(std::sync::atomic::Ordering::Relaxed);
                    (*node_ptr)
                        .data
                        .store(std::ptr::null_mut(), std::sync::atomic::Ordering::Relaxed);
                    ptr
                };
                let old_data_ptr = existing.swap_data(new_data_ptr);
                let old_disk = unsafe { (*old_data_ptr).disk };
                inner.add_dead_bytes(
                    old_disk.file_id as u32,
                    crate::entry::entry_size(size_of::<K>(), old_disk.len),
                );
                unsafe {
                    self.index
                        .collector()
                        .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
                    TypedNode::<K, T>::dealloc_node(node_ptr);
                }

                if HOOKS {
                    self.hook.on_write(
                        key,
                        Some(unsafe { &(*old_data).value }),
                        Some(unsafe { &(*new_data_ptr).value }),
                    );
                }

                Ok(Some(TypedRef {
                    _guard: guard,
                    data: old_data,
                    _marker: std::marker::PhantomData,
                }))
            }
        }
    }

    fn insert_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: &seize::LocalGuard<'_>,
        key: &K,
        value: T,
    ) -> DbResult<()> {
        if self.index.get(key.as_bytes(), guard).is_some() {
            return Err(DbError::KeyExists);
        }

        let mut buf = Vec::new();
        self.codec.encode_to(&value, &mut buf);
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;
        let height = random_height();
        let node_ptr = TypedNode::alloc(*key, value, disk_loc, height);
        self.index.insert(node_ptr, guard);

        let data = unsafe { &*node_ptr }.load_data();
        self.hook.on_write(key, None, Some(&data.value));
        Ok(())
    }

    fn delete_locked<'g>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: seize::LocalGuard<'g>,
        key: &K,
    ) -> DbResult<Option<TypedRef<'g, T>>> {
        self.delete_locked_inner::<true>(shard_id, inner, guard, key)
    }

    fn delete_locked_inner<'g, const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: seize::LocalGuard<'g>,
        key: &K,
    ) -> DbResult<Option<TypedRef<'g, T>>> {
        let node = match self.index.get(key.as_bytes(), &guard) {
            Some(n) => n,
            None => return Ok(None),
        };
        // SAFETY: load_data returns a pointer valid while guard is alive.
        let data = node.load_data() as *const TypedData<T>;

        if HOOKS {
            self.hook
                .on_write(key, Some(unsafe { &(*data).value }), None);
        }

        inner.append_entry(shard_id as u8, key.as_bytes(), &[], true)?;
        self.index.remove(key.as_bytes(), &guard);

        let disk = unsafe { (*data).disk };
        inner.add_dead_bytes(
            disk.file_id as u32,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );

        Ok(Some(TypedRef {
            _guard: guard,
            data,
            _marker: std::marker::PhantomData,
        }))
    }

    /// Compare-and-swap: if current value == expected, replace with new_value.
    /// Returns `Ok(())` on success, `Err(CasMismatch)` if current != expected,
    /// `Err(KeyNotFound)` if key doesn't exist.
    pub fn cas(&self, key: &K, expected: &T, new_value: T) -> DbResult<()>
    where
        T: PartialEq,
    {
        metrics::counter!("armdb.ops", "op" => "cas", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.cas");
        let shard_id = self.shard_for(key);
        let shard = &self.engine.shards()[shard_id];
        let mut inner = shard.lock();

        let guard = self.index.collector().enter();
        let existing = self
            .index
            .get(key.as_bytes(), &guard)
            .ok_or(DbError::KeyNotFound)?;

        let current_data = existing.load_data();
        if current_data.value != *expected {
            return Err(DbError::CasMismatch);
        }

        let mut buf = Vec::new();
        self.codec.encode_to(&new_value, &mut buf);
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

        let new_data = Box::new(TypedData {
            disk: disk_loc,
            value: new_value,
        });

        self.hook
            .on_write(key, Some(&current_data.value), Some(&new_data.value));

        let new_data_ptr = Box::into_raw(new_data);
        let old_data = existing.swap_data(new_data_ptr);
        let old_disk = unsafe { (*old_data).disk };
        inner.add_dead_bytes(
            old_disk.file_id as u32,
            crate::entry::entry_size(size_of::<K>(), old_disk.len),
        );
        unsafe {
            self.index
                .collector()
                .retire(old_data, seize::reclaim::boxed::<TypedData<T>>);
        }

        Ok(())
    }

    /// Atomically read-modify-write. Returns `Some(TypedRef)` to the **new** value
    /// if key existed, `None` otherwise.
    /// The closure must not be heavy (shard lock is held).
    pub fn update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<TypedRef<'_, T>>> {
        self.update_inner(key, f, false)
    }

    /// Like [`update()`](Self::update), but returns `Some(TypedRef)` to the **old** value.
    pub fn fetch_update(
        &self,
        key: &K,
        f: impl FnOnce(&T) -> T,
    ) -> DbResult<Option<TypedRef<'_, T>>> {
        self.update_inner(key, f, true)
    }

    fn update_inner(
        &self,
        key: &K,
        f: impl FnOnce(&T) -> T,
        return_old: bool,
    ) -> DbResult<Option<TypedRef<'_, T>>> {
        metrics::counter!("armdb.ops", "op" => "update", "tree" => "typed_tree").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.update");
        let shard_id = self.shard_for(key);
        let shard = &self.engine.shards()[shard_id];
        let mut inner = shard.lock();

        let guard = self.index.collector().enter();
        let existing = match self.index.get(key.as_bytes(), &guard) {
            Some(n) => n,
            None => return Ok(None),
        };

        let old_data = existing.load_data();
        let new_value = f(&old_data.value);

        let mut buf = Vec::new();
        self.codec.encode_to(&new_value, &mut buf);
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;

        let new_data = Box::new(TypedData {
            disk: disk_loc,
            value: new_value,
        });

        self.hook
            .on_write(key, Some(&old_data.value), Some(&new_data.value));

        let new_data_ptr = Box::into_raw(new_data);
        let old_data_ptr = existing.swap_data(new_data_ptr);
        let old_disk = unsafe { (*old_data_ptr).disk };
        inner.add_dead_bytes(
            old_disk.file_id as u32,
            crate::entry::entry_size(size_of::<K>(), old_disk.len),
        );
        unsafe {
            self.index
                .collector()
                .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
        }

        let data = if return_old {
            old_data_ptr
        } else {
            new_data_ptr
        };
        Ok(Some(TypedRef {
            _guard: guard,
            data,
            _marker: std::marker::PhantomData,
        }))
    }

    /// Check if a key exists.
    pub fn contains(&self, key: &K) -> bool {
        let guard = self.index.collector().enter();
        self.index.get(key.as_bytes(), &guard).is_some()
    }

    /// Return the first entry in index order, or `None` if empty.
    /// With `reversed=true` (default): the entry with the largest key.
    /// O(1) — follows head's level-0 pointer, skipping marked nodes.
    pub fn first(&self) -> Option<(K, TypedRef<'_, T>)> {
        let guard = self.index.collector().enter();
        let mut ptr = crate::skiplist::strip_mark(unsafe {
            (*self.index.head_ptr())
                .tower(0)
                .load(std::sync::atomic::Ordering::Acquire)
        });
        while !ptr.is_null() {
            let node = unsafe { &*ptr };
            if !node.is_marked() {
                let data = node.load_data() as *const TypedData<T>;
                return Some((
                    node.key,
                    TypedRef {
                        _guard: guard,
                        data,
                        _marker: std::marker::PhantomData,
                    },
                ));
            }
            ptr = crate::skiplist::strip_mark(
                node.tower(0).load(std::sync::atomic::Ordering::Acquire),
            );
        }
        None
    }

    /// Return the last entry in index order, or `None` if empty.
    /// With `reversed=true` (default): the entry with the smallest key.
    pub fn last(&self) -> Option<(K, TypedRef<'_, T>)> {
        let guard = self.index.collector().enter();
        let ptr = self.index.find_last();
        if ptr.is_null() {
            return None;
        }
        let node = unsafe { &*ptr };
        let data = node.load_data() as *const TypedData<T>;
        Some((
            node.key,
            TypedRef {
                _guard: guard,
                data,
                _marker: std::marker::PhantomData,
            },
        ))
    }

    // -- Range helpers (front/back pointer positioning) -------------------------

    fn resolve_front_asc(
        &self,
        bound: &Bound<&K>,
        guard: &seize::LocalGuard<'_>,
    ) -> *mut TypedNode<K, T> {
        match bound {
            Bound::Included(k) => self.index.find_first_ge(k.as_bytes(), guard),
            Bound::Excluded(k) => {
                let ge = self.index.find_first_ge(k.as_bytes(), guard);
                if !ge.is_null()
                    && !unsafe { &*ge }.is_marked()
                    && unsafe { &*ge }.key_bytes() == k.as_bytes()
                {
                    crate::skiplist::strip_mark(unsafe {
                        (*ge).tower(0).load(std::sync::atomic::Ordering::Acquire)
                    })
                } else {
                    ge
                }
            }
            Bound::Unbounded => crate::skiplist::strip_mark(unsafe {
                (*self.index.head_ptr())
                    .tower(0)
                    .load(std::sync::atomic::Ordering::Acquire)
            }),
        }
    }

    fn resolve_front_rev(
        &self,
        bound: &Bound<&K>,
        guard: &seize::LocalGuard<'_>,
    ) -> *mut TypedNode<K, T> {
        match bound {
            Bound::Included(k) => self.index.find_first_ge(k.as_bytes(), guard),
            Bound::Excluded(k) => {
                let ge = self.index.find_first_ge(k.as_bytes(), guard);
                if !ge.is_null()
                    && !unsafe { &*ge }.is_marked()
                    && unsafe { &*ge }.key_bytes() == k.as_bytes()
                {
                    crate::skiplist::strip_mark(unsafe {
                        (*ge).tower(0).load(std::sync::atomic::Ordering::Acquire)
                    })
                } else {
                    ge
                }
            }
            Bound::Unbounded => crate::skiplist::strip_mark(unsafe {
                (*self.index.head_ptr())
                    .tower(0)
                    .load(std::sync::atomic::Ordering::Acquire)
            }),
        }
    }

    fn prefix_bounds(&self, prefix: &[u8]) -> (K, Bound<K>) {
        if self.reversed {
            let mut search = K::zeroed();
            for b in search.as_bytes_mut().iter_mut() {
                *b = 0xFF;
            }
            search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
            let mut end_key = K::zeroed();
            end_key.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
            (search, Bound::Included(end_key))
        } else {
            let mut search = K::zeroed();
            search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
            let end = prefix_to_end_bound::<K>(prefix);
            (search, end)
        }
    }

    /// Iterate entries whose keys start with `prefix`.
    ///
    /// `reversed=true` (default): yields matching keys in DESC order.
    /// `next()` is O(1), `next_back()` is O(log n).
    pub fn prefix_iter(&self, prefix: &[u8]) -> TypedIter<'_, K, T> {
        let guard = self.index.collector().enter();
        let (search_key, end) = self.prefix_bounds(prefix);
        let front = self.index.find_first_ge(search_key.as_bytes(), &guard);
        TypedIter {
            list: &self.index,
            front,
            back: None,
            end,
            start: Bound::Included(search_key),
            reversed: self.reversed,
            done: false,
            _guard: guard,
        }
    }

    /// Iterate all entries in index order.
    ///
    /// `reversed=true` (default): DESC. `reversed=false`: ASC.
    /// `next()` is O(1), `next_back()` is O(log n).
    pub fn iter(&self) -> TypedIter<'_, K, T> {
        let guard = self.index.collector().enter();
        let front = crate::skiplist::strip_mark(unsafe {
            (*self.index.head_ptr())
                .tower(0)
                .load(std::sync::atomic::Ordering::Acquire)
        });
        TypedIter {
            list: &self.index,
            front,
            back: None,
            end: Bound::Unbounded,
            start: Bound::Unbounded,
            reversed: self.reversed,
            done: false,
            _guard: guard,
        }
    }

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    ///
    /// `reversed=true` (default): DESC within range. `reversed=false`: ASC.
    /// `next()` is O(1), `next_back()` is O(log n).
    pub fn range(&self, start: &K, end: &K) -> TypedIter<'_, K, T> {
        self.range_bounds(Bound::Included(start), Bound::Excluded(end))
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    ///
    /// Unlike [`range()`](Self::range), allows `Included`, `Excluded`, or `Unbounded`
    /// for each bound independently.
    ///
    /// `reversed=true` (default): DESC within range. `reversed=false`: ASC.
    /// `next()` is O(1), `next_back()` is O(log n).
    pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> TypedIter<'_, K, T> {
        let guard = self.index.collector().enter();
        if self.reversed {
            let front = self.resolve_front_rev(&end, &guard);
            TypedIter {
                list: &self.index,
                front,
                back: None,
                end: bound_owned(&start),
                start: bound_owned(&end),
                reversed: true,
                done: false,
                _guard: guard,
            }
        } else {
            let front = self.resolve_front_asc(&start, &guard);
            TypedIter {
                list: &self.index,
                front,
                back: None,
                end: bound_owned(&end),
                start: bound_owned(&start),
                reversed: false,
                done: false,
                _guard: guard,
            }
        }
    }

    pub fn len(&self) -> usize {
        self.index.len()
    }

    pub fn is_empty(&self) -> bool {
        self.index.is_empty()
    }

    /// Write hint files for all active shard files. Call during graceful shutdown.
    pub fn sync_hints(&self) -> DbResult<()> {
        for shard in self.engine.shards().iter() {
            shard.write_active_hint(size_of::<K>())?;
        }
        Ok(())
    }

    /// Iterate all entries and optionally mutate them. Call once at startup.
    ///
    /// The callback receives each (key, &T) and returns `MigrateAction`:
    /// - `Keep` — no change (fires `on_init` if `NEEDS_INIT`)
    /// - `Update(new_value)` — replace value (hook-free write, fires `on_init`)
    /// - `Delete` — remove entry (hook-free tombstone)
    ///
    /// Returns the number of mutated entries.
    pub fn migrate(&self, f: impl Fn(&K, &T) -> crate::MigrateAction<T>) -> DbResult<usize> {
        use crate::MigrateAction;

        let _guard = self.index.collector().enter();
        let mut current = crate::skiplist::strip_mark(unsafe {
            (*self.index.head_ptr())
                .tower(0)
                .load(std::sync::atomic::Ordering::Acquire)
        });
        let mut count = 0;
        while !current.is_null() {
            let node = unsafe { &*current };
            current = crate::skiplist::strip_mark(
                node.tower(0).load(std::sync::atomic::Ordering::Acquire),
            );
            if node.is_marked() {
                continue;
            }
            let data = node.load_data();
            match f(&node.key, &data.value) {
                MigrateAction::Keep => {
                    if H::NEEDS_INIT {
                        self.hook.on_init(&node.key, &data.value);
                    }
                }
                MigrateAction::Update(value) => {
                    if H::NEEDS_INIT {
                        self.hook.on_init(&node.key, &value);
                    }
                    let shard_id = self.shard_for(&node.key);
                    let mut inner = self.engine.shards()[shard_id].lock();
                    let g = self.index.collector().enter();
                    self.put_locked_inner::<false>(shard_id, &mut inner, g, &node.key, value)?;
                    count += 1;
                }
                MigrateAction::Delete => {
                    let shard_id = self.shard_for(&node.key);
                    let mut inner = self.engine.shards()[shard_id].lock();
                    let g = self.index.collector().enter();
                    self.delete_locked_inner::<false>(shard_id, &mut inner, g, &node.key)?;
                    count += 1;
                }
            }
        }

        tracing::info!(mutations = count, "typed_tree migration complete");
        Ok(count)
    }

    /// Replay `on_init` for every live entry. Used when no migration runs.
    pub(crate) fn replay_init(&self) {
        if !H::NEEDS_INIT {
            return;
        }
        let _guard = self.index.collector().enter();
        let mut current = crate::skiplist::strip_mark(unsafe {
            (*self.index.head_ptr())
                .tower(0)
                .load(std::sync::atomic::Ordering::Acquire)
        });
        while !current.is_null() {
            let node = unsafe { &*current };
            current = crate::skiplist::strip_mark(
                node.tower(0).load(std::sync::atomic::Ordering::Acquire),
            );
            if node.is_marked() {
                continue;
            }
            let data = node.load_data();
            self.hook.on_init(&node.key, &data.value);
        }
    }

    pub(crate) fn index(&self) -> &SkipList<TypedNode<K, T>> {
        &self.index
    }

    pub fn shard_for(&self, key: &K) -> usize {
        if self.shard_prefix_bits == 0 || self.shard_prefix_bits >= size_of::<K>() * 8 {
            let hash = xxhash_rust::xxh3::xxh3_64(key.as_bytes());
            return (hash as usize) % self.engine.shards().len();
        }

        let full_bytes = self.shard_prefix_bits / 8;
        let extra_bits = self.shard_prefix_bits % 8;

        let hash = if extra_bits == 0 {
            xxhash_rust::xxh3::xxh3_64(&key.as_bytes()[..full_bytes])
        } else {
            let mut buf = K::zeroed();
            buf.as_bytes_mut()[..full_bytes + 1].copy_from_slice(&key.as_bytes()[..full_bytes + 1]);
            let mask = !((1u8 << (8 - extra_bits)) - 1);
            buf.as_bytes_mut()[full_bytes] = key.as_bytes()[full_bytes] & mask;
            xxhash_rust::xxh3::xxh3_64(&buf.as_bytes()[..full_bytes + 1])
        };

        (hash as usize) % self.engine.shards().len()
    }
}

/// Handle for atomic multi-key operations within a single shard.
/// Obtained via [`TypedTree::atomic`]. The shard lock is held for the
/// lifetime of this struct — keep the closure short.
pub struct TypedShard<'a, K: Key, T: Send + Sync, C: Codec<T>, H: TypedWriteHook<K, T> = NoHook> {
    tree: &'a TypedTree<K, T, C, H>,
    inner: MutexGuard<'a, ShardInner>,
    shard_id: usize,
    guard: seize::LocalGuard<'a>,
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    TypedShard<'_, K, T, C, H>
{
    pub fn put(&mut self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
        self.check_shard(key)?;
        let guard = self.tree.index.collector().enter();
        self.tree
            .put_locked(self.shard_id, &mut self.inner, guard, key, value)
    }

    pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
        self.check_shard(key)?;
        self.tree
            .insert_locked(self.shard_id, &mut self.inner, &self.guard, key, value)
    }

    pub fn delete(&mut self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
        self.check_shard(key)?;
        let guard = self.tree.index.collector().enter();
        self.tree
            .delete_locked(self.shard_id, &mut self.inner, guard, key)
    }

    pub fn get(&self, key: &K) -> Option<&T> {
        let node = self.tree.index.get(key.as_bytes(), &self.guard)?;
        Some(&node.load_data().value)
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<&T> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

    pub fn contains(&self, key: &K) -> bool {
        self.tree.index.get(key.as_bytes(), &self.guard).is_some()
    }

    fn check_shard(&self, key: &K) -> DbResult<()> {
        if self.tree.shard_for(key) != self.shard_id {
            return Err(DbError::ShardMismatch);
        }
        Ok(())
    }
}

fn bound_owned<K: Copy>(b: &Bound<&K>) -> Bound<K> {
    match b {
        Bound::Included(k) => Bound::Included(**k),
        Bound::Excluded(k) => Bound::Excluded(**k),
        Bound::Unbounded => Bound::Unbounded,
    }
}

fn prefix_to_end_bound<K: Key>(prefix: &[u8]) -> Bound<K> {
    let mut incremented = prefix.to_vec();
    let mut carry = true;
    for byte in incremented.iter_mut().rev() {
        if carry {
            if *byte == 0xFF {
                *byte = 0x00;
            } else {
                *byte += 1;
                carry = false;
                break;
            }
        }
    }
    if carry {
        Bound::Unbounded
    } else {
        let mut end = K::zeroed();
        end.as_bytes_mut()[..incremented.len()].copy_from_slice(&incremented);
        Bound::Excluded(end)
    }
}

/// Iterator over entries in a `TypedTree`. Returned by `iter()`, `range()`, and `prefix_iter()`.
///
/// Weakly-consistent: concurrent inserts/updates may be visible during iteration.
/// Deleted entries are skipped. The `seize` guard prevents use-after-free.
pub struct TypedIter<'a, K: Key, T: Send + Sync> {
    list: &'a SkipList<TypedNode<K, T>>,
    front: *mut TypedNode<K, T>,
    /// `None` = not yet resolved (lazy). Computed on first `next_back()` call.
    back: Option<*mut TypedNode<K, T>>,
    end: Bound<K>,
    start: Bound<K>,
    reversed: bool,
    done: bool,
    _guard: seize::LocalGuard<'a>,
}

// SAFETY: TypedIter holds a seize guard that prevents node reclamation.
// The raw pointer is only dereferenced while the guard is alive.
unsafe impl<K: Key, T: Send + Sync> Send for TypedIter<'_, K, T> {}

impl<'a, K: Key, T: Send + Sync> Iterator for TypedIter<'a, K, T> {
    type Item = (K, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.done || self.front.is_null() {
                return None;
            }
            let node = unsafe { &*self.front };
            let converged = self.back.is_some_and(|back| std::ptr::eq(self.front, back));
            self.front = crate::skiplist::strip_mark(
                node.tower(0).load(std::sync::atomic::Ordering::Acquire),
            );
            if converged {
                self.done = true;
            }
            if node.is_marked() {
                if converged {
                    return None;
                }
                continue;
            }
            if !self.check_end(&node.key) {
                self.done = true;
                return None;
            }
            return Some((node.key, &node.load_data().value));
        }
    }
}

impl<'a, K: Key, T: Send + Sync> DoubleEndedIterator for TypedIter<'a, K, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.back.is_none() {
            self.back = Some(self.resolve_back());
            if self.front.is_null() {
                self.done = true;
            }
        }
        loop {
            let back = self.back.unwrap_or(std::ptr::null_mut());
            if self.done || back.is_null() {
                return None;
            }
            let node = unsafe { &*back };
            let key = node.key;
            let converged = std::ptr::eq(self.front, back);
            self.back = Some(self.list.find_last_lt(key.as_bytes()));
            if converged {
                self.done = true;
            }
            if node.is_marked() {
                if converged {
                    return None;
                }
                continue;
            }
            if !self.check_start(&key) {
                self.done = true;
                return None;
            }
            return Some((key, &node.load_data().value));
        }
    }
}

impl<K: Key, T: Send + Sync> TypedIter<'_, K, T> {
    /// Lazily resolve the back pointer for DoubleEndedIterator.
    fn resolve_back(&self) -> *mut TypedNode<K, T> {
        match &self.end {
            Bound::Unbounded => self.list.find_last(),
            Bound::Excluded(k) => self.list.find_last_lt(k.as_bytes()),
            Bound::Included(k) => {
                let ge = self.list.find_first_ge(k.as_bytes(), &self._guard);
                if !ge.is_null()
                    && !unsafe { &*ge }.is_marked()
                    && unsafe { &*ge }.key_bytes() == k.as_bytes()
                {
                    ge
                } else {
                    self.list.find_last_lt(k.as_bytes())
                }
            }
        }
    }

    /// Check if key is within the end bound (forward direction).
    #[inline(always)]
    fn check_end(&self, key: &K) -> bool {
        match &self.end {
            Bound::Unbounded => true,
            Bound::Excluded(end) => {
                if self.reversed {
                    key.as_bytes() > end.as_bytes()
                } else {
                    key.as_bytes() < end.as_bytes()
                }
            }
            Bound::Included(end) => {
                if self.reversed {
                    key.as_bytes() >= end.as_bytes()
                } else {
                    key.as_bytes() <= end.as_bytes()
                }
            }
        }
    }

    /// Check if key is within the start bound (backward direction).
    #[inline(always)]
    fn check_start(&self, key: &K) -> bool {
        match &self.start {
            Bound::Unbounded => true,
            Bound::Excluded(s) => {
                if self.reversed {
                    key.as_bytes() < s.as_bytes()
                } else {
                    key.as_bytes() > s.as_bytes()
                }
            }
            Bound::Included(s) => {
                if self.reversed {
                    key.as_bytes() <= s.as_bytes()
                } else {
                    key.as_bytes() >= s.as_bytes()
                }
            }
        }
    }
}

impl<'a, K: Key, T: Send + Sync> TypedIter<'a, K, T> {
    /// Collect all remaining entries. Convenience for backward compatibility.
    pub fn collect_vec(&mut self) -> Vec<(K, &'a T)> {
        self.collect()
    }
}

#[cfg(feature = "armour")]
impl<T, C, H> crate::armour::collection::Collection for TypedTree<T::SelfId, T, C, H>
where
    T: crate::CollectionMeta + Clone + Send + Sync,
    C: crate::Codec<T> + Sync,
    H: crate::hook::TypedWriteHook<T::SelfId, T>,
    T::SelfId: crate::Key + Ord,
{
    fn name(&self) -> &str {
        T::NAME
    }
    fn len(&self) -> usize {
        self.len()
    }
    fn compact(&self) -> crate::DbResult<usize> {
        self.compact()
    }
}