armdb 0.1.14

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

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

/// A tree with fixed-size keys and typed values `T`. Values are encoded via
/// a [`Codec`] and stored on disk (variable length), with a `BlockCache` for
/// reads — unlike [`TypedTree`](crate::TypedTree) which keeps values in memory.
///
/// Thin wrapper around [`VarTree<K, VarTypedHookAdapter<K, T, C, H>>`].
/// Each `VarTypedTree` owns its storage engine — one tree = one database directory.
///
/// # When to use
///
/// - Typed values too large to keep in RAM → `VarTypedTree` (disk-resident).
/// - Typed values that fit in RAM → [`TypedTree`](crate::TypedTree) (in-memory).
/// - Fixed-size zerocopy values → [`ZeroTree`](crate::ZeroTree).
///
/// # Error handling
///
/// Follows `VarTree` / `TypedTree` precedent (variant A):
/// - `get` / `first` / `last` → `None` if key missing **or** decode fails.
/// - Iterators skip entries with decode errors (`continue` in `next()`).
/// - `migrate` treats decode errors as `Keep` (logged via `tracing::warn!`).
///
/// # Write hooks
///
/// Uses [`TypedWriteHook<K, T>`] via [`VarTypedHookAdapter`]. The hook receives
/// `&T` directly; the adapter decodes raw bytes via the codec. `on_write` fires
/// on `put`/`insert`/`delete`/`cas`/`update`. Does **not** fire inside `atomic()`.
///
/// # Usage
///
/// ```ignore
/// let tree = VarTypedTree::<[u8; 16], Message, RapiraCodec>::open(
///     "data/messages",
///     Config::default(),
///     RapiraCodec,
/// )?;
/// tree.put(&key, &msg)?;
/// if let Some(m) = tree.get(&key) {
///     println!("{:?}", m);
/// }
/// tree.close()?;
/// ```
///
/// # Iteration
///
/// `iter()`, `range()`, and `prefix_iter()` return [`VarTypedIter`] with
/// `Item = (K, T)`. Each `next()` / `next_back()` may perform disk I/O on a
/// block-cache miss and performs a decode.
pub struct VarTypedTree<
    K: Key,
    T: Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T> = NoHook,
> {
    inner: VarTree<K, VarTypedHookAdapter<K, T, C, H>>,
    codec: C,
    _marker: PhantomData<fn() -> T>,
}

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

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

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

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

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

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

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

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

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

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

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

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

    /// Get and decode a value by key.
    /// Returns `None` if the key is missing or decoding fails.
    pub fn get(&self, key: &K) -> Option<T> {
        let bytes = self.inner.get(key)?;
        self.codec.decode_from(&bytes).ok()
    }

    /// Get a value by key, returning `Err(KeyNotFound)` if absent or decode fails.
    pub fn get_or_err(&self, key: &K) -> DbResult<T> {
        self.get(key).ok_or(DbError::KeyNotFound)
    }

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

    /// Return the first entry in index order, or `None` if empty or decode fails.
    pub fn first(&self) -> Option<(K, T)> {
        let (k, bytes) = self.inner.first()?;
        let v = self.codec.decode_from(&bytes).ok()?;
        Some((k, v))
    }

    /// Return the last entry in index order, or `None` if empty or decode fails.
    pub fn last(&self) -> Option<(K, T)> {
        let (k, bytes) = self.inner.last()?;
        let v = self.codec.decode_from(&bytes).ok()?;
        Some((k, v))
    }

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

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

    /// Insert a key-value pair only if the key does not exist.
    /// Returns `Err(KeyExists)` if the key is already present.
    pub fn insert(&self, key: &K, value: &T) -> DbResult<()> {
        let mut buf = Vec::new();
        self.codec.encode_to(value, &mut buf);
        self.inner.insert(key, &buf)
    }

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

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

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

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

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

    /// Atomically execute multiple operations on a single shard.
    /// All keys must route to the same shard as `shard_key`.
    /// The closure must be short — shard lock is held for its duration.
    /// Hook is **not** fired inside this block.
    pub fn atomic<R>(
        &self,
        shard_key: &K,
        f: impl FnOnce(&mut VarTypedShard<'_, K, T, C, H>) -> DbResult<R>,
    ) -> DbResult<R> {
        self.inner.atomic(shard_key, |var_shard| {
            // SAFETY: erase VarShard lifetime via `*mut ()` — see VarTypedShard doc.
            // The pointer is valid for the duration of this closure.
            let inner_ptr: *mut () = (var_shard as *mut VarShard<'_, _, _>).cast();
            let mut shard = VarTypedShard {
                tree: self,
                inner: inner_ptr,
                _marker: PhantomData,
            };
            f(&mut shard)
        })
    }

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

    /// Iterate entries whose keys start with `prefix`.
    ///
    /// `reversed=true` (default): yields matching keys in DESC order.
    /// `next()` is O(1), `next_back()` is O(log n). Disk I/O on cache miss.
    /// Entries with decode errors are skipped.
    pub fn prefix_iter(&self, prefix: &[u8]) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.prefix_iter(prefix),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

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

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    pub fn range(&self, start: &K, end: &K) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.range(start, end),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> VarTypedIter<'_, K, T, C, H> {
        VarTypedIter {
            inner: self.inner.range_bounds(start, end),
            codec: &self.codec,
            _marker: PhantomData,
        }
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    fn next(&mut self) -> Option<Self::Item> {
        for (k, bytes) in self.inner.by_ref() {
            if let Ok(v) = self.codec.decode_from(&bytes) {
                return Some((k, v));
            }
        }
        None
    }
}

impl<K: Key, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>> DoubleEndedIterator
    for VarTypedIter<'_, K, T, C, H>
{
    fn next_back(&mut self) -> Option<Self::Item> {
        while let Some((k, bytes)) = self.inner.next_back() {
            if let Ok(v) = self.codec.decode_from(&bytes) {
                return Some((k, v));
            }
        }
        None
    }
}

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

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

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

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

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

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

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

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

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

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

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

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