armdb 0.6.0

use std::marker::PhantomData;
use std::ops::Bound;

use crate::Key;
use crate::batch::{Applied, BatchWrite};
use crate::byte_view::ByteView;
use crate::codec::Codec;
use crate::compaction::CompactionIndex;
use crate::config::Config;
use crate::disk_loc::DiskLoc;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, TypedWriteHook, VarTypedHookAdapter};
use crate::var_tree::{VarIter, VarShard, VarTree};

/// A tree with fixed-size keys and typed values `T`. Values are encoded via
/// a [`Codec`] and stored on disk (variable length), with a block cache
/// (≤ 8 KB values) and a value cache (> 8 KB values) for reads — unlike
/// [`TypedTree`](crate::TypedTree) which keeps values in memory.
///
/// Thin wrapper around [`VarTree<K, VarTypedHookAdapter<K, T, C, H>>`].
/// Each `VarTypedTree` owns its storage engine — one tree = one database directory.
///
/// # When to use
///
/// - Typed values too large to keep in RAM → `VarTypedTree` (disk-resident).
/// - Typed values that fit in RAM → [`TypedTree`](crate::TypedTree) (in-memory).
/// - Fixed-size zerocopy values → [`ZeroTree`](crate::ZeroTree).
///
/// # Error handling
///
/// Follows `VarTree` / `TypedTree` precedent (variant A):
/// - `get` / `first` / `last` → `None` if key missing **or** decode fails.
/// - Iterators skip entries with decode errors (`continue` in `next()`).
/// - `migrate` treats decode errors as `Keep` (logged via `tracing::warn!`).
///
/// # Write hooks
///
/// Uses [`TypedWriteHook<K, T>`] via [`VarTypedHookAdapter`]. The hook receives
/// `&T` directly; the adapter decodes raw bytes via the codec. `on_write` fires
/// on `put`/`insert`/`delete`/`cas`/`compare_delete`/`update`/`fetch_update` and inside `atomic()`.
///
/// # Usage
///
/// ```ignore
/// let tree = VarTypedTree::<[u8; 16], Message, RapiraCodec>::open(
///     "data/messages",
///     Config::large_values().build(),
///     RapiraCodec,
/// )?;
/// tree.put(&key, &msg)?;
/// if let Some(m) = tree.get(&key) {
///     println!("{:?}", m);
/// }
/// tree.close()?;
/// ```
///
/// # Iteration
///
/// `iter()`, `range()`, and `prefix_iter()` return [`VarTypedIter`] with
/// `Item = (K, T)`. Each `next()` / `next_back()` may perform disk I/O on a
/// block-cache miss and performs a decode.
///
/// # Batch operations
///
/// [`get_many`](Self::get_many) does one sorted finger-walk over the index for a
/// list of keys; [`update_many`](Self::update_many) applies a per-key closure
/// grouped by shard (per-shard isolation, no global atomicity). See
/// `docs/superpowers/specs/26-06-24-armdb-batch-tree-ops.md`.
pub struct VarTypedTree<
    K: Key,
    T: Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T> = NoHook,
> {
    inner: VarTree<K, VarTypedHookAdapter<K, T, C, H>>,
    codec: C,
    _marker: PhantomData<fn() -> T>,
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone> VarTypedTree<K, T, C> {
    /// Open or create a `VarTypedTree` 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_inner(path, config, codec, NoHook)
    }
}

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

    fn open_hooked_inner(
        path: impl AsRef<std::path::Path>,
        config: Config,
        codec: C,
        hook: H,
    ) -> DbResult<Self> {
        let adapter = VarTypedHookAdapter {
            inner: hook,
            codec: codec.clone(),
            _marker: PhantomData,
        };
        let inner = VarTree::open_hooked(path, config, adapter)?;
        Ok(Self {
            inner,
            codec,
            _marker: PhantomData,
        })
    }

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

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

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

    /// Trigger a compaction pass across all shards.
    pub fn compact(&self) -> DbResult<usize> {
        self.inner.compact()
    }

    /// Write hint files for all active shard files. Call during graceful shutdown.
    pub fn sync_hints(&self) -> DbResult<()> {
        self.inner.sync_hints()
    }

    /// Pre-populate the block cache with blocks containing live values.
    pub fn warmup(&self) -> DbResult<()> {
        self.inner.warmup()
    }

    /// Access the underlying `VarTree` for raw byte operations.
    pub fn as_inner(&self) -> &VarTree<K, VarTypedHookAdapter<K, T, C, H>> {
        &self.inner
    }

    /// Access the codec used for encoding / decoding values.
    pub fn codec(&self) -> &C {
        &self.codec
    }

    // -- Reads ----------------------------------------------------------------

    /// Get and decode a value by key.
    /// Returns `None` if the key is missing or decoding fails.
    ///
    /// **O(log n)** — lock-free SkipList lookup; brief shard lock + `pread` on block-cache
    /// miss, then decode.
    pub fn get(&self, key: &K) -> Option<T> {
        let bytes = self.inner.get(key)?;
        self.codec.decode_from(&bytes).ok()
    }

    /// Get a value by key, returning `Err(KeyNotFound)` if absent or
    /// `Err(CorruptedEntry)` if present but undecodable.
    pub fn get_or_err(&self, key: &K) -> DbResult<T> {
        self.try_get(key)?.ok_or(DbError::KeyNotFound)
    }

    /// Strict read: `Ok(None)` only when absent; `Err` when present but the value
    /// cannot be read or decoded.
    pub fn try_get(&self, key: &K) -> DbResult<Option<T>> {
        match self.inner.try_get(key)? {
            Some(bytes) => self.codec.decode_from(&bytes).map(Some),
            None => Ok(None),
        }
    }

    /// Check if a key exists.
    pub fn contains(&self, key: &K) -> bool {
        self.inner.contains(key)
    }

    /// Return the first entry in index order, or `None` if empty or every boundary
    /// entry is unproducible. Skips unreadable/undecodable entries, matching
    /// `iter().next()`.
    pub fn first(&self) -> Option<(K, T)> {
        self.iter().next()
    }

    /// Return the last entry in index order (symmetric to `first`).
    pub fn last(&self) -> Option<(K, T)> {
        self.iter().next_back()
    }

    /// Strict first entry: `Ok(None)` only when there is no live entry; `Err` when
    /// the boundary entry cannot be read or decoded.
    pub fn try_first(&self) -> DbResult<Option<(K, T)>> {
        match self.inner.try_first()? {
            Some((k, bytes)) => Ok(Some((k, self.codec.decode_from(&bytes)?))),
            None => Ok(None),
        }
    }

    /// Strict last entry (symmetric to `try_first`).
    pub fn try_last(&self) -> DbResult<Option<(K, T)>> {
        match self.inner.try_last()? {
            Some((k, bytes)) => Ok(Some((k, self.codec.decode_from(&bytes)?))),
            None => Ok(None),
        }
    }

    // -- Writes ---------------------------------------------------------------

    /// Insert or update a key-value pair.
    ///
    /// Returns `true` if a previous value for `key` existed (overwrite),
    /// `false` for a fresh insert.
    pub fn put(&self, key: &K, value: &T) -> DbResult<bool> {
        let mut buf = Vec::new();
        self.codec.encode_to(value, &mut buf)?;
        self.inner.put(key, &buf)
    }

    /// 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<()> {
        let mut buf = Vec::new();
        self.codec.encode_to(value, &mut buf)?;
        self.inner.insert(key, &buf)
    }

    /// Delete a key. Returns `true` if the key existed.
    pub fn delete(&self, key: &K) -> DbResult<bool> {
        self.inner.delete(key)
    }

    /// Compare-and-swap based on encoded bytes. Relies on deterministic codec output.
    /// Returns `Ok(())` on success, `Err(CasMismatch)` if current != expected,
    /// `Err(KeyNotFound)` if key doesn't exist.
    ///
    /// **Caveat:** inherits `VarTree::cas` behavior — holds the shard lock while
    /// reading the current value; block-cache miss causes disk I/O under the lock.
    pub fn cas(&self, key: &K, expected: &T, new_value: &T) -> DbResult<()> {
        let mut exp_buf = Vec::new();
        self.codec.encode_to(expected, &mut exp_buf)?;
        let mut new_buf = Vec::new();
        self.codec.encode_to(new_value, &mut new_buf)?;
        self.inner.cas(key, &exp_buf, &new_buf)
    }

    /// Compare-and-delete based on encoded bytes. Relies on deterministic codec output.
    /// 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<()> {
        let mut exp_buf = Vec::new();
        self.codec.encode_to(expected, &mut exp_buf)?;
        self.inner.compare_delete(key, &exp_buf)
    }

    /// Atomically read-modify-write. Returns `Some(new_value)` if key existed,
    /// `None` if the key is absent; `Err(CorruptedEntry)` if the current value
    /// cannot be decoded.
    ///
    /// **Caveat:** inherits `VarTree::update` behavior — shard lock held during
    /// possible disk I/O.
    pub fn update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
        use std::cell::Cell;
        let out: Cell<Option<T>> = Cell::new(None);
        let result = self.inner.try_update_inner(
            key,
            |bytes| {
                let current = self.codec.decode_from(bytes)?;
                let new_val = f(&current);
                let mut buf = Vec::new();
                self.codec.encode_to(&new_val, &mut buf)?;
                out.set(Some(new_val));
                Ok(Some(ByteView::from_vec(buf)))
            },
            false,
        )?;
        if result.is_none() {
            return Ok(None);
        }
        Ok(out.into_inner())
    }

    /// Like [`update()`](Self::update), but returns `Some(old_value)` instead
    /// of the new one.
    pub fn fetch_update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
        use std::cell::Cell;
        let out: Cell<Option<T>> = Cell::new(None);
        let result = self.inner.try_update_inner(
            key,
            |bytes| {
                let current = self.codec.decode_from(bytes)?;
                let new_val = f(&current);
                let mut buf = Vec::new();
                self.codec.encode_to(&new_val, &mut buf)?;
                out.set(Some(current));
                Ok(Some(ByteView::from_vec(buf)))
            },
            true,
        )?;
        if result.is_none() {
            return Ok(None);
        }
        Ok(out.into_inner())
    }

    // -- Atomic ---------------------------------------------------------------

    /// 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 VarTypedShard<'_, K, T, C, H>) -> DbResult<R>,
    ) -> DbResult<R> {
        self.inner.atomic(shard_key, |var_shard| {
            // SAFETY: erase VarShard lifetime via `*mut ()` — see VarTypedShard doc.
            // The pointer is valid for the duration of this closure.
            let inner_ptr: *mut () = (var_shard as *mut VarShard<'_, _, _>).cast();
            let mut shard = VarTypedShard {
                tree: self,
                inner: inner_ptr,
                _marker: PhantomData,
            };
            f(&mut shard)
        })
    }

    // -- Iteration ------------------------------------------------------------

    /// 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). Disk I/O on cache miss.
    /// Entries with decode errors are skipped.
    pub fn prefix_iter(&self, prefix: &[u8]) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.prefix_iter(prefix),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

    /// Iterate all entries in index order. Disk I/O on cache miss.
    /// `next()` is O(1), `next_back()` is O(log n).
    /// Entries with decode errors are skipped.
    pub fn iter(&self) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.iter(),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    /// `next()` is O(1), `next_back()` is O(log n). Disk I/O on cache miss.
    pub fn range(&self, start: &K, end: &K) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.range(start, end),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    /// `next()` is O(1), `next_back()` is O(log n). Disk I/O on cache miss.
    pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.range_bounds(start, end),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

    // -- Info -----------------------------------------------------------------

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

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

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

    /// Encoded value byte length for `key`, or `None` if absent.
    /// Reads only the in-memory index entry (`DiskLoc::len`); no disk I/O.
    pub fn entry_len(&self, key: &K) -> Option<u32> {
        self.inner.entry_len(key)
    }

    // -- Batch ----------------------------------------------------------------

    /// Batch point-read returning decoded values, aligned to `keys`. **Lossy** —
    /// like single-key [`get`](Self::get), a decode failure yields `None`.
    pub fn get_many(&self, keys: &[K]) -> Vec<Option<T>> {
        self.inner
            .get_many(keys)
            .into_iter()
            .map(|opt| opt.and_then(|b| self.codec.decode_from(&b).ok()))
            .collect()
    }

    /// Batch read-modify-write over decoded values. See [`crate::ConstTree::update_many`].
    ///
    /// **Strict on codec failures**, consistent with single-key [`update`](Self::update):
    /// a codec decode of an old value or encode of a new value that fails aborts the
    /// batch with that `DbError` (the applied prefix stays committed, as for any
    /// mid-batch error). Codec errors are never silently collapsed into
    /// `Keep`/`Kept`. (Read-only [`get_many`](Self::get_many) stays lossy, mirroring
    /// single-key [`get`](Self::get).)
    pub fn update_many<P>(
        &self,
        items: Vec<(K, P)>,
        f: impl Fn(&K, Option<&T>, &P) -> BatchWrite<T>,
    ) -> DbResult<Vec<(K, Applied<T>)>> {
        use std::cell::RefCell;

        // The inner (raw-bytes) closure cannot return a `Result`, so a codec
        // failure is captured here and surfaced after the batch unwinds. Once an
        // error is recorded, remaining items short-circuit to `Keep` so the batch
        // stops mutating; the recorded error is then returned.
        let codec_err: RefCell<Option<DbError>> = RefCell::new(None);
        let inner_out = self.inner.update_many(items, |k, old_bytes, p| {
            if codec_err.borrow().is_some() {
                return BatchWrite::Keep;
            }
            let old_t: Option<T> = match old_bytes {
                Some(b) => match self.codec.decode_from(b) {
                    Ok(t) => Some(t),
                    Err(e) => {
                        *codec_err.borrow_mut() = Some(e);
                        return BatchWrite::Keep;
                    }
                },
                None => None,
            };
            match f(k, old_t.as_ref(), p) {
                BatchWrite::Set(v) => {
                    let mut buf = Vec::new();
                    match self.codec.encode_to(&v, &mut buf) {
                        Ok(()) => BatchWrite::Set(ByteView::from_vec(buf)),
                        Err(e) => {
                            *codec_err.borrow_mut() = Some(e);
                            BatchWrite::Keep
                        }
                    }
                }
                BatchWrite::Keep => BatchWrite::Keep,
                BatchWrite::Delete => BatchWrite::Delete,
            }
        })?;
        if let Some(e) = codec_err.into_inner() {
            return Err(e);
        }
        // On the success path every old value already decoded inside the closure
        // and every `new` was just encoded here, so these decodes succeed; a
        // failure would be genuine corruption and is propagated, not swallowed.
        let mut out = Vec::with_capacity(inner_out.len());
        for (k, applied) in inner_out {
            let applied_t = match applied {
                Applied::Written { old, new } => Applied::Written {
                    old: match old {
                        Some(b) => Some(self.codec.decode_from(&b)?),
                        None => None,
                    },
                    new: self.codec.decode_from(&new)?,
                },
                Applied::Kept => Applied::Kept,
                Applied::Deleted(b) => Applied::Deleted(self.codec.decode_from(&b)?),
            };
            out.push((k, applied_t));
        }
        Ok(out)
    }

    // -- Migration ------------------------------------------------------------

    /// Iterate all entries and optionally mutate them. Call once at startup.
    ///
    /// The callback receives each `(key, &T)` and returns `MigrateAction`:
    /// - `Keep` — no change
    /// - `Update(new_value)` — replace value (re-encoded via codec)
    /// - `Delete` — remove entry
    ///
    /// Entries that fail to decode are logged and left unchanged (`Keep`).
    /// Returns the number of mutated entries.
    pub fn migrate(&self, f: impl Fn(&K, &T) -> crate::MigrateAction<T>) -> DbResult<usize> {
        use crate::MigrateAction;
        self.inner
            .migrate(|key, bytes| match self.codec.decode_from(bytes) {
                Ok(current) => match f(key, &current) {
                    MigrateAction::Keep => MigrateAction::Keep,
                    MigrateAction::Update(new) => {
                        let mut buf = Vec::new();
                        match self.codec.encode_to(&new, &mut buf) {
                            Ok(()) => MigrateAction::Update(ByteView::from_vec(buf)),
                            Err(_) => {
                                tracing::warn!(
                                    "var_typed_tree migrate: encode error, keeping entry"
                                );
                                MigrateAction::Keep
                            }
                        }
                    }
                    MigrateAction::Delete => MigrateAction::Delete,
                },
                Err(_) => {
                    tracing::warn!("var_typed_tree migrate: decode error, keeping entry");
                    MigrateAction::Keep
                }
            })
    }

    /// Replay `on_init` for every live entry. Used by Db when no migration ran.
    pub(crate) fn replay_init(&self) {
        self.inner.replay_init();
    }
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> CompactionIndex<K>
    for VarTypedTree<K, T, C, H>
{
    fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
        self.inner.update_if_match(key, old_loc, new_loc)
    }

    fn invalidate_blocks(&self, shard_id: u8, file_id: u32, total_bytes: u64) {
        self.inner.invalidate_blocks(shard_id, file_id, total_bytes);
    }

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

    fn is_live(&self, shard_id: u8, key: &K, loc: DiskLoc) -> bool {
        self.inner.is_live(shard_id, key, loc)
    }
}

#[cfg(feature = "replication")]
impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>>
    crate::replication::ReplicationTarget for VarTypedTree<K, T, C, H>
{
    fn apply_entry(
        &self,
        shard_inner: &mut crate::shard::ShardInner,
        shard_id: u8,
        file_id: u32,
        entry_offset: u64,
        header: &crate::entry::EntryHeader,
        key: &[u8],
        value: &[u8],
    ) -> DbResult<crate::replication::ApplyOutcome> {
        self.inner.apply_entry(
            shard_inner,
            shard_id,
            file_id,
            entry_offset,
            header,
            key,
            value,
        )
    }

    fn try_apply_entry(
        &self,
        shard_inner: &mut crate::shard::ShardInner,
        shard_id: u8,
        file_id: u32,
        entry_offset: u64,
        header: &crate::entry::EntryHeader,
        raw_after_header: &[u8],
    ) -> DbResult<crate::replication::ApplyOutcome> {
        self.inner.try_apply_entry(
            shard_inner,
            shard_id,
            file_id,
            entry_offset,
            header,
            raw_after_header,
        )
    }

    fn key_len(&self) -> usize {
        self.inner.key_len()
    }
}

// ---------------------------------------------------------------------------
// VarTypedIter
// ---------------------------------------------------------------------------

/// Iterator over entries in a [`VarTypedTree`]. Returned by `iter()`, `range()`,
/// `range_bounds()`, and `prefix_iter()`. Wraps [`VarIter`] and decodes each
/// value via the codec. Entries with decode errors are skipped (see `tracing::debug!`).
pub struct VarTypedIter<'a, K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> {
    inner: VarIter<'a, K, VarTypedHookAdapter<K, T, C, H>>,
    codec: &'a C,
    _marker: PhantomData<fn() -> T>,
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> Iterator
    for VarTypedIter<'_, K, T, C, H>
{
    type Item = (K, T);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let (k, bytes) = self.inner.next()?;
            match self.codec.decode_from(&bytes) {
                Ok(v) => return Some((k, v)),
                Err(_) => {
                    tracing::debug!(
                        value_len = bytes.len(),
                        "var_typed_iter: decode error, skipping entry"
                    );
                }
            }
        }
    }
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> DoubleEndedIterator
    for VarTypedIter<'_, K, T, C, H>
{
    fn next_back(&mut self) -> Option<Self::Item> {
        while let Some((k, bytes)) = self.inner.next_back() {
            match self.codec.decode_from(&bytes) {
                Ok(v) => return Some((k, v)),
                Err(_) => {
                    tracing::debug!(
                        value_len = bytes.len(),
                        "var_typed_iter: decode error, skipping entry"
                    );
                }
            }
        }
        None
    }
}

// ---------------------------------------------------------------------------
// VarTypedShard
// ---------------------------------------------------------------------------

/// Handle for atomic multi-key operations within a single shard.
/// Obtained via [`VarTypedTree::atomic`]. The shard lock is held for the
/// lifetime of this struct — keep the closure short.
pub struct VarTypedShard<
    'tree,
    K: Key,
    T: Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T>,
> {
    tree: &'tree VarTypedTree<K, T, C, H>,
    // Type-erased raw pointer to a `VarShard<'_, K, VarTypedHookAdapter<K, T, C, H>>`.
    // Invariance of the inner `'_` lifetime would otherwise force it to `'static`
    // when stored. Only dereferenced inside the enclosing `atomic()` closure.
    inner: *mut (),
    _marker: PhantomData<&'tree mut ()>,
}

// SAFETY: raw pointer dereferenced only while the enclosing `atomic` closure
// is active; `VarShard` itself is `Send` and guards a mutex. `tree` is `Send+Sync`.
unsafe impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> Send
    for VarTypedShard<'_, K, T, C, H>
{
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>>
    VarTypedShard<'_, K, T, C, H>
{
    fn inner_mut(&mut self) -> &mut VarShard<'_, K, VarTypedHookAdapter<K, T, C, H>> {
        // SAFETY: pointer set in `VarTypedTree::atomic` from a valid `&mut VarShard`,
        // only dereferenced while the closure is live.
        unsafe { &mut *(self.inner as *mut VarShard<'_, K, VarTypedHookAdapter<K, T, C, H>>) }
    }

    fn inner_ref(&self) -> &VarShard<'_, K, VarTypedHookAdapter<K, T, C, H>> {
        // SAFETY: see `inner_mut`.
        unsafe { &*(self.inner as *const VarShard<'_, K, VarTypedHookAdapter<K, T, C, H>>) }
    }

    /// Insert or overwrite `key`. Returns `true` if the key already existed (overwrite),
    /// `false` for a fresh insert.
    pub fn put(&mut self, key: &K, value: &T) -> DbResult<bool> {
        let mut buf = Vec::new();
        self.tree.codec.encode_to(value, &mut buf)?;
        self.inner_mut().put(key, &buf)
    }

    pub fn insert(&mut self, key: &K, value: &T) -> DbResult<()> {
        let mut buf = Vec::new();
        self.tree.codec.encode_to(value, &mut buf)?;
        self.inner_mut().insert(key, &buf)
    }

    pub fn delete(&mut self, key: &K) -> DbResult<bool> {
        self.inner_mut().delete(key)
    }

    pub fn get(&self, key: &K) -> Option<T> {
        let bytes = self.inner_ref().get(key)?;
        self.tree.codec.decode_from(&bytes).ok()
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<T> {
        let bytes = self.inner_ref().get_or_err(key)?;
        self.tree.codec.decode_from(&bytes)
    }

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

    /// Atomically read-modify-write within the shard. Returns `Some(new_value)`
    /// if the key existed, `None` if absent.
    pub fn update(&mut self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
        use std::cell::Cell;
        let tree = self.tree; // &'tree VarTypedTree — Copy, avoids borrowing self in closure
        let out: Cell<Option<T>> = Cell::new(None);
        let result = self.inner_mut().update_inner(
            key,
            |bytes| {
                let current = tree.codec.decode_from(bytes)?;
                let new_val = f(&current);
                let mut buf = Vec::new();
                tree.codec.encode_to(&new_val, &mut buf)?;
                out.set(Some(new_val));
                Ok(Some(ByteView::from_vec(buf)))
            },
            false,
        )?;
        if result.is_none() {
            return Ok(None);
        }
        Ok(out.into_inner())
    }

    /// Like [`update`](Self::update), but returns `Some(old_value)` instead of
    /// the new one.
    pub fn fetch_update(&mut self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<T>> {
        use std::cell::Cell;
        let tree = self.tree; // Copy
        let out: Cell<Option<T>> = Cell::new(None);
        let result = self.inner_mut().update_inner(
            key,
            |bytes| {
                let current = tree.codec.decode_from(bytes)?;
                let new_val = f(&current);
                let mut buf = Vec::new();
                tree.codec.encode_to(&new_val, &mut buf)?;
                out.set(Some(current)); // capture the OLD value
                Ok(Some(ByteView::from_vec(buf)))
            },
            true,
        )?;
        if result.is_none() {
            return Ok(None);
        }
        Ok(out.into_inner())
    }
}

#[cfg(feature = "armour")]
impl<T, C, H> crate::armour::collection::Collection for VarTypedTree<T::SelfId, T, C, H>
where
    T: crate::CollectionMeta + Send + Sync,
    C: Codec<T> + Clone + 'static,
    H: TypedWriteHook<T::SelfId, T>,
    T::SelfId: crate::Key + Ord,
{
    fn name(&self) -> &str {
        T::NAME
    }
    fn len(&self) -> usize {
        self.len()
    }
    fn compact(&self) -> DbResult<usize> {
        self.compact()
    }
    fn flush(&self) -> DbResult<()> {
        self.flush_buffers()?;
        self.sync_hints()?;
        Ok(())
    }
    fn periodic_flush(&self) -> DbResult<()> {
        self.flush_buffers()
    }
}

// ==========================================================================
// MultiTx — cross-collection transaction support (feature `armour`)
//
// VarTypedTree wraps a VarTree whose hook is VarTypedHookAdapter; the multi-tx
// handle wraps the inner VarTx (whose guards keep the underlying VarShard state
// alive for 'a, satisfying R3) and layers the codec for `T <-> bytes`.
// ==========================================================================

/// Multi-shard transaction handle for [`VarTypedTree`] inside `Db::atomicN`.
#[cfg(feature = "armour")]
pub struct VarTypedTx<'a, K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> {
    inner: crate::var_tree::VarTx<'a, K, VarTypedHookAdapter<K, T, C, H>>,
    codec: &'a C,
}

#[cfg(feature = "armour")]
impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>>
    VarTypedTx<'_, K, T, C, H>
{
    pub fn try_get(&self, key: &K) -> DbResult<Option<T>> {
        match self.inner.try_get(key)? {
            Some(bytes) => self.codec.decode_from(&bytes).map(Some),
            None => Ok(None),
        }
    }

    pub fn try_contains(&self, key: &K) -> DbResult<bool> {
        self.inner.try_contains(key)
    }

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

    pub fn put(&mut self, key: &K, value: &T) -> DbResult<()> {
        let mut buf = Vec::new();
        self.codec.encode_to(value, &mut buf)?;
        self.inner.put(key, &buf)
    }

    pub fn insert(&mut self, key: &K, value: &T) -> DbResult<()> {
        let mut buf = Vec::new();
        self.codec.encode_to(value, &mut buf)?;
        self.inner.insert(key, &buf)
    }

    pub fn delete(&mut self, key: &K) -> DbResult<bool> {
        self.inner.delete(key)
    }
}

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

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

    fn begin_tx(&self) -> VarTypedTx<'_, K, T, C, H> {
        VarTypedTx {
            inner: crate::armour::multi_tx::MultiTx::begin_tx(&self.inner),
            codec: &self.codec,
        }
    }

    fn lock_shard_into<'a>(&'a self, shard_id: usize, tx: &mut VarTypedTx<'a, K, T, C, H>) {
        crate::armour::multi_tx::MultiTx::lock_shard_into(&self.inner, shard_id, &mut tx.inner)
    }

    fn release_locks(
        &self,
        tx: &mut VarTypedTx<'_, K, T, C, H>,
    ) -> crate::armour::multi_tx::SyncNeeds {
        crate::armour::multi_tx::MultiTx::release_locks(&self.inner, &mut tx.inner)
    }

    fn run_sync(&self, needs: crate::armour::multi_tx::SyncNeeds) -> DbResult<()> {
        crate::armour::multi_tx::MultiTx::run_sync(&self.inner, needs)
    }

    fn replay_hooks(&self, tx: VarTypedTx<'_, K, T, C, H>) {
        crate::armour::multi_tx::MultiTx::replay_hooks(&self.inner, tx.inner)
    }
}

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

    #[derive(Debug, Clone, PartialEq, Rapira)]
    struct Msg {
        id: u64,
        body: String,
        tags: Vec<String>,
    }

    type MsgTree = VarTypedTree<[u8; 16], Msg, RapiraCodec>;

    fn open(dir: &std::path::Path) -> MsgTree {
        VarTypedTree::open(dir, Config::test(), RapiraCodec).unwrap()
    }

    fn msg(id: u64) -> Msg {
        Msg {
            id,
            body: format!("hello from message #{id}"),
            tags: vec![format!("t{id}"), "bench".into()],
        }
    }

    #[test]
    fn compare_delete_match_mismatch_absent() {
        let dir = tempdir().unwrap();
        let tree = open(dir.path());

        let k = [1u8; 16];
        let v = msg(1);
        let other = msg(99);
        tree.put(&k, &v).unwrap();

        assert!(matches!(
            tree.compare_delete(&k, &other),
            Err(DbError::CasMismatch)
        ));
        assert_eq!(tree.get(&k), Some(v.clone()));

        assert!(tree.compare_delete(&k, &v).is_ok());
        assert!(tree.get(&k).is_none());

        assert!(matches!(
            tree.compare_delete(&k, &v),
            Err(DbError::KeyNotFound)
        ));
    }

    #[test]
    fn entry_len_returns_some_for_existing_key() {
        let tmp = tempdir().expect("tmp");
        let tree: VarTypedTree<[u8; 8], Vec<u8>, RapiraCodec> =
            VarTypedTree::open(tmp.path(), Config::test(), RapiraCodec).expect("open");
        let key = 1u64.to_be_bytes();
        tree.put(&key, &vec![10u8, 20, 30, 40]).expect("put");
        // RapiraCodec encodes Vec<u8> as 4-byte length prefix + data bytes.
        // vec![10, 20, 30, 40] → 4 + 4 = 8 encoded bytes.
        assert_eq!(tree.entry_len(&key), Some(8u32));
    }

    #[test]
    fn entry_len_returns_none_for_missing_key() {
        let tmp = tempdir().expect("tmp");
        let tree: VarTypedTree<[u8; 8], Vec<u8>, RapiraCodec> =
            VarTypedTree::open(tmp.path(), Config::test(), RapiraCodec).expect("open");
        let key = 99u64.to_be_bytes();
        assert_eq!(tree.entry_len(&key), None);
    }

    #[test]
    fn typed_decode_fault_surfaced_on_result_methods_raw_bytes_on_cas() {
        use crate::test_faults::{POISON, PoisonCodec};
        let dir = tempdir().unwrap();
        let tree: VarTypedTree<[u8; 8], u64, PoisonCodec> =
            VarTypedTree::open(dir.path(), Config::test(), PoisonCodec).unwrap();

        // Default reversed=true => first() is the largest key. Put a clean entry and
        // a poisoned boundary entry (key 2 > key 1).
        let k1 = 1u64.to_be_bytes();
        let k2 = 2u64.to_be_bytes();
        tree.put(&k1, &7u64).unwrap();
        tree.put(&k2, &POISON).unwrap();

        // Result methods surface CorruptedEntry, not KeyNotFound / Ok(None).
        assert!(matches!(
            tree.get_or_err(&k2),
            Err(DbError::CorruptedEntry { .. })
        ));
        assert!(matches!(
            tree.try_get(&k2),
            Err(DbError::CorruptedEntry { .. })
        ));
        assert!(matches!(
            tree.update(&k2, |v| v + 1),
            Err(DbError::CorruptedEntry { .. })
        ));
        assert!(matches!(
            tree.fetch_update(&k2, |v| v + 1),
            Err(DbError::CorruptedEntry { .. })
        ));

        // Lossy first() skips the poisoned boundary and returns the next valid entry,
        // matching iter().next().
        assert_eq!(tree.first(), Some((k1, 7u64)));
        assert_eq!(tree.first(), tree.iter().next());
        // Strict try_first surfaces the boundary decode error.
        assert!(matches!(
            tree.try_first(),
            Err(DbError::CorruptedEntry { .. })
        ));

        // cas/compare_delete compare RAW bytes — a decodable-but-poison value is not
        // a read error: mismatch on a different expected, success on the exact bytes.
        assert!(matches!(
            tree.cas(&k2, &123u64, &5u64),
            Err(DbError::CasMismatch)
        ));
        assert!(tree.cas(&k2, &POISON, &5u64).is_ok());
    }

    #[test]
    fn typed_shard_get_or_err_surfaces_decode_fault() {
        use crate::test_faults::{POISON, PoisonCodec};
        let dir = tempdir().unwrap();
        let tree: VarTypedTree<[u8; 8], u64, PoisonCodec> =
            VarTypedTree::open(dir.path(), Config::test(), PoisonCodec).unwrap();

        let k = 1u64.to_be_bytes();
        tree.put(&k, &POISON).unwrap();

        let got = tree.atomic(&k, |shard| shard.get_or_err(&k));
        assert!(matches!(got, Err(DbError::CorruptedEntry { .. })));

        let absent = 999u64.to_be_bytes();
        let got_absent = tree.atomic(&absent, |shard| shard.get_or_err(&absent));
        assert!(matches!(got_absent, Err(DbError::KeyNotFound)));
    }

    #[derive(Clone)]
    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 TRecHookState {
        writes: AtomicUsize,
        last_new: crate::sync::Mutex<Option<u64>>,
    }

    #[derive(Clone, Default)]
    struct TRecHook {
        state: Arc<TRecHookState>,
    }

    impl TypedWriteHook<[u8; 8], u64> for TRecHook {
        fn on_write(&self, _key: &[u8; 8], _old: Option<&u64>, new: Option<&u64>) {
            self.state.writes.fetch_add(1, AtomicOrdering::Relaxed);
            *crate::sync::lock(&self.state.last_new) = new.copied();
        }
    }

    fn open_var_typed_hooked(
        dir: &std::path::Path,
        hook: TRecHook,
    ) -> VarTypedTree<[u8; 8], u64, U64Codec, TRecHook> {
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        VarTypedTree::open_hooked(dir, cfg, U64Codec, hook).expect("open hooked")
    }

    #[test]
    fn var_typed_tree_atomic_fires_typed_hook() {
        let dir = tempdir().unwrap();
        let hook = TRecHook::default();
        let state = hook.state.clone();
        let tree = open_var_typed_hooked(dir.path(), hook);
        let k = 7u64.to_be_bytes();
        tree.atomic(&k, |s| {
            s.put(&k, &42)?;
            Ok(())
        })
        .expect("atomic");
        assert_eq!(state.writes.load(AtomicOrdering::Relaxed), 1);
        assert_eq!(*crate::sync::lock(&state.last_new), Some(42));
    }

    // ---------------------------------------------------------------------------
    // Tests — get_many / update_many (Task 5)
    // ---------------------------------------------------------------------------

    #[test]
    fn var_typed_get_many_and_update_many() {
        use crate::{Applied, BatchWrite};
        let dir = tempdir().unwrap();
        let cfg = Config::test();
        let tree = VarTypedTree::<[u8; 8], String, RapiraCodec>::open(dir.path(), cfg, RapiraCodec)
            .unwrap();
        tree.insert(&1u64.to_be_bytes(), &"a".to_string()).unwrap();
        tree.insert(&2u64.to_be_bytes(), &"b".to_string()).unwrap();

        let keys: Vec<[u8; 8]> = vec![2u64.to_be_bytes(), 9u64.to_be_bytes(), 1u64.to_be_bytes()];
        assert_eq!(
            tree.get_many(&keys),
            vec![Some("b".to_string()), None, Some("a".to_string())]
        );

        let items: Vec<([u8; 8], &str)> =
            vec![(1u64.to_be_bytes(), "a2"), (3u64.to_be_bytes(), "c")];
        let out = tree
            .update_many(items, |_k, _cur, p| BatchWrite::Set(p.to_string()))
            .unwrap();
        assert_eq!(
            out[0].1,
            Applied::Written {
                old: Some("a".to_string()),
                new: "a2".to_string()
            }
        );
        assert_eq!(
            out[1].1,
            Applied::Written {
                old: None,
                new: "c".to_string()
            }
        );
        assert_eq!(tree.get(&1u64.to_be_bytes()), Some("a2".to_string()));
        assert_eq!(tree.get(&3u64.to_be_bytes()), Some("c".to_string()));
    }

    #[test]
    fn var_typed_update_many_strict_on_codec_fault() {
        use crate::test_faults::{POISON, PoisonCodec};
        let dir = tempdir().unwrap();
        let tree: VarTypedTree<[u8; 8], u64, PoisonCodec> =
            VarTypedTree::open(dir.path(), Config::test(), PoisonCodec).unwrap();
        let k = 1u64.to_be_bytes();
        tree.put(&k, &POISON).unwrap();

        // Reading the poisoned old value during update_many must surface the decode
        // error (strict), not silently collapse to Applied::Kept.
        let res = tree.update_many(vec![(k, 5u64)], |_k, _old, v| BatchWrite::Set(*v));
        assert!(matches!(res, Err(DbError::CorruptedEntry { .. })));

        // The poisoned key was not overwritten by a hidden Keep path — raw bytes
        // remain, and a strict read still surfaces the corruption.
        assert!(matches!(
            tree.try_get(&k),
            Err(DbError::CorruptedEntry { .. })
        ));
    }

    #[test]
    fn var_typed_tree_shard_put_returns_existed() {
        let dir = tempdir().unwrap();
        let tree: VarTypedTree<[u8; 8], u64, U64Codec> =
            VarTypedTree::open(dir.path(), Config::test(), U64Codec).unwrap();
        let k = 1u64.to_be_bytes();
        let (fresh, overwrite) = tree
            .atomic(&k, |s| {
                let fresh = s.put(&k, &10u64)?;
                let overwrite = s.put(&k, &20u64)?;
                Ok((fresh, overwrite))
            })
            .expect("atomic");
        assert!(!fresh, "fresh insert must return false");
        assert!(overwrite, "overwrite must return true");
    }

    #[test]
    fn var_typed_tree_shard_update_and_fetch_update() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        let tree: VarTypedTree<[u8; 8], u64, U64Codec> =
            VarTypedTree::open(dir.path(), cfg, U64Codec).unwrap();
        let k = 1u64.to_be_bytes();
        tree.put(&k, &100u64).unwrap();

        // After put: k -> 100
        // update: 100 + 1 = 101 stored, returns Some(101) [new]
        // fetch_update: 101 + 1 = 102 stored, returns Some(101) [old]
        // missing key returns None
        let (upd, fetched, missing) = tree
            .atomic(&k, |s| {
                let upd = s.update(&k, |old| old + 1)?;
                let fetched = s.fetch_update(&k, |old| old + 1)?;
                let missing = s.update(&2u64.to_be_bytes(), |old| *old)?;
                Ok((upd, fetched, missing))
            })
            .expect("atomic");

        assert_eq!(upd, Some(101), "update returns new value");
        assert_eq!(fetched, Some(101), "fetch_update returns old value");
        assert_eq!(missing, None, "absent key returns None");
        assert_eq!(
            tree.get(&k),
            Some(102),
            "final stored value after two increments"
        );
    }

    #[test]
    fn var_typed_tree_atomic_update_fires_hook() {
        let dir = tempdir().unwrap();
        let hook = TRecHook::default();
        let state = hook.state.clone();
        let tree = open_var_typed_hooked(dir.path(), hook);
        let k = 7u64.to_be_bytes();

        // Seed the value so update finds the key.
        tree.put(&k, &10u64).unwrap();
        // Reset counter after the initial put.
        state.writes.store(0, AtomicOrdering::Relaxed);

        tree.atomic(&k, |s| s.update(&k, |old| old + 5))
            .expect("atomic update");

        // The hook must have fired once for the update (new = 15).
        assert_eq!(
            state.writes.load(AtomicOrdering::Relaxed),
            1,
            "hook fired once"
        );
        assert_eq!(
            *crate::sync::lock(&state.last_new),
            Some(15),
            "hook saw the new value"
        );
    }
}