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::{self, Mutex, 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`.
///
/// `TypedRef` is `!Send` because it holds a `seize::LocalGuard` which is
/// thread-local. Use `.clone_value()` or copy the data out if you need to
/// send the value across threads.
///
/// ```compile_fail,E0277
/// fn assert_send<T: Send>() {}
/// // TypedRef must NOT be Send — LocalGuard is !Send.
/// assert_send::<armdb::TypedRef<'_, u64>>();
/// ```
pub struct TypedRef<'a, T> {
    _guard: seize::LocalGuard<'a>,
    data: *const TypedData<T>,
    _marker: std::marker::PhantomData<&'a 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 }
    }
}

impl<T: Clone> TypedRef<'_, T> {
    /// Clone the referenced value.
    pub fn clone_value(&self) -> T {
        (**self).clone()
    }
}

/// Deferred hook/retire work returned by _locked methods.
/// Caller drops shard lock, fires hook, then retires pointers.
struct MutationResult<T> {
    /// Raw pointer for TypedRef return. None if key was new / not found.
    ref_ptr: Option<*const TypedData<T>>,
    /// Hook args: (old_value_ptr, new_value_ptr). None = don't fire hook.
    hook: Option<(Option<*const T>, Option<*const T>)>,
    /// Pointer to retire via seize after hook. None = nothing to retire.
    retire: Option<*mut TypedData<T>>,
}

/// 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`/`compare_delete`/`update`
/// and on each mutation inside `atomic()` (replayed after the lock is released, in application order).
/// 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::balanced().build(),
///     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,
    encode_bufs: Vec<Mutex<Vec<u8>>>,
}

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_hooked(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> {
        let config = config.with_resolved_hints(false);
        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 shard_count = engine.shards().len();
        let encode_bufs = (0..shard_count).map(|_| Mutex::new(Vec::new())).collect();

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

        // 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 outcome = recover_typed_tree::<K, T, C>(
            &shard_dir_refs,
            &shard_ids,
            tree.index(),
            &tree.codec,
            hints,
            #[cfg(feature = "encryption")]
            tree.engine.cipher(),
        )?;
        for tail in &outcome.active_tails {
            tree.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
        }
        for (shard_idx, dead) in outcome.shard_dead_bytes {
            tree.engine.shards()[shard_idx].install_dead_bytes(dead);
        }
        let max_gsn = outcome.max_gsn;

        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)
    }

    fn is_live(&self, _shard_id: u8, key: &K, loc: DiskLoc) -> bool {
        let guard = self.index.collector().enter();
        self.index
            .get(key.as_bytes(), &guard)
            .is_some_and(|node| node.load_data().disk == loc)
    }
}

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.
    ///
    /// **O(log n)** — SkipList point lookup.
    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::new(guard, data))
    }

    /// 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 guard = self.index.collector().enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            self.put_locked(shard_id, &mut inner, &guard, key, value)?
        };
        self.fire_hook_and_retire(key, &mr);
        Ok(mr.ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    /// 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 guard = self.index.collector().enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            self.insert_locked::<true>(shard_id, &mut inner, &guard, key, value)?
        };
        self.fire_hook_and_retire(key, &mr);
        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 guard = self.index.collector().enter();
        let mr = {
            let mut inner = self.engine.shards()[shard_id].lock();
            self.delete_locked(shard_id, &mut inner, &guard, key)?
        };
        self.fire_hook_and_retire(key, &mr);
        Ok(mr.ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    /// 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,
            events: Vec::new(),
        };
        let result = f(&mut shard);
        let TypedShard {
            inner,
            guard,
            events,
            ..
        } = shard;
        drop(inner);
        if H::NEEDS_WRITE {
            for (key, mr) in &events {
                self.fire_hook_and_retire(key, mr);
            }
        }
        drop(guard);
        result
    }

    /// Fire deferred on_write and retire old data. Called after lock release.
    ///
    /// SAFETY: caller must hold a live `seize::LocalGuard` that was entered
    /// before the mutation. This guarantees old/new value pointers are valid
    /// (not yet reclaimed) even if a reentrant hook retires them.
    fn fire_hook_and_retire(&self, key: &K, mr: &MutationResult<T>) {
        if let Some((old, new)) = mr.hook {
            unsafe {
                self.hook.on_write(key, old.map(|p| &*p), new.map(|p| &*p));
            }
        }
        if let Some(ptr) = mr.retire {
            unsafe {
                self.index
                    .collector()
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
    }

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

    fn put_locked_inner<const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: &seize::LocalGuard<'_>,
        key: &K,
        value: T,
    ) -> DbResult<MutationResult<T>> {
        let mut buf = sync::lock(&self.encode_bufs[shard_id]);
        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
        if let Some(existing) = self.index.get(key.as_bytes(), guard) {
            let old_data = existing.load_data() as *const TypedData<T>;
            let new_data = Box::new(TypedData {
                disk: disk_loc,
                value,
            });
            let new_data_ptr = Box::into_raw(new_data);

            // Commit: swap pointer in index
            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,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            if HOOKS {
                return Ok(MutationResult {
                    ref_ptr: Some(old_data),
                    hook: Some((
                        Some(unsafe { &(*old_data).value as *const T }),
                        Some(unsafe { &(*new_data_ptr).value as *const T }),
                    )),
                    retire: Some(old_data_ptr),
                });
            }
            unsafe {
                self.index
                    .collector()
                    .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
            return Ok(MutationResult {
                ref_ptr: Some(old_data),
                hook: None,
                retire: None,
            });
        }

        // 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 => {
                let hook = if HOOKS {
                    let data = unsafe { &*node_ptr }.load_data();
                    Some((None, Some(&data.value as *const T)))
                } else {
                    None
                };
                Ok(MutationResult {
                    ref_ptr: None,
                    hook,
                    retire: None,
                })
            }
            InsertResult::Exists(existing) => {
                // Race: another thread 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,
                    crate::entry::entry_size(size_of::<K>(), old_disk.len),
                );
                // Dealloc the unused node (not the data — data was moved to existing)
                unsafe {
                    TypedNode::<K, T>::dealloc_node(node_ptr);
                }

                if HOOKS {
                    Ok(MutationResult {
                        ref_ptr: Some(old_data),
                        hook: Some((
                            Some(unsafe { &(*old_data).value as *const T }),
                            Some(unsafe { &(*new_data_ptr).value as *const T }),
                        )),
                        retire: Some(old_data_ptr),
                    })
                } else {
                    unsafe {
                        self.index
                            .collector()
                            .retire(old_data_ptr, seize::reclaim::boxed::<TypedData<T>>);
                    }
                    Ok(MutationResult {
                        ref_ptr: Some(old_data),
                        hook: None,
                        retire: None,
                    })
                }
            }
        }
    }

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

        let mut buf = sync::lock(&self.encode_bufs[shard_id]);
        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 hook = if HOOKS {
            let data = unsafe { &*node_ptr }.load_data();
            Some((None, Some(&data.value as *const T)))
        } else {
            None
        };
        Ok(MutationResult {
            ref_ptr: None,
            hook,
            retire: None,
        })
    }

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

    fn delete_locked_inner<const HOOKS: bool>(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        guard: &seize::LocalGuard<'_>,
        key: &K,
    ) -> DbResult<MutationResult<T>> {
        let node = match self.index.get(key.as_bytes(), guard) {
            Some(n) => n,
            None => {
                return Ok(MutationResult {
                    ref_ptr: None,
                    hook: None,
                    retire: None,
                });
            }
        };
        let data = node.load_data() as *const TypedData<T>;

        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,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );

        // Never retire data separately here — index.remove() already retired
        // the node, and TypedNode::reclaim will free the data inside it.
        // The guard keeps the node (and its data) alive until dropped.
        let hook = if HOOKS {
            Some((Some(unsafe { &(*data).value as *const T }), None))
        } else {
            None
        };
        Ok(MutationResult {
            ref_ptr: Some(data),
            hook,
            retire: None,
        })
    }

    /// 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 guard = self.index.collector().enter();

        let (hook_old, hook_new, retire_ptr) = {
            let shard = &self.engine.shards()[shard_id];
            let mut inner = shard.lock();

            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 = sync::lock(&self.encode_bufs[shard_id]);
            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,
            });

            let new_data_ptr = Box::into_raw(new_data);
            // Commit: swap
            let old_data = existing.swap_data(new_data_ptr);
            let old_disk = unsafe { (*old_data).disk };
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            (
                &current_data.value as *const T,
                unsafe { &(*new_data_ptr).value as *const T },
                old_data,
            )
        };

        unsafe {
            self.hook.on_write(key, Some(&*hook_old), Some(&*hook_new));
            self.index
                .collector()
                .retire(retire_ptr, seize::reclaim::boxed::<TypedData<T>>);
        }

        Ok(())
    }

    /// Compare-and-delete: if the current value == `expected`, delete the key.
    /// Returns `Ok(())` on success, `Err(CasMismatch)` if current != expected,
    /// `Err(KeyNotFound)` if the key doesn't exist.
    pub fn compare_delete(&self, key: &K, expected: &T) -> DbResult<()>
    where
        T: PartialEq,
    {
        metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "typed_tree")
            .increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("typed_tree.compare_delete");
        let shard_id = self.shard_for(key);
        let guard = self.index.collector().enter();

        let mr = {
            let shard = &self.engine.shards()[shard_id];
            let mut inner = shard.lock();

            let existing = self
                .index
                .get(key.as_bytes(), &guard)
                .ok_or(DbError::KeyNotFound)?;
            if existing.load_data().value != *expected {
                return Err(DbError::CasMismatch);
            }

            self.delete_locked_inner::<true>(shard_id, &mut inner, &guard, key)?
        };

        self.fire_hook_and_retire(key, &mr);
        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 guard = self.index.collector().enter();

        let result = {
            let shard = &self.engine.shards()[shard_id];
            let mut inner = shard.lock();

            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 = sync::lock(&self.encode_bufs[shard_id]);
            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,
            });

            let new_data_ptr = Box::into_raw(new_data);
            // Commit: swap
            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,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );

            (old_data as *const TypedData<T>, new_data_ptr, old_data_ptr)
        };

        let (old_data, new_data_ptr, retire_ptr) = result;
        unsafe {
            self.hook
                .on_write(key, Some(&(*old_data).value), Some(&(*new_data_ptr).value));
            self.index
                .collector()
                .retire(retire_ptr, seize::reclaim::boxed::<TypedData<T>>);
        }

        let ref_ptr = if return_old {
            old_data
        } else {
            new_data_ptr as *const TypedData<T>
        };
        Ok(Some(TypedRef::new(guard, ref_ptr)))
    }

    /// 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::new(guard, data)));
            }
            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 mut ptr = self.index.find_last(&guard);
        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::new(guard, data)));
            }
            ptr = self.index.find_last_lt(node.key_bytes(), &guard);
        }
        None
    }

    // -- 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]) -> Option<(K, Bound<K>)> {
        if prefix.len() > size_of::<K>() {
            return None;
        }
        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);
            Some((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);
            Some((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 Some((search_key, end)) = self.prefix_bounds(prefix) else {
            return TypedIter {
                list: &self.index,
                front: std::ptr::null_mut(),
                back: None,
                end: Bound::Unbounded,
                start: Bound::Unbounded,
                reversed: self.reversed,
                done: true,
                _guard: guard,
            };
        };
        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.
    ///
    /// **Must run before any other writer obtains a handle to this collection.**
    /// Concurrent writes during migration produce undefined results.
    ///
    /// 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) => {
                    let shard_id = self.shard_for(&node.key);
                    let init_ptr = {
                        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)?;
                        H::NEEDS_INIT.then(|| node.load_data() as *const TypedData<T>)
                    };
                    if let Some(ptr) = init_ptr {
                        let data = unsafe { &*ptr };
                        self.hook.on_init(&node.key, &data.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>,
    events: Vec<(K, MutationResult<T>)>,
}

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 mr = self.tree.put_locked_inner::<true>(
            self.shard_id,
            &mut self.inner,
            &self.guard,
            key,
            value,
        )?;
        let ref_ptr = mr.ref_ptr;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.tree
                    .index
                    .collector()
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        let guard = self.tree.index.collector().enter();
        Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
        self.check_shard(key)?;
        let mr = self.tree.insert_locked::<true>(
            self.shard_id,
            &mut self.inner,
            &self.guard,
            key,
            value,
        )?;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.tree
                    .index
                    .collector()
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
        self.check_shard(key)?;
        let mr = self.tree.delete_locked_inner::<true>(
            self.shard_id,
            &mut self.inner,
            &self.guard,
            key,
        )?;
        let ref_ptr = mr.ref_ptr;
        if H::NEEDS_WRITE {
            self.events.push((*key, mr));
        } else if let Some(ptr) = mr.retire {
            unsafe {
                self.tree
                    .index
                    .collector()
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
        let guard = self.tree.index.collector().enter();
        Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
    }

    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>,
}

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(), &self._guard));
            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(&self._guard),
            Bound::Excluded(k) => self.list.find_last_lt(k.as_bytes(), &self._guard),
            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(), &self._guard)
                }
            }
        }
    }

    /// 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()
    }
    fn flush(&self) -> crate::DbResult<()> {
        self.flush_buffers()?;
        self.sync_hints()?;
        Ok(())
    }
    fn periodic_flush(&self) -> crate::DbResult<()> {
        self.flush_buffers()
    }
}

// ==========================================================================
// MultiTx — cross-collection transaction support (feature `armour`)
//
// Engine-direct family: locks `engine.shards()[id]` per shard, holds one seize
// guard, has NO durability sync seam (release_locks → SyncNeeds::none()), and
// replays via fire_hook_and_retire — which also retires deferred old values, so
// it must run for every logged mutation (mirrors the single-shard atomic()).
// ==========================================================================

/// Multi-shard transaction handle for [`TypedTree`] inside `Db::atomicN`.
#[cfg(feature = "armour")]
pub struct TypedTx<'a, K: Key, T: Send + Sync, C: Codec<T>, H: TypedWriteHook<K, T> = NoHook> {
    tree: &'a TypedTree<K, T, C, H>,
    inners: Vec<(usize, MutexGuard<'a, ShardInner>)>,
    seize: seize::LocalGuard<'a>,
    log: Vec<(K, MutationResult<T>)>,
}

#[cfg(feature = "armour")]
impl<'a, K: Key, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    TypedTx<'a, K, T, C, H>
{
    fn position(&self, key: &K) -> DbResult<usize> {
        let sid = self.tree.shard_for(key);
        self.inners
            .iter()
            .position(|(s, _)| *s == sid)
            .ok_or(DbError::ShardMismatch)
    }

    pub fn try_get(&self, key: &K) -> DbResult<Option<&T>> {
        self.position(key)?;
        Ok(self
            .tree
            .index
            .get(key.as_bytes(), &self.seize)
            .map(|n| &n.load_data().value))
    }

    pub fn try_contains(&self, key: &K) -> DbResult<bool> {
        self.position(key)?;
        Ok(self.tree.index.get(key.as_bytes(), &self.seize).is_some())
    }

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

    /// Retire a deferred old value immediately when no hook will replay it.
    fn retire_if_no_hook(&self, mr: &MutationResult<T>) {
        if let Some(ptr) = mr.retire {
            unsafe {
                self.tree
                    .index
                    .collector()
                    .retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
            }
        }
    }

    pub fn put(&mut self, key: &K, value: T) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner) = &mut self.inners[i];
        let mr = self
            .tree
            .put_locked_inner::<true>(*sid, inner, &self.seize, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }

    pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner) = &mut self.inners[i];
        let mr = self
            .tree
            .insert_locked::<true>(*sid, inner, &self.seize, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner) = &mut self.inners[i];
        let mr = self
            .tree
            .delete_locked_inner::<true>(*sid, inner, &self.seize, key)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, mr));
        } else {
            self.retire_if_no_hook(&mr);
        }
        Ok(())
    }
}

#[cfg(feature = "armour")]
impl<K: Key, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
    crate::armour::multi_tx::MultiTx for TypedTree<K, T, C, H>
{
    type Key = K;
    type Tx<'a>
        = TypedTx<'a, K, T, C, H>
    where
        Self: 'a;

    fn shard_for_key(&self, key: &K) -> usize {
        self.shard_for(key)
    }

    fn begin_tx(&self) -> TypedTx<'_, K, T, C, H> {
        TypedTx {
            tree: self,
            inners: Vec::new(),
            seize: self.index.collector().enter(),
            log: Vec::new(),
        }
    }

    fn lock_shard_into<'a>(&'a self, shard_id: usize, tx: &mut TypedTx<'a, K, T, C, H>) {
        tx.inners
            .push((shard_id, self.engine.shards()[shard_id].lock()));
    }

    fn release_locks(
        &self,
        tx: &mut TypedTx<'_, K, T, C, H>,
    ) -> crate::armour::multi_tx::SyncNeeds {
        tx.inners.clear(); // drops engine shard locks; no per-write sync seam
        crate::armour::multi_tx::SyncNeeds::none()
    }

    fn run_sync(&self, _needs: crate::armour::multi_tx::SyncNeeds) -> DbResult<()> {
        Ok(())
    }

    fn replay_hooks(&self, tx: TypedTx<'_, K, T, C, H>) {
        if H::NEEDS_WRITE {
            for (k, mr) in &tx.log {
                self.fire_hook_and_retire(k, mr);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::Codec;
    use crate::config::Config;
    use crate::hook::TypedWriteHook;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};

    struct U64Codec;

    impl Codec<u64> for U64Codec {
        fn encode_to(&self, value: &u64, buf: &mut Vec<u8>) -> DbResult<()> {
            buf.clear();
            buf.extend_from_slice(&value.to_be_bytes());
            Ok(())
        }

        fn decode_from(&self, bytes: &[u8]) -> DbResult<u64> {
            let arr: [u8; 8] = bytes
                .try_into()
                .map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
            Ok(u64::from_be_bytes(arr))
        }
    }

    #[derive(Default)]
    struct TRecHook {
        writes: AtomicUsize,
        #[allow(clippy::type_complexity)]
        seq: crate::sync::Mutex<Vec<(u64, Option<u64>, Option<u64>)>>,
    }

    impl TypedWriteHook<[u8; 8], u64> for TRecHook {
        fn on_write(&self, key: &[u8; 8], old: Option<&u64>, new: Option<&u64>) {
            self.writes.fetch_add(1, AtomicOrdering::Relaxed);
            crate::sync::lock(&self.seq).push((
                u64::from_be_bytes(*key),
                old.copied(),
                new.copied(),
            ));
        }
    }

    fn open_typed_hooked<H: TypedWriteHook<[u8; 8], u64>>(
        dir: &std::path::Path,
        cfg: Config,
        hook: H,
    ) -> TypedTree<[u8; 8], u64, U64Codec, H> {
        TypedTree::open_hooked(dir, cfg, U64Codec, hook).expect("open hooked")
    }

    #[test]
    fn typed_tree_atomic_fires_hooks_in_order() {
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        let tree = open_typed_hooked(dir.path(), cfg, TRecHook::default());
        let k = 5u64.to_be_bytes();
        tree.atomic(&k, |s| {
            s.put(&k, 10)?; // new key
            s.put(&k, 20)?; // update, old=10
            s.delete(&k)?; // delete, old=20
            Ok(())
        })
        .expect("atomic");
        assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 3);
        let seq = crate::sync::lock(&tree.hook.seq).clone();
        assert_eq!(seq[0], (5, None, Some(10)));
        assert_eq!(seq[1], (5, Some(10), Some(20)));
        assert_eq!(seq[2], (5, Some(20), None));
    }

    #[test]
    fn typed_tree_atomic_fires_for_applied_on_err() {
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        let tree = open_typed_hooked(dir.path(), cfg, TRecHook::default());
        let k = 1u64.to_be_bytes();
        let r: DbResult<()> = tree.atomic(&k, |s| {
            s.put(&k, 99)?;
            Err(DbError::KeyNotFound)
        });
        assert!(r.is_err());
        assert_eq!(tree.hook.writes.load(AtomicOrdering::Relaxed), 1);
    }
}
