use std::hash::Hash;
use std::mem::size_of;
use std::ops::Bound;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};

use crate::Key;
use crate::byte_view::ByteView;
use crate::cache::{BlockCache, BlockKey};
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, WriteHook};
use crate::io::aligned_buf::AlignedBuf;
use crate::map_index::{
    Dir, KeyPage, MapIter, MapIterView, OrdKey, Page, ShardIndex, ShardSource, dir_from,
    key_index_batch,
};
use crate::recovery::recover_var_map;
use crate::shard::ShardInner;
use crate::sync::{self, Mutex, MutexGuard};
use crate::value_cache::{ValueCache, ValueKey};

const MAX_STALE_RETRIES: usize = 3;

/// A map with fixed-size keys and variable-length values.
/// Uses per-shard HashMap for O(1) lookup. Values are stored on disk,
/// cached at 4096-byte block granularity via `BlockCache`.
/// Ordered scans are available when opened with `iterable(true)` via
/// [`iter_view`](Self::iter_view). Use [`VarTree`](crate::VarTree) for prefix/range
/// scans without the iterable flag.
///
/// Each `VarMap` owns its storage engine — one map = one database directory.
///
/// # Write hooks
///
/// Uses [`WriteHook<K>`]. `on_write` fires on `put`/`insert`/`delete`/`cas`/`compare_delete`/`update`/`fetch_update`
/// and inside `atomic()` (buffered during the closure, replayed after the shard lock is released).
/// When `NEEDS_OLD_VALUE = false`, old value is omitted from `on_write` to skip disk I/O.
///
/// `on_init` fires once per live entry during `migrate()` or `replay_init()`
/// (enable via `NEEDS_INIT = true`).
/// Batch item for [`ShardSource`] — carries the `DiskLoc` observed under the index
/// lock so [`VarMapIterView::retain`] can compare-delete by loc without re-reading disk.
#[doc(hidden)]
pub struct VarMapIterVal {
    pub(crate) loc: DiskLoc,
    pub(crate) value: ByteView,
}

/// Outcome of a locked read-modify-write operation.
enum UpdateOutcome {
    /// Key was not present in the index.
    Missing,
    /// Closure returned `None` — no write should happen.
    Unchanged(ByteView),
    /// Closure returned a new value — entry was written to disk.
    Updated { old: ByteView, new: ByteView },
}

pub struct VarMap<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    indexes: Vec<Mutex<ShardIndex<K, DiskLoc>>>,
    engine: Engine,
    block_cache: BlockCache,
    value_cache: ValueCache,
    compaction_threshold: f64,
    shard_prefix_bits: usize,
    hook: H,
    iterable: bool,
    reversed: bool,
    entry_count: AtomicUsize,
}

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

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
    /// Open or create a `VarMap` with a write hook for secondary index maintenance.
    pub fn open_hooked(
        path: impl AsRef<std::path::Path>,
        config: Config,
        hook: H,
    ) -> DbResult<Self> {
        let config = config.with_resolved_hints(true);
        if config.var_cache_warn_needed() {
            tracing::warn!(
                "Var collection opened with no cache (block_cache=0 and value_cache=0): every get hits disk"
            );
        }
        if config.app_cache_without_direct_io() {
            tracing::warn!(
                "Var collection has an app cache enabled (block_cache/value_cache) but direct_io is off: \
                 values are cached twice (app cache + OS page cache). Enable direct_io to avoid redundant \
                 RAM — exception: keeping the app cache for plaintext caching under encryption."
            );
        }
        Self::open_inner(path, config, hook)
    }

    fn open_inner(path: impl AsRef<std::path::Path>, config: Config, hook: H) -> DbResult<Self> {
        let compaction_threshold = config.compaction_threshold;
        let shard_prefix_bits = config.shard_prefix_bits;
        let iterable = config.iterable;
        let reversed = config.reversed;
        let block_cache = BlockCache::new(&config.block_cache);
        let value_cache = ValueCache::new(&config.value_cache);
        let engine = Engine::open(path, config)?;

        let shard_count = engine.shards().len();
        let mut indexes = Vec::with_capacity(shard_count);
        for _ in 0..shard_count {
            // Built once after recovery via `set_iterable_and_rebuild` — see the
            // ConstMap `open_inner` for the rationale.
            indexes.push(Mutex::new(ShardIndex::new(false)));
        }

        let map = Self {
            indexes,
            engine,
            block_cache,
            value_cache,
            compaction_threshold,
            shard_prefix_bits,
            hook,
            iterable,
            reversed,
            entry_count: AtomicUsize::new(0),
        };

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

        let hints = map.engine.hints();
        let outcome = recover_var_map::<K>(
            &shard_dir_refs,
            &shard_ids,
            map.indexes(),
            hints,
            #[cfg(feature = "encryption")]
            map.engine.cipher(),
        )?;
        for tail in &outcome.active_tails {
            map.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
        }
        for (shard_idx, dead) in outcome.shard_dead_bytes {
            map.engine.shards()[shard_idx].install_dead_bytes(dead);
        }
        let max_gsn = outcome.max_gsn;

        let initial_len: usize = map.indexes.iter().map(|m| sync::lock(m).len()).sum();
        map.entry_count.store(initial_len, AtomicOrdering::Relaxed);

        map.engine
            .gsn()
            .fetch_max(max_gsn + 1, AtomicOrdering::Relaxed);
        if hints {
            for shard in map.engine.shards().iter() {
                shard.set_key_len(size_of::<K>());
            }
        }
        if iterable {
            for idx in &map.indexes {
                sync::lock(idx).set_iterable_and_rebuild();
            }
        }
        tracing::info!(
            key_size = size_of::<K>(),
            entries = map.len(),
            "var_map recovered"
        );

        Ok(map)
    }

    /// 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 + Send + Sync + Hash + Eq, H: WriteHook<K>> CompactionIndex<K> for VarMap<K, H> {
    fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
        let mut index = sync::lock(&self.indexes[self.shard_for(key)]);
        if let Some(disk) = index.get_mut(key)
            && *disk == old_loc
        {
            *disk = new_loc;
            return true;
        }
        false
    }

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

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

    fn is_live(&self, shard_id: u8, key: &K, loc: DiskLoc) -> bool {
        let index = sync::lock(&self.indexes[shard_id as usize]);
        index.get(key).is_some_and(|disk| *disk == loc)
    }
}

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
    /// Trigger a compaction pass across all shards.
    pub fn compact(&self) -> DbResult<usize> {
        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 value by key. O(1) average index lookup; checks block cache first, then disk.
    /// Takes a brief per-shard mutex to look up the location; the `pread` on a cache miss
    /// runs outside the lock.
    pub fn get(&self, key: &K) -> Option<ByteView> {
        metrics::counter!("armdb.ops", "op" => "get", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.get");
        let result = self.try_get(key);
        #[cfg(feature = "hot-path-tracing")]
        if let Err(e) = &result {
            tracing::error!("VarMap get error: {:?}", e);
        }
        result.unwrap_or_default()
    }

    /// Strict read: `Ok(None)` only when the key is genuinely absent; `Err` when the
    /// key exists but its value cannot be read. The result-returning primitive for
    /// `VarMap` (no node helper — re-locks the shard index per stale-retry).
    pub fn try_get(&self, key: &K) -> DbResult<Option<ByteView>> {
        let shard_id = self.shard_for(key);
        for _ in 0..MAX_STALE_RETRIES {
            let (disk, cached) = {
                let index = sync::lock(&self.indexes[shard_id]);
                let disk = match index.get(key) {
                    Some(d) => *d,
                    None => return Ok(None),
                };
                // Probe the value cache while holding the index lock: a hit then
                // references the current loc ⇒ file live ⇒ bytes are current.
                let cached = if disk.is_value_cache_routed() {
                    self.value_cache.get(&ValueKey {
                        shard_id: shard_id as u8,
                        file_id: disk.file_id,
                        offset: disk.offset,
                    })
                } else {
                    None
                };
                (disk, cached)
            };
            if let Some(v) = cached {
                metrics::counter!("armdb.value_cache.hit").increment(1);
                return Ok(Some(v));
            }
            let is_large = disk.is_value_cache_routed();
            if is_large {
                metrics::counter!("armdb.value_cache.miss").increment(1);
            }
            match self.read_value_cached_inner(&disk, shard_id as u8) {
                Ok(v) => {
                    if is_large {
                        self.value_cache.insert(
                            ValueKey {
                                shard_id: shard_id as u8,
                                file_id: disk.file_id,
                                offset: disk.offset,
                            },
                            v.clone(),
                        );
                    }
                    return Ok(Some(v));
                }
                Err(DbError::StaleDiskLoc) => {
                    metrics::counter!("armdb.read.stale_retry", "tree" => "var_map").increment(1);
                    continue;
                }
                Err(e) => return Err(e),
            }
        }
        Err(DbError::StaleDiskLoc)
    }

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

    /// Insert or update a key-value pair.
    ///
    /// Returns `true` if a previous value for `key` existed (i.e. this was an
    /// overwrite), `false` for a fresh insert. The flag is read directly from
    /// the index update and adds no extra I/O or value decoding.
    pub fn put(&self, key: &K, value: &[u8]) -> DbResult<bool> {
        metrics::counter!("armdb.ops", "op" => "put", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.put");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);
        let old_value = if H::NEEDS_OLD_VALUE {
            if let Some(disk) = index.get(key) {
                Some(self.read_value_locked_result(disk, shard_id as u8, &inner)?)
            } else {
                None
            }
        } else {
            None
        };
        let existed = self.put_locked(shard_id, &mut inner, &mut index, key, value)?;
        drop(index);
        drop(inner);
        self.hook.on_write(key, old_value.as_deref(), Some(value));
        Ok(existed)
    }

    /// 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: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "insert", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.insert");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);
        self.insert_locked(shard_id, &mut inner, &mut index, key, value)?;
        drop(index);
        drop(inner);
        self.hook.on_write(key, None, Some(value));
        Ok(())
    }

    /// Delete a key. Returns `true` if the key existed.
    pub fn delete(&self, key: &K) -> DbResult<bool> {
        metrics::counter!("armdb.ops", "op" => "delete", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.delete");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);
        let old_value = if H::NEEDS_OLD_VALUE {
            if let Some(disk) = index.get(key) {
                Some(self.read_value_locked_result(disk, shard_id as u8, &inner)?)
            } else {
                None
            }
        } else {
            None
        };
        let existed = self.delete_locked(shard_id, &mut inner, &mut index, key)?;
        drop(index);
        drop(inner);
        if existed {
            self.hook.on_write(key, old_value.as_deref(), None);
        }
        Ok(existed)
    }

    /// 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.
    ///
    /// **Caveat:** Holds the shard lock while reading the current value. On a
    /// block-cache miss this performs disk I/O under the lock, blocking all
    /// writes to the same shard. Pre-warm the cache to avoid latency spikes.
    pub fn cas(&self, key: &K, expected: &[u8], new_value: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "cas", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.cas");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);

        let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
        let current = self.read_value_locked_result(&disk, shard_id as u8, &inner)?;
        if current.as_ref() != expected {
            return Err(DbError::CasMismatch);
        }

        let (new_disk_loc, _gsn) =
            inner.append_entry(shard_id as u8, key.as_bytes(), new_value, false)?;
        inner.add_dead_bytes(
            disk.file_id,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );
        let _old = index.upsert(*key, new_disk_loc);
        debug_assert!(_old.is_some(), "cas: key must exist in index");

        drop(index);
        drop(inner);
        self.hook.on_write(
            key,
            if H::NEEDS_OLD_VALUE {
                Some(&*current)
            } else {
                None
            },
            Some(new_value),
        );
        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.
    ///
    /// **Caveat:** Holds the shard lock while reading the current value. On a
    /// block-cache miss this performs disk I/O under the lock, blocking all
    /// writes to the same shard. Pre-warm the cache to avoid latency spikes.
    pub fn compare_delete(&self, key: &K, expected: &[u8]) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.compare_delete");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);

        let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
        let current = self.read_value_locked_result(&disk, shard_id as u8, &inner)?;
        if current.as_ref() != expected {
            return Err(DbError::CasMismatch);
        }

        let removed = self.delete_locked(shard_id, &mut inner, &mut index, key)?;
        debug_assert!(
            removed,
            "compare_delete: key must exist after a successful match"
        );

        drop(index);
        drop(inner);
        self.hook.on_write(
            key,
            if H::NEEDS_OLD_VALUE {
                Some(&*current)
            } else {
                None
            },
            None,
        );
        Ok(())
    }

    /// Compare-and-delete by `DiskLoc`: delete only if the live entry still
    /// points at `expected`. Used by [`VarMapIterView::retain`] without
    /// re-reading the value from disk under the shard lock.
    pub(crate) fn compare_delete_if_loc(&self, key: &K, expected: DiskLoc) -> DbResult<()> {
        metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "var_map").increment(1);
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);

        let disk = *index.get(key).ok_or(DbError::KeyNotFound)?;
        if disk != expected {
            return Err(DbError::CasMismatch);
        }

        let old_value = if H::NEEDS_OLD_VALUE {
            Some(self.read_value_locked_result(&disk, shard_id as u8, &inner)?)
        } else {
            None
        };
        let removed = self.delete_locked(shard_id, &mut inner, &mut index, key)?;
        debug_assert!(
            removed,
            "compare_delete_if_loc: key must exist after a successful loc match"
        );

        drop(index);
        drop(inner);
        self.hook.on_write(key, old_value.as_deref(), None);
        Ok(())
    }

    /// Atomically read-modify-write. Returns `Some(new_value)` if key existed, `None` otherwise.
    /// The closure receives the current value and returns a new ByteView.
    /// The closure must not be heavy (shard lock is held).
    ///
    /// **Caveat:** Holds the shard lock while reading the current value. On a
    /// block-cache miss this performs disk I/O under the lock, blocking all
    /// writes to the same shard. Pre-warm the cache to avoid latency spikes.
    pub fn update(&self, key: &K, f: impl FnOnce(&[u8]) -> ByteView) -> DbResult<Option<ByteView>> {
        self.update_inner(key, f, false)
    }

    /// Like [`update()`](Self::update), but returns `Some(old_value)` instead of the new one.
    pub fn fetch_update(
        &self,
        key: &K,
        f: impl FnOnce(&[u8]) -> ByteView,
    ) -> DbResult<Option<ByteView>> {
        self.update_inner(key, f, true)
    }

    fn try_update_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, DiskLoc>,
        key: &K,
        f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
    ) -> DbResult<UpdateOutcome> {
        let disk = match index.get(key) {
            Some(d) => *d,
            None => return Ok(UpdateOutcome::Missing),
        };
        let current = self.read_value_locked_result(&disk, shard_id as u8, inner)?;

        let new_value = match f(&current)? {
            Some(v) => v,
            None => return Ok(UpdateOutcome::Unchanged(current)),
        };

        let (new_disk_loc, _gsn) =
            inner.append_entry(shard_id as u8, key.as_bytes(), &new_value, false)?;
        inner.add_dead_bytes(
            disk.file_id,
            crate::entry::entry_size(size_of::<K>(), disk.len),
        );
        let _old = index.upsert(*key, new_disk_loc);
        debug_assert!(_old.is_some(), "update: key must exist in index");

        Ok(UpdateOutcome::Updated {
            old: current,
            new: new_value,
        })
    }

    pub(crate) fn try_update_inner(
        &self,
        key: &K,
        f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
        return_old: bool,
    ) -> DbResult<Option<ByteView>> {
        metrics::counter!("armdb.ops", "op" => "update", "tree" => "var_map").increment(1);
        #[cfg(feature = "hot-path-tracing")]
        tracing::trace!("var_map.update");
        let shard_id = self.shard_for(key);
        let mut inner = self.engine.shards()[shard_id].lock();
        let mut index = sync::lock(&self.indexes[shard_id]);

        match self.try_update_locked(shard_id, &mut inner, &mut index, key, f)? {
            UpdateOutcome::Missing => Ok(None),
            UpdateOutcome::Unchanged(current) => Ok(Some(current)),
            UpdateOutcome::Updated { old, new } => {
                drop(index);
                drop(inner);
                self.hook.on_write(
                    key,
                    if H::NEEDS_OLD_VALUE {
                        Some(&*old)
                    } else {
                        None
                    },
                    Some(&*new),
                );
                Ok(Some(if return_old { old } else { new }))
            }
        }
    }

    fn update_inner(
        &self,
        key: &K,
        f: impl FnOnce(&[u8]) -> ByteView,
        return_old: bool,
    ) -> DbResult<Option<ByteView>> {
        self.try_update_inner(key, |bytes| Ok(Some(f(bytes))), return_old)
    }

    /// Check if a key exists.
    pub fn contains(&self, key: &K) -> bool {
        let index = sync::lock(&self.indexes[self.shard_for(key)]);
        index.contains_key(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> {
        let index = sync::lock(&self.indexes[self.shard_for(key)]);
        index.get(key).map(|disk| disk.len)
    }

    /// Total number of entries across all shards.
    pub fn len(&self) -> usize {
        self.entry_count.load(AtomicOrdering::Relaxed)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

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

    // -- Ordered iteration ----------------------------------------------------

    /// Ordered scan view. `None` unless the collection was opened with
    /// `iterable(true)`. Direction follows `Config::reversed` (default DESC).
    ///
    /// **Occasional admin/maintenance API, not a hot path.** Use it to browse a
    /// collection or bulk-clean stale rows (`retain`); for scan-heavy access reach
    /// for the `*Tree` sibling — a Map's strength is O(1) point lookup, and the
    /// companion key index costs extra memory plus O(log N) per write. Unlike a
    /// `Tree` scan it is **not a point-in-time snapshot**: the k-way merge drops
    /// each shard lock between batches, so concurrent writes may be only partially
    /// reflected (run on a quiet DB if you need a coherent view).
    pub fn iter_view(&self) -> Option<VarMapIterView<'_, K, H>> {
        self.iterable.then(|| VarMapIterView {
            map: self,
            inner: MapIterView::new(self, dir_from(self.reversed)),
        })
    }

    /// Pre-populate the block cache with blocks containing live values.
    ///
    /// Walks the index to collect unique block offsets, sorts them for sequential I/O,
    /// then reads each block into the cache. Only blocks with live data are loaded.
    pub fn warmup(&self) -> DbResult<()> {
        use std::collections::BTreeSet;

        let mut blocks: BTreeSet<(u8, u32, u64)> = BTreeSet::new();
        for (i, index) in self.indexes.iter().enumerate() {
            let map = sync::lock(index);
            for disk in map.values() {
                // Large values are served from the value_cache, not the block
                // cache — skip them so warmup doesn't waste block-cache budget.
                if disk.is_value_cache_routed() {
                    continue;
                }
                let block_offset = disk.offset as u64 & !4095;
                blocks.insert((i as u8, disk.file_id, block_offset));
            }
        }

        for (shard_id, file_id, block_offset) in &blocks {
            let key = BlockKey {
                shard_id: *shard_id,
                file_id: *file_id,
                block_offset: *block_offset,
            };
            if self.block_cache.get(&key).is_some() {
                continue;
            }
            let shard = &self.engine.shards()[*shard_id as usize];
            let (buf, is_full_block) = shard.read_block(*file_id, *block_offset)?;
            if is_full_block {
                self.block_cache.insert(key, Arc::new(buf));
            }
        }

        Ok(())
    }

    // -- Read-only scan -------------------------------------------------------

    /// Read-only pass over all live entries. Order is unspecified; entries
    /// inserted/removed concurrently may be missed or visited once — intended
    /// for schema validation on a quiet database.
    ///
    /// Lock order matches `replay_init`: snapshot keys under the per-shard index
    /// lock, release it, then read each value via `self.try_get` (which
    /// re-acquires the index lock only briefly). A key that is genuinely absent
    /// by read time was concurrently deleted — skipped silently; only a real
    /// read failure is logged with a `warn`.
    pub fn for_each(&self, mut f: impl FnMut(K, ByteView)) {
        for i in 0..self.engine.shards().len() {
            let keys: Vec<K> = {
                let index = sync::lock(&self.indexes[i]);
                index.keys().copied().collect()
            };
            for key in keys {
                match self.try_get(&key) {
                    Ok(Some(value)) => f(key, value),
                    // Concurrently deleted between the key snapshot and the
                    // read — benign, skip silently.
                    Ok(None) => {}
                    Err(e) => {
                        tracing::warn!(
                            key = ?key.as_bytes(),
                            error = ?e,
                            "var_map for_each: skipping entry — value read failed"
                        );
                    }
                }
            }
        }
    }

    /// Iterate all entries and optionally mutate them. Call once at startup.
    ///
    /// - `Keep` — no change (fires `on_init` if `H::NEEDS_INIT`); not counted
    /// - `Update(new_value)` — replace value (hook-free write, fires `on_init`)
    /// - `Delete` — remove entry (hook-free tombstone, no hooks)
    ///
    /// `on_write` is NEVER fired during migrate — see `docs/hooks.md` in the armdb crate.
    /// Returns the number of mutated entries.
    pub fn migrate(
        &self,
        f: impl Fn(&K, &[u8]) -> crate::MigrateAction<ByteView>,
    ) -> DbResult<usize> {
        use crate::MigrateAction;

        let mut count = 0;
        for i in 0..self.engine.shards().len() {
            let keys: Vec<K> = {
                let index = sync::lock(&self.indexes[i]);
                index.keys().copied().collect()
            };
            for key in keys {
                let value = match self.get(&key) {
                    Some(v) => v,
                    None => {
                        tracing::warn!(
                            key = ?key.as_bytes(),
                            "var_map migrate: skipping entry — value read failed"
                        );
                        continue;
                    }
                };
                match f(&key, &value) {
                    MigrateAction::Keep => {
                        if H::NEEDS_INIT {
                            self.hook.on_init(&key, &value);
                        }
                    }
                    MigrateAction::Update(new_value) => {
                        // Keys were collected from indexes[i], so they always
                        // hash to shard i — no need to recompute via shard_for.
                        {
                            let mut inner = self.engine.shards()[i].lock();
                            let mut index = sync::lock(&self.indexes[i]);
                            self.put_locked(i, &mut inner, &mut index, &key, &new_value)?;
                        }
                        if H::NEEDS_INIT {
                            self.hook.on_init(&key, &new_value);
                        }
                        count += 1;
                    }
                    MigrateAction::Delete => {
                        let mut inner = self.engine.shards()[i].lock();
                        let mut index = sync::lock(&self.indexes[i]);
                        self.delete_locked(i, &mut inner, &mut index, &key)?;
                        count += 1;
                    }
                }
            }
        }

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

    /// Replay `on_init` for every live entry. Used when no migration runs
    /// (Db calls this after `run_migration` returns `false`).
    ///
    /// Lock order: snapshot keys under per-shard index lock, release the
    /// lock, then call `self.get(key)` which re-acquires the index lock
    /// briefly to copy the `DiskLoc` and reads the value WITHOUT holding
    /// the index lock. This matches the lock order of public writes
    /// (shard → index) and avoids the inverted order that would deadlock
    /// against an active writer.
    pub(crate) fn replay_init(&self) {
        if !H::NEEDS_INIT {
            return;
        }
        let mut count = 0usize;
        for i in 0..self.engine.shards().len() {
            let keys: Vec<K> = {
                let index = sync::lock(&self.indexes[i]);
                index.keys().copied().collect()
            };
            for key in keys {
                let value = match self.get(&key) {
                    Some(v) => v,
                    None => {
                        tracing::warn!(
                            key = ?key.as_bytes(),
                            "var_map replay_init: skipping entry — value read failed"
                        );
                        continue;
                    }
                };
                self.hook.on_init(&key, &value);
                count += 1;
            }
        }
        tracing::debug!(replayed = count, "var_map replay_init complete");
    }

    /// 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 VarMapShard<'_, K, H>) -> DbResult<R>,
    ) -> DbResult<R> {
        let shard_id = self.shard_for(shard_key);
        let inner = self.engine.shards()[shard_id].lock();
        let index = sync::lock(&self.indexes[shard_id]);
        let mut shard = VarMapShard {
            tree: self,
            inner,
            index,
            shard_id,
            events: Vec::new(),
        };
        let result = f(&mut shard);
        let VarMapShard {
            inner,
            index,
            events,
            ..
        } = shard;
        drop(index);
        drop(inner);
        if H::NEEDS_WRITE {
            for (k, old, new) in &events {
                self.hook.on_write(k, old.as_deref(), new.as_deref());
            }
        }
        result
    }

    pub(crate) fn indexes(&self) -> &[Mutex<ShardIndex<K, DiskLoc>>] {
        &self.indexes
    }

    /// Returns `true` if `key` already had a value (overwrite), `false` if newly inserted.
    fn put_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, DiskLoc>,
        key: &K,
        value: &[u8],
    ) -> DbResult<bool> {
        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;

        if let Some(old_disk) = index.upsert(*key, disk_loc) {
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );
            Ok(true)
        } else {
            self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
            Ok(false)
        }
    }

    fn insert_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, DiskLoc>,
        key: &K,
        value: &[u8],
    ) -> DbResult<()> {
        if index.contains_key(key) {
            return Err(DbError::KeyExists);
        }

        let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), value, false)?;
        index.upsert(*key, disk_loc);
        self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
        Ok(())
    }

    fn delete_locked(
        &self,
        shard_id: usize,
        inner: &mut ShardInner,
        index: &mut ShardIndex<K, DiskLoc>,
        key: &K,
    ) -> DbResult<bool> {
        if !index.contains_key(key) {
            return Ok(false);
        }

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

        if let Some(old_disk) = index.remove(key) {
            inner.add_dead_bytes(
                old_disk.file_id,
                crate::entry::entry_size(size_of::<K>(), old_disk.len),
            );
            self.entry_count.fetch_sub(1, AtomicOrdering::Relaxed);
            Ok(true)
        } else {
            Ok(false)
        }
    }

    fn read_value_cached_inner(&self, disk: &DiskLoc, shard_id: u8) -> DbResult<ByteView> {
        let len = disk.len as usize;
        let start = (disk.offset & 4095) as usize;

        // Large values (>2 blocks) bypass the block-cache path entirely;
        // caching for large values is done in `try_get` under the index lock.
        // This check MUST sit before the cache lookup: a warm first block would
        // otherwise route a large value into `extract_from_block`, which only
        // supports two blocks and would panic on `next[..second_len]` when
        // `second_len > 4096`. StaleDiskLoc propagates so VarMap::get's
        // per-shard-index retry loop takes a fresh DiskLoc snapshot.
        if start + len > 8192 {
            {
                let shard = &self.engine.shards()[shard_id as usize];
                let inner = shard.lock();
                if inner.active.file_id == disk.file_id
                    && let Some(bytes) = inner.write_buf.read(disk.offset as u64, len)
                {
                    return Ok(ByteView::new(bytes));
                }
            }
            let shard = &self.engine.shards()[shard_id as usize];
            return Ok(ByteView::from_vec(shard.read_value(
                disk.file_id,
                disk.offset,
                len,
            )?));
        }

        let block_offset = disk.offset as u64 & !4095;
        let cache_key = BlockKey {
            shard_id,
            file_id: disk.file_id,
            block_offset,
        };

        if let Some(block) = self.block_cache.get(&cache_key) {
            metrics::counter!("armdb.block_cache.hit").increment(1);
            return Self::extract_from_block(&block, start, len, || {
                self.get_or_read_block(shard_id, disk.file_id, block_offset + 4096)
            });
        }

        {
            let shard = &self.engine.shards()[shard_id as usize];
            let inner = shard.lock();
            if inner.active.file_id == disk.file_id
                && let Some(bytes) = inner.write_buf.read(disk.offset as u64, len)
            {
                return Ok(ByteView::new(bytes));
            }
        }

        metrics::counter!("armdb.block_cache.miss").increment(1);
        let block = self.get_or_read_block(shard_id, disk.file_id, block_offset)?;
        Self::extract_from_block(&block, start, len, || {
            self.get_or_read_block(shard_id, disk.file_id, block_offset + 4096)
        })
    }

    fn extract_from_block(
        block: &AlignedBuf,
        start: usize,
        len: usize,
        next_block: impl FnOnce() -> DbResult<Arc<AlignedBuf>>,
    ) -> DbResult<ByteView> {
        debug_assert!(
            start + len <= 8192,
            "extract_from_block supports at most 2 blocks (8192 bytes)"
        );
        if start + len <= 4096 {
            Ok(ByteView::new(&block[start..start + len]))
        } else {
            let next = next_block()?;
            let first_part = &block[start..];
            let second_len = len - first_part.len();
            let mut combined = Vec::with_capacity(len);
            combined.extend_from_slice(first_part);
            combined.extend_from_slice(&next[..second_len]);
            Ok(ByteView::from_vec(combined))
        }
    }

    fn get_or_read_block(
        &self,
        shard_id: u8,
        file_id: u32,
        block_offset: u64,
    ) -> DbResult<Arc<AlignedBuf>> {
        let key = BlockKey {
            shard_id,
            file_id,
            block_offset,
        };
        if let Some(cached) = self.block_cache.get(&key) {
            return Ok(cached);
        }
        let shard = &self.engine.shards()[shard_id as usize];
        let (buf, is_full_block) = shard.read_block(file_id, block_offset)?;
        let arc = Arc::new(buf);
        if is_full_block {
            self.block_cache.insert(key, arc.clone());
        }
        Ok(arc)
    }

    /// Read value when shard lock is already held, propagating `DbError`.
    /// Used by both `read_value_locked` (Option wrapper) and the large-value
    /// fallback in `read_value_cached_inner`.
    fn read_value_locked_result(
        &self,
        disk: &DiskLoc,
        shard_id: u8,
        inner: &ShardInner,
    ) -> DbResult<ByteView> {
        let len = disk.len as usize;

        // 1. Write buffer (for unflushed data)
        if inner.active.file_id == disk.file_id
            && let Some(bytes) = inner.write_buf.read(disk.offset as u64, len)
        {
            return Ok(ByteView::new(bytes));
        }

        // 2. Block cache (lock-free, single-block fast path)
        let block_offset = disk.offset as u64 & !4095;
        let start = (disk.offset & 4095) as usize;
        if start + len <= 4096 {
            let cache_key = BlockKey {
                shard_id,
                file_id: disk.file_id,
                block_offset,
            };
            if let Some(block) = self.block_cache.get(&cache_key) {
                return Ok(ByteView::new(&block[start..start + len]));
            }

            // 2b. Cache miss — read block under held lock, populate cache
            let (buf, is_full_block) = inner.read_block_locked(disk.file_id, block_offset)?;
            let arc = Arc::new(buf);
            if is_full_block {
                self.block_cache.insert(cache_key, arc.clone());
            }
            return Ok(ByteView::new(&arc[start..start + len]));
        }

        // 3. Multi-block: disk read via encryption-aware helper (no cache).
        let bytes = inner.read_value_from_disk_locked(disk)?;
        Ok(ByteView::new(&bytes))
    }

    /// Read value when shard lock is already held. Used by CAS/update.
    ///
    /// Thin `Option` wrapper around `read_value_locked_result` that preserves
    /// the existing call-site contract: `StaleDiskLoc` is logged at error
    /// level and collapsed to `None`; other errors collapse to `None` silently.
    fn read_value_locked(
        &self,
        disk: &DiskLoc,
        shard_id: u8,
        inner: &ShardInner,
    ) -> Option<ByteView> {
        match self.read_value_locked_result(disk, shard_id, inner) {
            Ok(v) => Some(v),
            Err(DbError::StaleDiskLoc) => {
                tracing::error!(
                    file_id = disk.file_id,
                    shard_id,
                    "stale DiskLoc under shard lock - programming bug",
                );
                None
            }
            Err(_) => None,
        }
    }

    pub fn shard_for(&self, key: &K) -> usize {
        crate::shard_for_key(key, self.shard_prefix_bits, self.engine.shards().len())
    }
}

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> ShardSource for VarMap<K, H> {
    type Key = K;
    type Val = VarMapIterVal;

    fn shard_count(&self) -> usize {
        self.engine.shards().len()
    }

    fn next_batch(
        &self,
        shard: usize,
        after: Option<&K>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<K>>,
        range_end: Bound<OrdKey<K>>,
    ) -> Vec<(K, VarMapIterVal)> {
        // Keys whose value read fails (concurrently deleted → `Ok(None)`, or
        // unreadable → `Err`) are skipped, but the cursor must still advance past
        // them: otherwise a full batch of skipped keys would make the iterator
        // believe the shard is exhausted and silently drop everything after it.
        // So we keep scanning forward (advancing `cursor`) until we have `n`
        // readable entries or the companion is genuinely exhausted in range.
        let mut out = Vec::with_capacity(n);
        let mut cursor: Option<K> = after.copied();
        loop {
            let snapshots: Vec<(K, DiskLoc)> = {
                let index = sync::lock(&self.indexes[shard]);
                let Some(ki) = index.key_index() else {
                    return out;
                };
                let keys = key_index_batch(ki, cursor.as_ref(), dir, n, range_start, range_end);
                // Under the same lock, every key from the companion is present in
                // the map, so this never drops a key — it only reads the `DiskLoc`.
                keys.into_iter()
                    .filter_map(|k| index.get(&k).map(|disk| (k, *disk)))
                    .collect()
            };

            if snapshots.is_empty() {
                return out; // shard exhausted in range
            }
            let scanned = snapshots.len();
            let last_scanned = snapshots.last().map(|(k, _)| *k);

            for (key, loc) in snapshots {
                match self.try_get(&key) {
                    Ok(Some(value)) => out.push((key, VarMapIterVal { loc, value })),
                    // Concurrently deleted between the key snapshot and the read.
                    Ok(None) => {}
                    Err(e) => {
                        tracing::warn!(
                            key = ?key.as_bytes(),
                            error = ?e,
                            "var_map iter: skipping entry — value read failed"
                        );
                    }
                }
            }

            // Enough collected, or the companion returned a short (final) chunk.
            if out.len() >= n || scanned < n {
                return out;
            }
            // Full chunk but some keys were skipped — advance past it and continue.
            cursor = last_scanned;
        }
    }

    /// Keys straight from the companion `BTreeSet` — no `try_get`, so no off-disk
    /// read. This is the cheap path that makes keys-only iteration worthwhile for
    /// `VarMap`/`VarTypedMap`.
    fn next_keys_batch(
        &self,
        shard: usize,
        after: Option<&K>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<K>>,
        range_end: Bound<OrdKey<K>>,
    ) -> Vec<K> {
        let index = sync::lock(&self.indexes[shard]);
        let Some(ki) = index.key_index() else {
            return Vec::new();
        };
        key_index_batch(ki, after, dir, n, range_start, range_end)
    }
}

impl<'a, K, H> VarMapIterView<'a, K, H>
where
    K: Key + Send + Sync + Hash + Eq,
    H: WriteHook<K>,
{
    /// Full ordered scan yielding `(K, ByteView)`.
    pub fn iter(&self) -> VarMapValueIter<'a, K, H> {
        VarMapValueIter {
            inner: self.inner.iter(),
        }
    }

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    pub fn range(&self, start: &K, end: &K) -> VarMapValueIter<'a, K, H> {
        VarMapValueIter {
            inner: self.inner.range(start, end),
        }
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> VarMapValueIter<'a, K, H> {
        VarMapValueIter {
            inner: self.inner.range_bounds(start, end),
        }
    }

    /// Iterate entries whose keys start with `prefix`.
    pub fn prefix_iter(&self, prefix: &[u8]) -> VarMapValueIter<'a, K, H> {
        VarMapValueIter {
            inner: self.inner.prefix_iter(prefix),
        }
    }

    /// Keyset pagination in the collection's configured direction.
    pub fn paginate(&self, after: Option<&K>, limit: usize) -> Page<K, ByteView> {
        let iter = match (after, dir_from(self.map.reversed)) {
            (None, _) => self.iter(),
            (Some(k), Dir::Asc) => self.range_bounds(Bound::Excluded(k), Bound::Unbounded),
            (Some(k), Dir::Desc) => self.range_bounds(Bound::Unbounded, Bound::Excluded(k)),
        };
        let items: Vec<_> = iter.take(limit).collect();
        let next = if items.len() == limit {
            items.last().map(|(k, _)| *k)
        } else {
            None
        };
        Page { items, next }
    }

    // -- Keys-only iteration -------------------------------------------------
    //
    // These read keys straight from the companion `BTreeSet`, so — unlike value
    // iteration — they perform **no off-disk read**. This is the main reason to
    // reach for keys-only on Var collections.

    /// Full ordered **keys-only** scan — no off-disk value read.
    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
        self.inner.keys()
    }

    /// Keys-only iteration over `[start, end)`.
    pub fn keys_range(&self, start: &K, end: &K) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_range(start, end)
    }

    /// Keys-only iteration over a bounded range.
    pub fn keys_range_bounds(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
    ) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_range_bounds(start, end)
    }

    /// Keys-only iteration over keys starting with `prefix`.
    pub fn keys_prefix(&self, prefix: &[u8]) -> impl Iterator<Item = K> + '_ {
        self.inner.keys_prefix(prefix)
    }

    /// Keys-only pagination — cheap, no off-disk value read.
    pub fn keys_paginate(&self, after: Option<&K>, limit: usize) -> KeyPage<K> {
        self.inner.keys_paginate(after, limit)
    }

    /// Keep entries where `f` returns `true`; delete the rest via loc-revalidating
    /// compare-delete (no disk re-read for the equality check).
    ///
    /// If the `DiskLoc` changed between iteration and delete (`CasMismatch`) or the
    /// key was removed concurrently (`KeyNotFound`), the delete is skipped (R5).
    /// Unreadable values are skipped and not deleted (R4). Not atomic across the
    /// collection; a genuine storage error aborts the scan and is returned.
    pub fn retain(&self, mut f: impl FnMut(&K, &ByteView) -> bool) -> DbResult<()> {
        for (key, entry) in self.inner.iter() {
            if f(&key, &entry.value) {
                continue;
            }
            match self.map.compare_delete_if_loc(&key, entry.loc) {
                Ok(()) => {}
                Err(DbError::CasMismatch) | Err(DbError::KeyNotFound) => {
                    tracing::debug!("retain: compare_delete skipped (concurrent change)");
                }
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

/// View over an iterable [`VarMap`]; exposes ordered iteration yielding `ByteView` values.
pub struct VarMapIterView<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    map: &'a VarMap<K, H>,
    inner: MapIterView<'a, VarMap<K, H>>,
}

/// Iterator over entries in an iterable [`VarMap`], yielding `(K, ByteView)`.
pub struct VarMapValueIter<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> {
    inner: MapIter<'a, VarMap<K, H>>,
}

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> Iterator for VarMapValueIter<'_, K, H> {
    type Item = (K, ByteView);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next().map(|(k, entry)| (k, entry.value))
    }
}

#[cfg(feature = "replication")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> crate::replication::ReplicationTarget
    for VarMap<K, 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> {
        use crate::replication::ApplyOutcome;

        let key = K::from_bytes(key);

        let disk = DiskLoc::new(
            file_id,
            (entry_offset + size_of::<crate::entry::EntryHeader>() as u64 + size_of::<K>() as u64)
                as u32,
            header.value_len,
        );

        if header.is_tombstone() {
            let old = sync::lock(&self.indexes[self.shard_for(&key)]).remove(&key);
            match old {
                Some(old_disk) => {
                    self.entry_count.fetch_sub(1, AtomicOrdering::Relaxed);
                    Ok(ApplyOutcome::TombstoneRemoved(old_disk))
                }
                None => Ok(ApplyOutcome::Inserted), // no-op tombstone — no dead bytes
            }
        } else {
            let old = sync::lock(&self.indexes[self.shard_for(&key)]).upsert(key, disk);
            match old {
                Some(old_disk) => Ok(ApplyOutcome::Replaced(old_disk)),
                None => {
                    self.entry_count.fetch_add(1, AtomicOrdering::Relaxed);
                    Ok(ApplyOutcome::Inserted)
                }
            }
        }
    }

    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> {
        use crate::replication::ApplyOutcome;

        if raw_after_header.len() < size_of::<K>() + header.value_len as usize {
            return Ok(ApplyOutcome::NotMatched);
        }
        let key = &raw_after_header[..size_of::<K>()];
        let value = &raw_after_header[size_of::<K>()..size_of::<K>() + header.value_len as usize];
        let crc = crate::entry::compute_crc32(header.gsn, header.value_len, key, value);
        if crc != header.crc32 {
            return Ok(ApplyOutcome::NotMatched);
        }
        self.apply_entry(
            shard_inner,
            shard_id,
            file_id,
            entry_offset,
            header,
            key,
            value,
        )
    }

    fn key_len(&self) -> usize {
        size_of::<K>()
    }
}

#[cfg(feature = "replication")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMap<K, H> {
    /// Install SPSC replication producers into every shard and start a
    /// `ReplicationServer` bound to `bind_addr`.
    ///
    /// # Single-call contract
    ///
    /// Each call installs fresh SPSC producers, replacing any previously
    /// installed ones. Call this at most once per `VarMap` instance — a
    /// second call will orphan the in-flight producer of any active streaming
    /// connection on the first server, which will then observe an empty ring
    /// buffer and silently stop forwarding entries.
    pub fn start_replication_server(
        &self,
        bind_addr: std::net::SocketAddr,
        signal: crate::shutdown::ShutdownSignal,
    ) -> crate::error::DbResult<crate::replication::ReplicationServer> {
        let consumers = self.install_replication_producers()?;
        crate::replication::ReplicationServer::start(
            bind_addr,
            self.engine.shards().clone(),
            consumers,
            self.engine.config().max_file_size,
            signal,
        )
    }

    fn install_replication_producers(
        &self,
    ) -> crate::error::DbResult<Vec<rtrb::Consumer<crate::replication::ReplicationEntry>>> {
        const SPSC_CAPACITY: usize = 4096;
        let shards = self.engine.shards();
        let mut consumers = Vec::with_capacity(shards.len());
        for shard in shards.iter() {
            let (p, c) = rtrb::RingBuffer::new(SPSC_CAPACITY);
            shard.set_replication_producer(p);
            consumers.push(c);
        }
        Ok(consumers)
    }

    /// Start a `ReplicationClient` that streams entries from `leader_addr`
    /// into `registry`. Symmetric to [`Self::start_replication_server`].
    ///
    /// `key_len` is derived from `size_of::<K>()`.
    pub fn start_replication_client(
        &self,
        leader_addr: std::net::SocketAddr,
        registry: std::sync::Arc<crate::replication::ReplicationRegistry>,
        signal: crate::shutdown::ShutdownSignal,
    ) -> crate::error::DbResult<crate::replication::ReplicationClient> {
        crate::replication::ReplicationClient::start(
            leader_addr,
            self.engine.shards().clone(),
            registry,
            size_of::<K>() as u16,
            signal,
        )
    }
}

#[cfg(feature = "replication")]
impl<K, H> VarMap<K, H>
where
    K: Key + Send + Sync + Hash + Eq + 'static,
    H: WriteHook<K> + Send + Sync + 'static,
{
    /// Wrap a shared handle to this map as a `Box<dyn ReplicationTarget>`.
    ///
    /// The returned box holds an `Arc` clone — the caller retains full read
    /// access to the original map through the `Arc` while the registry owns
    /// the box. This is the intended pattern for follower-side wiring:
    ///
    /// ```ignore
    /// let follower = Arc::new(VarMap::<[u8; 8]>::open(path, cfg)?);
    /// let registry = ReplicationRegistry::new(follower.as_replication_target());
    /// // `follower` remains usable for .get() etc.
    /// ```
    pub fn as_replication_target(
        self: &std::sync::Arc<Self>,
    ) -> Box<dyn crate::replication::ReplicationTarget> {
        Box::new(std::sync::Arc::clone(self))
    }
}

/// Handle for atomic multi-key operations within a single shard.
/// Obtained via [`VarMap::atomic`]. The shard + index locks are held for the
/// lifetime of this struct — keep the closure short.
pub struct VarMapShard<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    tree: &'a VarMap<K, H>,
    inner: MutexGuard<'a, ShardInner>,
    index: MutexGuard<'a, ShardIndex<K, DiskLoc>>,
    shard_id: usize,
    events: Vec<(K, Option<ByteView>, Option<ByteView>)>,
}

impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMapShard<'_, K, 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: &[u8]) -> DbResult<bool> {
        self.check_shard(key)?;
        let old = if H::NEEDS_OLD_VALUE && H::NEEDS_WRITE {
            if let Some(disk) = self.index.get(key) {
                Some(
                    self.tree
                        .read_value_locked_result(disk, self.shard_id as u8, &self.inner)?,
                )
            } else {
                None
            }
        } else {
            None
        };
        let existed =
            self.tree
                .put_locked(self.shard_id, &mut self.inner, &mut self.index, key, value)?;
        if H::NEEDS_WRITE {
            self.events.push((*key, old, Some(ByteView::from(value))));
        }
        Ok(existed)
    }

    pub fn insert(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
        self.check_shard(key)?;
        self.tree
            .insert_locked(self.shard_id, &mut self.inner, &mut self.index, key, value)?;
        if H::NEEDS_WRITE {
            self.events.push((*key, None, Some(ByteView::from(value))));
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<bool> {
        self.check_shard(key)?;
        let old = if H::NEEDS_OLD_VALUE && H::NEEDS_WRITE {
            if let Some(disk) = self.index.get(key) {
                Some(
                    self.tree
                        .read_value_locked_result(disk, self.shard_id as u8, &self.inner)?,
                )
            } else {
                None
            }
        } else {
            None
        };
        let existed =
            self.tree
                .delete_locked(self.shard_id, &mut self.inner, &mut self.index, key)?;
        if existed && H::NEEDS_WRITE {
            self.events.push((*key, old, None));
        }
        Ok(existed)
    }

    pub fn get(&self, key: &K) -> Option<ByteView> {
        let disk = *self.index.get(key)?;
        self.tree
            .read_value_locked(&disk, self.shard_id as u8, &self.inner)
    }

    pub fn get_or_err(&self, key: &K) -> DbResult<ByteView> {
        let disk = *self.index.get(key).ok_or(DbError::KeyNotFound)?;
        self.tree
            .read_value_locked_result(&disk, self.shard_id as u8, &self.inner)
    }

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

    /// Atomically read-modify-write a key within the shard.
    /// Returns `Some(new_value)` if the key existed, `None` if absent.
    pub fn update(
        &mut self,
        key: &K,
        f: impl FnOnce(&[u8]) -> ByteView,
    ) -> DbResult<Option<ByteView>> {
        self.update_inner(key, |b| Ok(Some(f(b))), false)
    }

    /// 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(&[u8]) -> ByteView,
    ) -> DbResult<Option<ByteView>> {
        self.update_inner(key, |b| Ok(Some(f(b))), true)
    }

    /// Core update primitive used by both `update`/`fetch_update` (public) and
    /// `VarTypedMapShard` (next task). Closure returns `None` to abort the write.
    pub(crate) fn update_inner(
        &mut self,
        key: &K,
        f: impl FnOnce(&[u8]) -> DbResult<Option<ByteView>>,
        return_old: bool,
    ) -> DbResult<Option<ByteView>> {
        self.check_shard(key)?;
        match self.tree.try_update_locked(
            self.shard_id,
            &mut self.inner,
            &mut self.index,
            key,
            f,
        )? {
            UpdateOutcome::Missing => Ok(None),
            UpdateOutcome::Unchanged(current) => Ok(Some(current)),
            UpdateOutcome::Updated { old, new } => {
                if H::NEEDS_WRITE {
                    let event_old = if H::NEEDS_OLD_VALUE {
                        Some(old.clone())
                    } else {
                        None
                    };
                    self.events.push((*key, event_old, Some(new.clone())));
                }
                Ok(Some(if return_old { old } else { new }))
            }
        }
    }

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

// ==========================================================================
// MultiTx — cross-collection transaction support (feature `armour`)
//
// Engine-direct variable map: per shard holds (engine ShardInner, index of
// DiskLocs); reads decode through the shard's `inner`; no seize guard, no sync
// seam; hooks replay from one collection-wide log.
// ==========================================================================

#[cfg(feature = "armour")]
type VarMapGuard<'a, K> = (
    usize,
    MutexGuard<'a, ShardInner>,
    MutexGuard<'a, ShardIndex<K, DiskLoc>>,
);

/// Multi-shard transaction handle for [`VarMap`] inside `Db::atomicN`.
#[cfg(feature = "armour")]
pub struct VarMapTx<'a, K: Key + Send + Sync + Hash + Eq, H: WriteHook<K> = NoHook> {
    tree: &'a VarMap<K, H>,
    shards: Vec<VarMapGuard<'a, K>>,
    log: Vec<(K, Option<ByteView>, Option<ByteView>)>,
}

#[cfg(feature = "armour")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> VarMapTx<'_, K, H> {
    fn position(&self, key: &K) -> DbResult<usize> {
        let sid = self.tree.shard_for(key);
        self.shards
            .iter()
            .position(|(s, _, _)| *s == sid)
            .ok_or(DbError::ShardMismatch)
    }

    pub fn try_get(&self, key: &K) -> DbResult<Option<ByteView>> {
        let i = self.position(key)?;
        let (sid, inner, index) = &self.shards[i];
        let Some(disk) = index.get(key) else {
            return Ok(None);
        };
        Ok(Some(
            self.tree
                .read_value_locked_result(disk, *sid as u8, inner)?,
        ))
    }

    pub fn try_contains(&self, key: &K) -> DbResult<bool> {
        let i = self.position(key)?;
        Ok(self.shards[i].2.contains_key(key))
    }

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

    pub fn put(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
        let i = self.position(key)?;
        let old = if H::NEEDS_OLD_VALUE && H::NEEDS_WRITE {
            let (sid, inner, index) = &self.shards[i];
            if let Some(disk) = index.get(key) {
                Some(
                    self.tree
                        .read_value_locked_result(disk, *sid as u8, inner)?,
                )
            } else {
                None
            }
        } else {
            None
        };
        let (sid, inner, index) = &mut self.shards[i];
        self.tree.put_locked(*sid, inner, index, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, old, Some(ByteView::from(value))));
        }
        Ok(())
    }

    pub fn insert(&mut self, key: &K, value: &[u8]) -> DbResult<()> {
        let i = self.position(key)?;
        let (sid, inner, index) = &mut self.shards[i];
        self.tree.insert_locked(*sid, inner, index, key, value)?;
        if H::NEEDS_WRITE {
            self.log.push((*key, None, Some(ByteView::from(value))));
        }
        Ok(())
    }

    pub fn delete(&mut self, key: &K) -> DbResult<bool> {
        let i = self.position(key)?;
        let old = if H::NEEDS_OLD_VALUE && H::NEEDS_WRITE {
            let (sid, inner, index) = &self.shards[i];
            if let Some(disk) = index.get(key) {
                Some(
                    self.tree
                        .read_value_locked_result(disk, *sid as u8, inner)?,
                )
            } else {
                None
            }
        } else {
            None
        };
        let (sid, inner, index) = &mut self.shards[i];
        let existed = self.tree.delete_locked(*sid, inner, index, key)?;
        if existed && H::NEEDS_WRITE {
            self.log.push((*key, old, None));
        }
        Ok(existed)
    }
}

#[cfg(feature = "armour")]
impl<K: Key + Send + Sync + Hash + Eq, H: WriteHook<K>> crate::armour::multi_tx::MultiTx
    for VarMap<K, H>
{
    type Key = K;
    type Tx<'a>
        = VarMapTx<'a, K, H>
    where
        Self: 'a;

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

    fn begin_tx(&self) -> VarMapTx<'_, K, H> {
        VarMapTx {
            tree: self,
            shards: Vec::new(),
            log: Vec::new(),
        }
    }

    fn lock_shard_into<'a>(&'a self, shard_id: usize, tx: &mut VarMapTx<'a, K, H>) {
        // engine shard inner first, then the per-shard index — matching atomic().
        let inner = self.engine.shards()[shard_id].lock();
        let index = sync::lock(&self.indexes[shard_id]);
        tx.shards.push((shard_id, inner, index));
    }

    fn release_locks(&self, tx: &mut VarMapTx<'_, K, H>) -> crate::armour::multi_tx::SyncNeeds {
        tx.shards.clear(); // drops index + engine guards; 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: VarMapTx<'_, K, H>) {
        if H::NEEDS_WRITE {
            for (k, old, new) in &tx.log {
                self.hook.on_write(k, old.as_deref(), new.as_deref());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Config;
    use crate::compaction::compact_shard;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
    use tempfile::tempdir;

    /// Test hook with counters for on_write / on_init calls and the last observed new value for
    /// on_write. NEEDS_INIT and NEEDS_OLD_VALUE are parameterized by const generics.
    #[derive(Default)]
    struct CountingHook<const NEEDS_INIT: bool, const NEEDS_OLD: bool> {
        writes: AtomicUsize,
        writes_with_old: AtomicUsize,
        inits: AtomicUsize,
        last_write_new: crate::sync::Mutex<Option<Vec<u8>>>,
        last_init_value: crate::sync::Mutex<Option<Vec<u8>>>,
    }

    impl<const NEEDS_INIT: bool, const NEEDS_OLD: bool> WriteHook<[u8; 8]>
        for CountingHook<NEEDS_INIT, NEEDS_OLD>
    {
        const NEEDS_OLD_VALUE: bool = NEEDS_OLD;
        const NEEDS_INIT: bool = NEEDS_INIT;

        fn on_write(&self, _key: &[u8; 8], old: Option<&[u8]>, new: Option<&[u8]>) {
            self.writes.fetch_add(1, AtomicOrdering::Relaxed);
            if old.is_some() {
                self.writes_with_old.fetch_add(1, AtomicOrdering::Relaxed);
            }
            *crate::sync::lock(&self.last_write_new) = new.map(<[u8]>::to_vec);
        }

        fn on_init(&self, _key: &[u8; 8], value: &[u8]) {
            self.inits.fetch_add(1, AtomicOrdering::Relaxed);
            *crate::sync::lock(&self.last_init_value) = Some(value.to_vec());
        }
    }

    fn open_test_map_hooked<const NEEDS_INIT: bool, const NEEDS_OLD: bool>(
        dir: &std::path::Path,
        hook: CountingHook<NEEDS_INIT, NEEDS_OLD>,
    ) -> VarMap<[u8; 8], CountingHook<NEEDS_INIT, NEEDS_OLD>> {
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 8192;
        cfg.write_buffer_size = 8192;
        cfg.compaction_threshold = 0.0;
        VarMap::open_hooked(dir, cfg, hook).expect("open hooked test map")
    }

    fn open_test_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 8192;
        cfg.write_buffer_size = 8192;
        cfg.compaction_threshold = 0.0;
        VarMap::open(dir, cfg).expect("open test map")
    }

    #[test]
    fn open_resolves_hints_true_for_var() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = crate::Config::test_no_hints();
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).unwrap();
        assert_eq!(map.config().hints, Some(true));
    }

    #[test]
    fn put_reports_whether_key_existed() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());
        let key = 42u64.to_be_bytes();

        // Fresh insert -> false.
        assert!(!map.put(&key, b"v1").expect("put"));
        // Overwrite of the same key -> true.
        assert!(map.put(&key, b"v2").expect("put"));
        // Survives a flush (overwrite still detected after entries hit disk).
        map.flush_buffers().expect("flush");
        assert!(map.put(&key, b"v3").expect("put"));
        // A different key is a fresh insert again.
        assert!(!map.put(&7u64.to_be_bytes(), b"x").expect("put"));
    }

    #[test]
    fn for_each_visits_all_entries() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        map.put(&1u64.to_be_bytes(), b"one").unwrap();
        map.put(&2u64.to_be_bytes(), b"two").unwrap();
        map.put(&3u64.to_be_bytes(), b"three").unwrap();

        let mut seen: std::collections::HashMap<[u8; 8], Vec<u8>> =
            std::collections::HashMap::new();
        map.for_each(|key, value| {
            seen.insert(key, value.to_vec());
        });

        assert_eq!(seen.len(), 3);
        assert_eq!(seen[&1u64.to_be_bytes()], b"one".to_vec());
        assert_eq!(seen[&2u64.to_be_bytes()], b"two".to_vec());
        assert_eq!(seen[&3u64.to_be_bytes()], b"three".to_vec());
    }

    fn put_until_compactable(map: &VarMap<[u8; 8]>, key: [u8; 8]) -> DiskLoc {
        // Capture the FIRST put's DiskLoc — file 0 accumulates 100% dead bytes
        // as later overwrites move the live pointer to other files, so
        // compaction at threshold=0.0 erases it. The captured DiskLoc is
        // genuinely stale afterward. Each put is 280 bytes; with
        // max_file_size=8192, ~30 puts cross the first rotation boundary.
        // Use 65 puts so the snap file is sealed as immutable AND fully dead.
        map.put(&key, &[0u8; 256]).expect("first put");
        let snap = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
        for i in 1..65u8 {
            map.put(&key, &[i; 256]).expect("overwrite");
        }
        map.put(&key, b"final-value-payload-XX").expect("final put");
        snap
    }

    #[test]
    fn var_map_len_tracks_mutations() {
        let dir = tempdir().unwrap();
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), Config::balanced().build()).unwrap();
        assert_eq!(map.len(), 0);
        assert!(map.is_empty());

        let k1 = [1u8; 8];
        let k2 = [2u8; 8];
        map.put(&k1, b"a").unwrap();
        assert_eq!(map.len(), 1);

        map.put(&k2, b"b").unwrap();
        assert_eq!(map.len(), 2);

        // Overwrite — len stays the same
        map.put(&k1, b"c").unwrap();
        assert_eq!(map.len(), 2);

        map.delete(&k1).unwrap();
        assert_eq!(map.len(), 1);

        map.delete(&k2).unwrap();
        assert_eq!(map.len(), 0);
        assert!(map.is_empty());

        // Delete non-existent — no change
        map.delete(&k1).unwrap();
        assert_eq!(map.len(), 0);
    }

    #[test]
    fn var_map_len_survives_reopen() {
        let dir = tempdir().unwrap();
        {
            let map: VarMap<[u8; 8]> =
                VarMap::open(dir.path(), Config::balanced().build()).unwrap();
            for i in 0u64..50 {
                map.put(&i.to_le_bytes(), b"val").unwrap();
            }
            assert_eq!(map.len(), 50);
        }
        // Reopen
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), Config::balanced().build()).unwrap();
        assert_eq!(map.len(), 50);
    }

    #[test]
    fn read_value_cached_inner_returns_stale_after_compaction() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 7u64.to_be_bytes();
        let snap = put_until_compactable(&map, key);

        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let _ = compact_shard(shard, &map, 0.0).expect("compaction");

        match map.read_value_cached_inner(&snap, shard_id) {
            Err(DbError::StaleDiskLoc) => {}
            Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
            Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
        }
    }

    #[test]
    fn get_during_compaction_returns_some() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 11u64.to_be_bytes();
        let _snap = put_until_compactable(&map, key);
        let shard = &map.engine.shards()[0];
        let _ = compact_shard(shard, &map, 0.0).expect("compaction");

        let v = map.get(&key).expect("post-compaction get");
        assert_eq!(v.as_bytes(), b"final-value-payload-XX");
    }

    #[test]
    fn get_or_read_block_returns_stale_for_unknown_file_id() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        match map.get_or_read_block(0, 9999, 0) {
            Err(DbError::StaleDiskLoc) => {}
            Ok(_) => panic!("expected StaleDiskLoc, got Ok"),
            Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
        }
    }

    #[test]
    fn retry_limit_returns_none_on_persistent_stale() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 99u64.to_be_bytes();
        map.put(&key, b"payload").expect("put");

        // Force every read_value_cached_inner call to see Stale by manually
        // clearing inner.immutable AND moving active.file_id out of the way.
        // The HashMap still points at the original (now-unreachable) file_id,
        // so each retry hits Shard::read_block which returns StaleDiskLoc.

        {
            let shard_id = map.shard_for(&key) as u8;
            let shard = &map.engine.shards()[shard_id as usize];
            let mut inner = shard.lock();
            inner.immutable = Vec::new();
            inner.active.file_id = u32::MAX;
        }

        assert!(
            map.get(&key).is_none(),
            "MAX_STALE_RETRIES must terminate the retry loop and return None"
        );
    }

    #[test]
    fn var_map_replay_init_fires_on_init_per_live_key_raw() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());

        for i in 0u64..5 {
            map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
        }
        map.delete(&3u64.to_be_bytes()).expect("delete");

        map.hook.writes.store(0, AtomicOrdering::Relaxed);
        map.hook.inits.store(0, AtomicOrdering::Relaxed);

        map.replay_init();

        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 4);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
    }

    #[test]
    fn var_map_replay_init_no_hook_is_noop() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        for i in 0u64..3 {
            map.put(&i.to_be_bytes(), &[i as u8; 8]).expect("put");
        }
        map.replay_init();
        assert!(map.get(&0u64.to_be_bytes()).is_some());
    }

    #[test]
    fn var_map_migrate_keep_fires_on_init_not_on_write_raw() {
        use crate::MigrateAction;
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());

        for i in 0u64..4 {
            map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
        }
        map.hook.writes.store(0, AtomicOrdering::Relaxed);
        map.hook.inits.store(0, AtomicOrdering::Relaxed);

        let mutated = map.migrate(|_, _| MigrateAction::Keep).expect("migrate");

        assert_eq!(mutated, 0);
        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 4);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
    }

    #[test]
    fn var_map_migrate_update_fires_on_init_with_new_value_raw() {
        use crate::MigrateAction;
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());

        let key = 42u64.to_be_bytes();
        map.put(&key, b"old-value").expect("put");
        map.hook.writes.store(0, AtomicOrdering::Relaxed);
        map.hook.inits.store(0, AtomicOrdering::Relaxed);

        let new = ByteView::new(b"new-value");
        let mutated = map
            .migrate(move |_, _| MigrateAction::Update(new.clone()))
            .expect("migrate");

        assert_eq!(mutated, 1);
        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 1);
        let g = crate::sync::lock(&map.hook.last_init_value);
        assert_eq!(g.as_deref(), Some(b"new-value".as_ref()));
        drop(g);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
        assert_eq!(map.get(&key).unwrap().as_bytes(), b"new-value");
    }

    #[test]
    fn var_map_migrate_delete_fires_no_hooks_raw() {
        use crate::MigrateAction;
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());

        let key = 7u64.to_be_bytes();
        map.put(&key, b"x").expect("put");
        map.hook.writes.store(0, AtomicOrdering::Relaxed);
        map.hook.inits.store(0, AtomicOrdering::Relaxed);

        let mutated = map.migrate(|_, _| MigrateAction::Delete).expect("migrate");

        assert_eq!(mutated, 1);
        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
        assert!(map.get(&key).is_none());
    }

    #[test]
    fn var_map_public_put_still_fires_on_write_once() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, false>(dir.path(), CountingHook::default());

        map.put(&1u64.to_be_bytes(), b"v").expect("put");
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
    }

    #[test]
    fn var_map_migrate_no_init_hook_is_silent_for_keep_and_update() {
        use crate::MigrateAction;
        let dir = tempdir().unwrap();
        // NEEDS_INIT = false: on_init must never be called.
        let map = open_test_map_hooked::<false, false>(dir.path(), CountingHook::default());

        for i in 0u64..3 {
            map.put(&i.to_be_bytes(), &[i as u8; 16]).expect("put");
        }
        map.hook.writes.store(0, AtomicOrdering::Relaxed);
        map.hook.inits.store(0, AtomicOrdering::Relaxed);

        map.migrate(|_, _| MigrateAction::Keep)
            .expect("migrate keep");
        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);

        let new = ByteView::new(b"new");
        map.migrate(move |_, _| MigrateAction::Update(new.clone()))
            .expect("migrate update");
        assert_eq!(map.hook.inits.load(AtomicOrdering::Relaxed), 0);
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 0);
    }

    #[test]
    fn var_map_atomic_fires_hooks() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
        let key = 1u64.to_be_bytes();
        map.atomic(&key, |shard| {
            shard.put(&key, b"a")?; // new key
            shard.delete(&key)?; // delete, old=b"a"
            Ok(())
        })
        .expect("atomic");
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
        assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1); // delete saw old
    }

    #[test]
    fn var_map_atomic_fires_for_applied_on_err() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
        let key = 1u64.to_be_bytes();
        let r: DbResult<()> = map.atomic(&key, |shard| {
            shard.put(&key, b"x")?;
            Err(DbError::KeyNotFound)
        });
        assert!(r.is_err());
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
    }

    #[test]
    fn var_map_atomic_same_key_old_from_disk() {
        // NEEDS_OLD_VALUE=true: the old value of the 2nd put must be the 1st put's value,
        // read back through the write buffer under the lock.
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<false, true>(dir.path(), CountingHook::default());
        let key = 1u64.to_be_bytes();
        map.atomic(&key, |shard| {
            shard.put(&key, b"first")?;
            shard.put(&key, b"second")?;
            Ok(())
        })
        .expect("atomic");
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
        assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1); // 2nd put saw old
        assert_eq!(
            crate::sync::lock(&map.hook.last_write_new).clone(),
            Some(b"second".to_vec())
        );
    }

    #[test]
    fn var_map_shard_put_returns_existed() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
        let key = 1u64.to_be_bytes();
        let (fresh, overwrite) = map
            .atomic(&key, |shard| {
                let fresh = shard.put(&key, b"a")?;
                let overwrite = shard.put(&key, b"b")?;
                Ok((fresh, overwrite))
            })
            .expect("atomic");
        assert!(!fresh);
        assert!(overwrite);
    }

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

        let k = [1u8; 8];
        map.put(&k, b"hello").unwrap();

        assert!(matches!(
            map.compare_delete(&k, b"WRONG"),
            Err(DbError::CasMismatch)
        ));
        assert_eq!(map.get(&k).as_deref(), Some(&b"hello"[..]));

        assert!(map.compare_delete(&k, b"hello").is_ok());
        assert!(map.get(&k).is_none());

        assert!(matches!(
            map.compare_delete(&k, b"hello"),
            Err(DbError::KeyNotFound)
        ));
    }

    #[test]
    fn compare_delete_hook_needs_old_true_gets_old() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<false, true>(dir.path(), CountingHook::default());

        let k = [2u8; 8];
        map.put(&k, b"payload").unwrap();
        let with_old_before = map.hook.writes_with_old.load(AtomicOrdering::Relaxed);

        map.compare_delete(&k, b"payload").unwrap();

        assert_eq!(
            map.hook.writes_with_old.load(AtomicOrdering::Relaxed) - with_old_before,
            1
        );
    }

    #[test]
    fn compare_delete_hook_needs_old_false_gets_none() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<false, false>(dir.path(), CountingHook::default());

        let k = [3u8; 8];
        map.put(&k, b"payload").unwrap();
        let writes_before = map.hook.writes.load(AtomicOrdering::Relaxed);
        let with_old_before = map.hook.writes_with_old.load(AtomicOrdering::Relaxed);

        map.compare_delete(&k, b"payload").unwrap();

        assert_eq!(
            map.hook.writes.load(AtomicOrdering::Relaxed) - writes_before,
            1
        );
        assert_eq!(
            map.hook.writes_with_old.load(AtomicOrdering::Relaxed) - with_old_before,
            0
        );
    }

    /// `_result` returns Ok when the value sits in the active write buffer.
    #[test]
    fn read_value_locked_result_ok_from_write_buf() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 1u64.to_be_bytes();
        let payload = b"in-write-buffer-value";
        map.put(&key, payload).expect("put");

        let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let inner = shard.lock();
        let v = map
            .read_value_locked_result(&disk, shard_id, &inner)
            .expect("write-buf read must succeed");
        assert_eq!(v.as_bytes(), payload);
    }

    /// `_result` returns Ok when the value is on an immutable file and the
    /// entry sits within a single 4 KiB block.
    #[test]
    fn read_value_locked_result_ok_from_disk_immutable() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 2u64.to_be_bytes();
        let payload = b"single-block-immutable";
        map.put(&key, payload).expect("put");
        // Each put is 280 bytes; with max_file_size=8192, need >30 puts to
        // cross the rotation boundary. Use keys 100..135 to avoid overwriting key 2.
        for i in 100u64..135 {
            map.put(&i.to_be_bytes(), &[i as u8; 256]).expect("rotator");
        }

        let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
        // `open_test_map` uses default CacheConfig (max_size = 0) — cache is
        // disabled, so step 2 returns None and the read exercises step 3 (disk).

        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let inner = shard.lock();
        assert_ne!(
            disk.file_id, inner.active.file_id,
            "test setup failed: key=2 entry is still in the active file's write buffer",
        );
        let v = map
            .read_value_locked_result(&disk, shard_id, &inner)
            .expect("disk read must succeed");
        assert_eq!(v.as_bytes(), payload);
    }

    /// `_result` propagates `StaleDiskLoc` from the disk read.
    #[test]
    fn read_value_locked_result_propagates_stale_disk_loc() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 3u64.to_be_bytes();
        let snap = put_until_compactable(&map, key);

        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let _ = compact_shard(shard, &map, 0.0).expect("compaction");

        let inner = shard.lock();
        match map.read_value_locked_result(&snap, shard_id, &inner) {
            Err(DbError::StaleDiskLoc) => {}
            Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
            Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
        }
    }

    /// Backward-compat: Option wrapper still returns None on StaleDiskLoc.
    /// No tracing assertion: the None is the load-bearing invariant.
    #[test]
    fn read_value_locked_returns_none_on_stale() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());

        let key = 4u64.to_be_bytes();
        let snap = put_until_compactable(&map, key);

        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let _ = compact_shard(shard, &map, 0.0).expect("compaction");

        let inner = shard.lock();
        assert!(map.read_value_locked(&snap, shard_id, &inner).is_none());
    }

    /// Pins the "size check sits before cache lookup" invariant for VarMap.
    /// Without the early `start + len > 8192` check, the cached path would
    /// hit `extract_from_block` and panic on `next[..second_len]`.
    #[test]
    fn large_value_with_first_block_cached_uses_fallback() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        // Enable the block cache so `cache.insert` actually takes effect
        // (default CacheConfig has max_size = 0, making the cache a no-op).
        cfg.block_cache.max_size = 1 << 20;
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");

        let key = 42u64.to_be_bytes();
        let payload: Vec<u8> = (0..20_000u32).map(|i| i as u8).collect();
        map.put(&key, &payload).expect("put large");
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");

        let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
        let start = (disk.offset & 4095) as usize;
        assert!(
            start + disk.len as usize > 8192,
            "test precondition: large value must span >2 blocks",
        );

        let block_offset = disk.offset as u64 & !4095;
        let shard_id = map.shard_for(&key) as u8;
        let cache_key = BlockKey {
            shard_id,
            file_id: disk.file_id,
            block_offset,
        };
        map.block_cache
            .insert(cache_key, Arc::new(AlignedBuf::zeroed(4096)));
        assert!(
            map.block_cache.get(&cache_key).is_some(),
            "cache must be enabled for warm-cache scenario",
        );

        let v = map
            .read_value_cached_inner(&disk, shard_id)
            .expect("large value must read via locked fallback");
        assert_eq!(v.as_bytes(), payload.as_slice());
    }

    #[test]
    fn extract_from_block_single_block() {
        let mut block = AlignedBuf::zeroed(4096);
        for (i, byte) in block.iter_mut().enumerate() {
            *byte = i as u8;
        }
        let v = VarMap::<[u8; 8]>::extract_from_block(&block, 100, 50, || {
            panic!("next_block must not be called for single-block reads")
        })
        .expect("ok");
        let expected: Vec<u8> = (100u8..150u8).collect();
        assert_eq!(v.as_bytes(), expected.as_slice());
    }

    #[test]
    fn extract_from_block_two_blocks_exact() {
        let mut first = AlignedBuf::zeroed(4096);
        for byte in first.iter_mut() {
            *byte = 0xAA;
        }
        let mut second = AlignedBuf::zeroed(4096);
        for byte in second.iter_mut() {
            *byte = 0xBB;
        }
        let v = VarMap::<[u8; 8]>::extract_from_block(&first, 4095, 4097, || Ok(Arc::new(second)))
            .expect("ok");
        let bytes = v.as_bytes();
        assert_eq!(bytes.len(), 4097);
        assert_eq!(bytes[0], 0xAA);
        assert_eq!(bytes[1], 0xBB);
        assert_eq!(bytes[4096], 0xBB);
    }

    #[test]
    fn extract_from_block_two_blocks_partial() {
        let mut first = AlignedBuf::zeroed(4096);
        for byte in first.iter_mut() {
            *byte = 0x11;
        }
        let mut second = AlignedBuf::zeroed(4096);
        for byte in second.iter_mut() {
            *byte = 0x22;
        }
        let v = VarMap::<[u8; 8]>::extract_from_block(&first, 4000, 200, || Ok(Arc::new(second)))
            .expect("ok");
        let bytes = v.as_bytes();
        assert_eq!(bytes.len(), 200);
        assert!(bytes[..96].iter().all(|b| *b == 0x11));
        assert!(bytes[96..].iter().all(|b| *b == 0x22));
    }

    fn open_large_value_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        VarMap::open(dir, cfg).expect("open large-value test map")
    }

    fn build_large_payload(seed: u8) -> Vec<u8> {
        (0..50_000u32)
            .map(|i| (i as u8).wrapping_add(seed))
            .collect()
    }

    /// `open_large_value_map` uses default CacheConfig (max_size = 0), so
    /// the block cache is disabled and the read goes through the large-value
    /// locked fallback → disk read.
    #[test]
    fn large_value_read_via_locked_fallback() {
        let dir = tempdir().unwrap();
        let map = open_large_value_map(dir.path());

        let key = 100u64.to_be_bytes();
        let payload = build_large_payload(0);
        map.put(&key, &payload).expect("put large");
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");

        let v = map
            .get(&key)
            .expect("read must succeed via locked fallback");
        assert_eq!(v.as_bytes(), payload.as_slice());
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn varmap_strict_reads_and_mutators_surface_read_fault() {
        use crate::test_faults::{big_value, corrupt_tags};
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.encryption_key = Some([7u8; 32]);
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");

        let key = 1u64.to_be_bytes();
        map.put(&key, &big_value(0xCD)).unwrap();
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");
        corrupt_tags(&map.engine.shard_dirs()[0]);

        assert!(map.get(&key).is_none());
        assert!(matches!(
            map.get_or_err(&key),
            Err(DbError::EncryptionError(_))
        ));
        assert!(matches!(
            map.try_get(&key),
            Err(DbError::EncryptionError(_))
        ));
        assert!(matches!(
            map.cas(&key, b"x", b"y"),
            Err(DbError::EncryptionError(_))
        ));
        assert!(matches!(
            map.compare_delete(&key, b"x"),
            Err(DbError::EncryptionError(_))
        ));
        assert!(matches!(
            map.update(&key, |_| ByteView::new(b"z")),
            Err(DbError::EncryptionError(_))
        ));

        let absent = 999u64.to_be_bytes();
        assert!(matches!(map.try_get(&absent), Ok(None)));
        assert!(matches!(map.get_or_err(&absent), Err(DbError::KeyNotFound)));
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn large_value_read_encrypted() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.encryption_key = Some([7u8; 32]);
        // Cache stays at default (max_size = 0, disabled), so the locked
        // fallback always routes through pread_value_encrypted.
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");

        let key = 101u64.to_be_bytes();
        let payload = build_large_payload(0xAB);
        map.put(&key, &payload).expect("put encrypted large");
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");

        let v = map
            .get(&key)
            .expect("encrypted large read must succeed via pread_value_encrypted");
        assert_eq!(v.as_bytes(), payload.as_slice());
    }

    #[test]
    fn large_value_stale_disk_loc_deterministic() {
        let dir = tempdir().unwrap();
        let map = open_large_value_map(dir.path());

        let key = 102u64.to_be_bytes();
        let payload = build_large_payload(0x42);
        map.put(&key, &payload).expect("first large put");

        let snap = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");

        // Rotate the active file so the large entry moves to an immutable file.
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate after large put");

        // Overwrite many times to move the live pointer off the original file
        // (making it 100% dead so compaction erases it).
        for i in 1..20u8 {
            map.put(&key, &[i; 256]).expect("overwrite");
        }
        map.put(&key, b"live-after-compaction").expect("final put");

        let shard_id = map.shard_for(&key) as u8;
        let shard = &map.engine.shards()[shard_id as usize];
        let _ = compact_shard(shard, &map, 0.0).expect("compaction");

        match map.read_value_cached_inner(&snap, shard_id) {
            Err(DbError::StaleDiskLoc) => {}
            Ok(v) => panic!("expected StaleDiskLoc, got Ok({:?})", v.as_bytes()),
            Err(e) => panic!("expected StaleDiskLoc, got Err({e})"),
        }

        let v = map
            .get(&key)
            .expect("public path must retry and return live value");
        assert_eq!(v.as_bytes(), b"live-after-compaction");
    }

    /// `_result` step 2: single-block cache fast path.
    #[test]
    fn read_value_locked_result_ok_from_cache_single_block() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 8192;
        cfg.write_buffer_size = 8192;
        cfg.compaction_threshold = 0.0;
        cfg.block_cache.max_size = 1 << 20;
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");

        let key = 9u64.to_be_bytes();
        let payload = b"small-single-block-value";
        map.put(&key, payload).expect("put");
        // Need enough writes to push write_offset past max_file_size=8192.
        for i in 100u64..135 {
            map.put(&i.to_be_bytes(), &[i as u8; 256]).expect("rotator");
        }

        let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");

        let start = (disk.offset & 4095) as usize;
        let len = disk.len as usize;
        assert!(
            start + len <= 4096,
            "test precondition: value must fit in a single block",
        );

        // Pre-warm the cache BEFORE acquiring the shard lock.
        // get_or_read_block → shard.read_block → sync::lock(&inner), so it
        // must not be called while the shard lock is already held.
        let block_offset = disk.offset as u64 & !4095;
        let shard_id = map.shard_for(&key) as u8;
        let cache_key = BlockKey {
            shard_id,
            file_id: disk.file_id,
            block_offset,
        };
        let block = map
            .get_or_read_block(shard_id, disk.file_id, block_offset)
            .expect("read block");
        map.block_cache.insert(cache_key, block);
        assert!(
            map.block_cache.get(&cache_key).is_some(),
            "cache must contain the block for step 2 to fire",
        );

        let shard = &map.engine.shards()[shard_id as usize];
        let inner = shard.lock();
        assert_ne!(
            disk.file_id, inner.active.file_id,
            "test setup failed: key entry is still in the active write buffer",
        );

        let v = map
            .read_value_locked_result(&disk, shard_id, &inner)
            .expect("cache-hit read must succeed");
        assert_eq!(v.as_bytes(), payload);
    }

    /// End-to-end roundtrip with file_id above u16::MAX.
    ///
    /// Uses write_buffer_size=128 KiB so a single 512-byte entry fits in one
    /// flush cycle; max_file_size=4096 forces per-entry rotation; cache
    /// enabled so the fast path is also exercised.  After bumping
    /// next_file_id to 70_000 and rotating, the active file gets an id that
    /// exceeds u16::MAX.  The entry is flushed + rotated to make it
    /// immutable, then `get` must return the correct bytes via the locked
    /// disk-read path.  Before the file_id u32 widening this would either
    /// panic or return wrong bytes because the high bits were truncated.
    #[test]
    fn var_map_get_with_file_id_above_u16() {
        let dir = tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.block_cache.max_size = 1 << 20;
        let map: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open");

        let shard = &map.engine.shards()[0];

        // Bump file id well past u16::MAX and rotate so the next active file
        // carries the new id.
        shard.set_next_file_id(70_000);
        shard.rotate_active_for_test(8).expect("first rotate");
        assert!(
            shard.active_file_id() >= 70_000,
            "active_file_id should be >= 70_000 after rotation"
        );

        let key = 42u64.to_be_bytes();
        let value = vec![0xC3u8; 512];
        map.put(&key, &value).expect("put");

        // Confirm the DiskLoc has file_id > u16::MAX.
        {
            let disk = *sync::lock(&map.indexes[0]).get(&key).expect("indexed");
            assert!(
                disk.file_id > u16::MAX as u32,
                "DiskLoc.file_id must be above u16::MAX, got {}",
                disk.file_id,
            );
        }

        // Flush the write buffer and rotate so the entry lands on an
        // immutable file — reads will go through Shard::read_block.
        shard.flush().expect("flush");
        shard.rotate_active_for_test(8).expect("second rotate");

        let got = map.get(&key).expect("get must return Some");
        assert_eq!(got.as_bytes(), value.as_slice());
    }

    #[test]
    fn var_map_iter_reads_values_after_unlock() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(true)
            .build();
        let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
        m.put(&[1, 0], b"alpha").unwrap();
        m.put(&[2, 0], b"beta").unwrap();
        let got: Vec<(u8, Vec<u8>)> = m
            .iter_view()
            .unwrap()
            .iter()
            .map(|(k, v)| (k[0], v.as_ref().to_vec()))
            .collect();
        assert_eq!(got, [(2, b"beta".to_vec()), (1, b"alpha".to_vec())]);
    }

    #[test]
    fn var_map_iter_view_none_unless_iterable() {
        let dir = tempfile::tempdir().unwrap();
        let m = open_test_map(dir.path());
        assert!(m.iter_view().is_none());
    }

    #[test]
    fn var_map_keys_only_skips_value_read() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = Config::balanced()
            .shard_count(2)
            .hints(true)
            .iterable(true)
            .build();
        let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
        m.put(&[1, 0], b"alpha").unwrap();
        m.put(&[2, 0], b"beta").unwrap();
        // keys() reads only the companion BTreeSet — no off-disk value read.
        let keys: Vec<u8> = m.iter_view().unwrap().keys().map(|k| k[0]).collect();
        assert_eq!(keys, vec![2, 1], "DESC keys");
        // keys_paginate carries keys only.
        let page = m.iter_view().unwrap().keys_paginate(None, 1);
        assert_eq!(page.keys, vec![[2, 0]]);
        assert_eq!(page.next, Some([2, 0]));
    }

    #[test]
    fn var_map_retain_deletes_rejected_entries() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = Config::balanced().shard_count(1).iterable(true).build();
        let m = VarMap::<[u8; 2]>::open(dir.path(), cfg).unwrap();
        m.put(&[1, 0], b"odd").unwrap();
        m.put(&[2, 0], b"even").unwrap();
        m.put(&[3, 0], b"odd2").unwrap();
        m.iter_view()
            .unwrap()
            .retain(|_k, v| v.as_ref() == b"even")
            .unwrap();
        let got: Vec<u8> = m.iter_view().unwrap().iter().map(|(k, _)| k[0]).collect();
        assert_eq!(got, vec![2]);
        assert_eq!(m.len(), 1);
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn var_map_iter_skips_unreadable_key() {
        use crate::test_faults::{big_value, corrupt_tags};
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.encryption_key = Some([7u8; 32]);
        cfg.iterable = true;
        let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");

        let k1 = 1u64.to_be_bytes();
        let k2 = 2u64.to_be_bytes();
        m.put(&k2, &big_value(0x22)).unwrap();
        m.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");
        corrupt_tags(&m.engine.shard_dirs()[0]);
        // k2 is on an immutable encrypted file (tags corrupted). k1 stays in the
        // active write buffer and remains readable without touching disk tags.
        m.put(&k1, b"readable").unwrap();

        let got: Vec<u8> = m.iter_view().unwrap().iter().map(|(k, _)| k[7]).collect();
        assert_eq!(got, vec![1]);
        assert_eq!(m.len(), 2, "unreadable entry must not be deleted by iter");
        assert!(matches!(m.try_get(&k2), Err(DbError::EncryptionError(_))));

        // keys() never reads values, so the unreadable k2 still appears — the
        // whole point of keys-only iteration for maintenance scans.
        let mut keys: Vec<u8> = m.iter_view().unwrap().keys().map(|k| k[7]).collect();
        keys.sort_unstable();
        assert_eq!(keys, vec![1, 2], "keys-only returns the unreadable key too");
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn var_map_retain_skips_unreadable_entry() {
        use crate::test_faults::{big_value, corrupt_tags};
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.encryption_key = Some([7u8; 32]);
        cfg.iterable = true;
        let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");

        let k1 = 1u64.to_be_bytes();
        let k2 = 2u64.to_be_bytes();
        m.put(&k2, &big_value(0x22)).unwrap();
        m.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");
        corrupt_tags(&m.engine.shard_dirs()[0]);
        m.put(&k1, b"keep").unwrap();

        m.iter_view().unwrap().retain(|_k, _v| false).unwrap();
        // k1 was readable and deleted; k2 was skipped (unreadable) and remains indexed.
        assert_eq!(m.len(), 1);
        assert!(m.get(&k1).is_none());
        assert!(matches!(m.try_get(&k2), Err(DbError::EncryptionError(_))));
    }

    /// Regression: when a whole `next_batch` (≥ the internal batch size of 64) is
    /// unreadable, the iterator must keep scanning past it instead of treating the
    /// empty batch as shard exhaustion and dropping every later key.
    #[cfg(feature = "encryption")]
    #[test]
    fn var_map_iter_advances_past_full_unreadable_batch() {
        use crate::test_faults::{big_value, corrupt_tags};
        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 4 * 1024 * 1024;
        cfg.write_buffer_size = 4 * 1024 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.encryption_key = Some([9u8; 32]);
        cfg.iterable = true;
        let m: VarMap<[u8; 8]> = VarMap::open(dir.path(), cfg).expect("open enc");

        // 65 keys (2..=66) — more than one full iteration batch (64) — all made
        // unreadable. They occupy the TOP of DESC order.
        for i in 2u64..=66 {
            m.put(&i.to_be_bytes(), &big_value(i as u8)).unwrap();
        }
        m.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");
        corrupt_tags(&m.engine.shard_dirs()[0]);

        // One readable key that sorts LAST in DESC order, after the unreadable run.
        let readable = 1u64.to_be_bytes();
        m.put(&readable, b"alive").unwrap();

        // Pre-fix: the first 64-key batch was entirely unreadable, so next_batch
        // returned empty, the shard was wrongly marked exhausted, and key 1 was lost.
        let got: Vec<[u8; 8]> = m.iter_view().unwrap().iter().map(|(k, _)| k).collect();
        assert_eq!(
            got,
            vec![readable],
            "iter must scan past a full unreadable batch"
        );

        // keys-only never reads values, so all 66 keys are present.
        assert_eq!(m.iter_view().unwrap().keys().count(), 66);
    }

    // -- Value cache tests -------------------------------------------------------

    fn open_value_cache_map(dir: &std::path::Path) -> VarMap<[u8; 8]> {
        let mut cfg = Config::test();
        cfg.shard_count = 1;
        cfg.max_file_size = 128 * 1024;
        cfg.write_buffer_size = 128 * 1024;
        cfg.compaction_threshold = 0.0;
        cfg.value_cache.max_size = 4 << 20;
        VarMap::open(dir, cfg).expect("open value-cache test map")
    }

    fn large_payload(seed: u8) -> Vec<u8> {
        (0..50_000u32)
            .map(|i| (i as u8).wrapping_add(seed))
            .collect()
    }

    #[test]
    fn map_large_value_miss_then_cached() {
        let dir = tempdir().unwrap();
        let map = open_value_cache_map(dir.path());
        let key = 1u64.to_be_bytes();
        let payload = large_payload(0x33);
        map.put(&key, &payload).unwrap();
        // First read: cache miss → value read from disk, then inserted into value_cache.
        assert_eq!(map.get(&key).expect("first get").as_bytes(), &payload[..]);
        let disk = {
            let idx = sync::lock(&map.indexes[0]);
            *idx.get(&key).unwrap()
        };
        let vkey = crate::value_cache::ValueKey {
            shard_id: 0,
            file_id: disk.file_id,
            offset: disk.offset,
        };
        assert!(
            map.value_cache.get(&vkey).is_some(),
            "value cached after read"
        );
        // Second read: cache hit.
        assert_eq!(map.get(&key).expect("second get").as_bytes(), &payload[..]);
    }

    #[test]
    fn map_large_value_read_then_edit_returns_new() {
        let dir = tempdir().unwrap();
        let map = open_value_cache_map(dir.path());
        let key = 7u64.to_be_bytes();
        let v1 = large_payload(0x11);
        let v2 = large_payload(0x22);
        map.put(&key, &v1).unwrap();
        assert_eq!(map.get(&key).expect("get v1").as_bytes(), &v1[..]);
        // Overwrite: new DiskLoc → old cache entry is stale; try_get sees new loc
        // under the index lock, so the old cached bytes are never returned.
        map.put(&key, &v2).unwrap();
        assert_eq!(map.get(&key).expect("get v2").as_bytes(), &v2[..]);
    }

    #[test]
    fn map_large_value_read_then_compact_returns_live() {
        let dir = tempdir().unwrap();
        let map = open_value_cache_map(dir.path());
        let key = 102u64.to_be_bytes();
        let payload = large_payload(0x42);
        map.put(&key, &payload).unwrap();
        assert_eq!(map.get(&key).expect("initial get").as_bytes(), &payload[..]);
        map.engine.shards()[0]
            .rotate_active_for_test(8)
            .expect("rotate");
        for i in 1..20u8 {
            map.put(&key, &[i; 256]).unwrap();
        }
        map.put(&key, b"live-after-compaction").unwrap();
        let _ = compact_shard(&map.engine.shards()[0], &map, 0.0).expect("compaction");
        assert_eq!(
            map.get(&key).expect("post-compaction get").as_bytes(),
            b"live-after-compaction"
        );
    }

    #[test]
    fn var_map_shard_update_and_fetch_update() {
        let dir = tempdir().unwrap();
        let map = open_test_map(dir.path());
        let key = 1u64.to_be_bytes();
        map.put(&key, b"a").unwrap();
        let (upd, fetched, missing) = map
            .atomic(&key, |shard| {
                let upd = shard.update(&key, |old| {
                    let mut v = old.to_vec();
                    v.push(b'b');
                    ByteView::from(v.as_slice())
                })?;
                let fetched = shard.fetch_update(&key, |old| {
                    let mut v = old.to_vec();
                    v.push(b'c');
                    ByteView::from(v.as_slice())
                })?;
                let missing = shard.update(&2u64.to_be_bytes(), |old| ByteView::from(old))?;
                Ok((upd, fetched, missing))
            })
            .expect("atomic");
        assert_eq!(upd.as_deref(), Some(&b"ab"[..]));
        assert_eq!(fetched.as_deref(), Some(&b"ab"[..]));
        assert!(missing.is_none());
        assert_eq!(map.get(&key).as_deref(), Some(&b"abc"[..]));
    }

    #[test]
    fn var_map_atomic_update_fires_hook() {
        let dir = tempdir().unwrap();
        let map = open_test_map_hooked::<true, true>(dir.path(), CountingHook::default());
        let key = 1u64.to_be_bytes();
        map.atomic(&key, |shard| {
            shard.put(&key, b"hello")?;
            shard.update(&key, |old| {
                let mut v = old.to_vec();
                v.push(b'!');
                ByteView::from(v.as_slice())
            })?;
            Ok(())
        })
        .expect("atomic");
        // put fires hook(old=None, new=Some); update fires hook(old=Some, new=Some)
        assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 2);
        assert_eq!(map.hook.writes_with_old.load(AtomicOrdering::Relaxed), 1);
        assert_eq!(map.get(&key).as_deref(), Some(&b"hello!"[..]));
    }
}
