haematite 0.6.1

use std::collections::BTreeMap;
use std::fmt;
use std::sync::{Arc, Mutex, MutexGuard};

use crate::store::NodeStore;
use crate::tree::{Cursor, Hash, TreeError};
use crate::wal::{LookupResult, WalBuffer};

use super::operation_error::BranchGetError;
use super::registry::BranchRegistryGuard;
use super::time::Timestamp;

/// Shard identifier used by the branch handle.
///
/// Re-exported from the platform-neutral [`crate::ids`] module so the wasm sync
/// codec and the native branch layer share one identical type.
pub use crate::ids::ShardId;

/// The shard used by the single-shard fork helpers.
pub const DEFAULT_SHARD_ID: ShardId = 0;

/// Errors raised by branch-handle construction and branch-local operations.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BranchError {
    /// A multi-shard branch was requested without any shard roots.
    NoShards,
    /// A shard id appeared more than once while constructing a branch.
    DuplicateShard { shard_id: ShardId },
    /// A branch operation targeted a shard not present in the handle.
    UnknownShard { shard_id: ShardId },
    /// A branch shard's shared state mutex was poisoned by a panicking holder.
    BufferPoisoned { shard_id: ShardId },
    /// `fork_from` was given an unregistered parent handle: the parent's
    /// heads carry no prune pin, so anchoring a child to them is the same
    /// unprotected-anchor hazard `fork_at` hard-refuses (§16.4 Q3). Fork the
    /// parent through a registered constructor, or use `fork_at` against a
    /// protected root.
    UnregisteredParent,
    /// Reading the shared content-addressed tree failed.
    Tree(TreeError),
}

impl fmt::Display for BranchError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NoShards => write!(formatter, "branch requires at least one shard"),
            Self::DuplicateShard { shard_id } => write!(
                formatter,
                "branch shard {shard_id} was provided more than once"
            ),
            Self::UnknownShard { shard_id } => {
                write!(formatter, "branch shard {shard_id} is not present")
            }
            Self::BufferPoisoned { shard_id } => {
                write!(
                    formatter,
                    "branch shard {shard_id} shared state is poisoned"
                )
            }
            Self::UnregisteredParent => write!(
                formatter,
                "fork_from requires a registered parent handle (its heads carry no prune pin \
                 otherwise); fork the parent via a registered constructor or use fork_at"
            ),
            Self::Tree(error) => write!(formatter, "branch tree read error: {error}"),
        }
    }
}

impl std::error::Error for BranchError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Tree(error) => Some(error),
            Self::NoShards
            | Self::DuplicateShard { .. }
            | Self::UnknownShard { .. }
            | Self::BufferPoisoned { .. }
            | Self::UnregisteredParent => None,
        }
    }
}

impl From<TreeError> for BranchError {
    fn from(error: TreeError) -> Self {
        Self::Tree(error)
    }
}

/// Durable-name binding carried by handles opened on a named branch
/// (BRANCH-COMMIT-PATH.md §3, as amended by §16.2).
///
/// Set only by `create_branch`/`open_branch` (A1 stage 3); anonymous `fork*`
/// handles carry no binding. `created` is the branch's creation identity: the
/// create-time timestamp, immutable for the record's life, presented alongside
/// `seq` on every durable advance so a stale handle from a removed-and-
/// recreated branch of the same name can never hijack the new generation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RefBinding {
    /// Full branch name, matching the durable HBR1 record.
    pub name: String,
    /// Creation identity (§16.2): the timestamp captured at branch creation.
    pub created: Timestamp,
    /// The durable commit sequence this handle last observed.
    pub seq: u64,
}

/// Handle to a forked branch.
///
/// A handle stores only root hashes and branch-local WAL buffers. It never owns
/// or copies tree nodes; reads fall through to the caller-provided shared
/// content-addressed [`NodeStore`] at the branch's recorded root hash.
///
/// Clones share per-shard state: each shard's `{current_root, buffer}` pair
/// lives behind one mutex shared by every clone, so a commit's root advance and
/// buffer reset are observed together or not at all — no clone can ever see a
/// "new root + undrained buffer" torn read (BRANCH-COMMIT-PATH.md §3).
#[derive(Debug, Clone)]
pub struct BranchHandle {
    primary_shard: ShardId,
    primary: BranchShard,
    additional_shards: BTreeMap<ShardId, BranchShard>,
    registry_guard: Option<Arc<BranchRegistryGuard>>,
    /// Durable-name binding, shared by all clones. `None` for anonymous
    /// handles; set only by `create_branch`/`open_branch` (A1 stage 3).
    pub(crate) binding: Option<Arc<Mutex<RefBinding>>>,
}

impl BranchHandle {
    /// Create a single-shard branch at `root`.
    ///
    /// The branch's current root is initially equal to its fork-point root and
    /// its WAL buffer is empty.
    #[must_use]
    pub fn new(root: Hash) -> Self {
        Self {
            primary_shard: DEFAULT_SHARD_ID,
            primary: BranchShard::new(root),
            additional_shards: BTreeMap::new(),
            registry_guard: None,
            binding: None,
        }
    }

    /// Create a per-shard branch from committed shard roots.
    ///
    /// This records each shard root and allocates one distinct empty WAL buffer
    /// per shard. No tree store is accepted here, so construction cannot read or
    /// copy tree nodes.
    pub fn from_shard_roots<I>(roots: I) -> Result<Self, BranchError>
    where
        I: IntoIterator<Item = (ShardId, Hash)>,
    {
        Self::from_anchor_head_pairs(
            roots
                .into_iter()
                .map(|(shard_id, root)| (shard_id, root, root)),
        )
    }

    /// Construct a handle whose fork anchors and current heads differ — the
    /// `open_branch` shape, where a branch reopens at its recorded committed
    /// heads while `fork_point` stays the recorded divergence anchor.
    ///
    /// Entries are `(shard_id, fork_anchor, head)`.
    pub(crate) fn from_anchor_head_pairs<I>(entries: I) -> Result<Self, BranchError>
    where
        I: IntoIterator<Item = (ShardId, Hash, Hash)>,
    {
        let mut iter = entries.into_iter();
        let Some((primary_shard, primary_anchor, primary_head)) = iter.next() else {
            return Err(BranchError::NoShards);
        };

        let mut additional_shards = BTreeMap::new();
        for (shard_id, anchor, head) in iter {
            if shard_id == primary_shard
                || additional_shards
                    .insert(shard_id, BranchShard::with_head(anchor, head))
                    .is_some()
            {
                return Err(BranchError::DuplicateShard { shard_id });
            }
        }

        Ok(Self {
            primary_shard,
            primary: BranchShard::with_head(primary_anchor, primary_head),
            additional_shards,
            registry_guard: None,
            binding: None,
        })
    }

    /// The fork-point root hash for the primary shard.
    #[must_use]
    pub const fn fork_point(&self) -> Hash {
        self.primary.fork_point
    }

    /// The current root hash for the primary shard.
    ///
    /// Read under the shard's state mutex, so the value is always one a commit
    /// installed together with its buffer reset — never half a generation.
    #[must_use]
    pub fn current_root(&self) -> Hash {
        self.primary.current_root()
    }

    /// The shard id used by [`fork_point`](Self::fork_point) and
    /// [`current_root`](Self::current_root).
    #[must_use]
    pub const fn primary_shard(&self) -> ShardId {
        self.primary_shard
    }

    /// True when this handle is bound to a durable named branch.
    ///
    /// Named handles are produced by `create_branch`/`open_branch` (A1 stage 3)
    /// and carry the [`RefBinding`] a durable commit presents to the ref
    /// store's CAS; anonymous `fork*` handles return false.
    #[must_use]
    pub const fn is_named(&self) -> bool {
        self.binding.is_some()
    }

    /// Number of shards coordinated by this branch handle.
    #[must_use]
    pub fn shard_count(&self) -> usize {
        self.additional_shards.len().saturating_add(1)
    }

    /// Return true when this handle contains `shard_id`.
    #[must_use]
    pub fn contains_shard(&self, shard_id: ShardId) -> bool {
        self.shard(shard_id).is_some()
    }

    /// Return the fork-point root hash for a specific shard.
    #[must_use]
    pub fn shard_fork_point(&self, shard_id: ShardId) -> Option<Hash> {
        self.shard(shard_id).map(|shard| shard.fork_point)
    }

    /// Return the current root hash for a specific shard.
    ///
    /// Read under the shard's state mutex; see [`current_root`](Self::current_root).
    #[must_use]
    pub fn shard_current_root(&self, shard_id: ShardId) -> Option<Hash> {
        self.shard(shard_id).map(BranchShard::current_root)
    }

    /// Return true when `shard_id` has buffered, uncommitted mutations.
    pub fn shard_is_dirty(&self, shard_id: ShardId) -> Result<bool, BranchError> {
        let state = self.lock_state(shard_id)?;
        Ok(!state.buffer.is_empty())
    }

    /// Iterate the shard ids present in this handle.
    pub fn shard_ids(&self) -> impl Iterator<Item = ShardId> + '_ {
        std::iter::once(self.primary_shard).chain(self.additional_shards.keys().copied())
    }

    /// Buffer a branch-local put on `shard_id`.
    ///
    /// The mutation is appended only to the branch's shard WAL buffer. It is not
    /// propagated to any parent WAL buffer or tree.
    pub fn put<K, V>(&self, shard_id: ShardId, key: K, value: V) -> Result<(), BranchError>
    where
        K: AsRef<[u8]>,
        V: AsRef<[u8]>,
    {
        let mut state = self.lock_state(shard_id)?;
        state.buffer.put(key, value);
        drop(state);
        Ok(())
    }

    /// Buffer a branch-local delete on `shard_id`.
    ///
    /// A buffered delete shadows any value stored in the shared tree at the
    /// branch root.
    pub fn delete<K>(&self, shard_id: ShardId, key: K) -> Result<(), BranchError>
    where
        K: AsRef<[u8]>,
    {
        let mut state = self.lock_state(shard_id)?;
        state.buffer.delete(key);
        drop(state);
        Ok(())
    }

    /// Read a key through the branch view for `shard_id`.
    ///
    /// The branch WAL buffer is checked first. Only keys that are not buffered
    /// fall through to the shared content-addressed tree at the branch's current
    /// shard root, so parent writes after the fork are not visible here.
    pub fn get<S>(
        &self,
        shard_id: ShardId,
        store: &S,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, BranchError>
    where
        S: NodeStore + ?Sized,
    {
        self.get_with_source(shard_id, store, key)
            .map_err(|error| match error {
                BranchGetError::Branch(error) => error,
                BranchGetError::Tree(error) => BranchError::Tree(error.into_tree_error()),
            })
    }

    pub(crate) fn get_with_source<S>(
        &self,
        shard_id: ShardId,
        store: &S,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, BranchGetError<S::Error>>
    where
        S: NodeStore + ?Sized,
    {
        // The buffer lookup and the fallthrough root are taken under ONE state
        // guard: a commit swaps {current_root, buffer} atomically, so the root
        // this read falls through to is always the generation whose buffer
        // produced the miss — never a stale root paired with a newer buffer or
        // vice versa (BRANCH-COMMIT-PATH.md §3 clone coherence).
        let (lookup, root) = {
            let state = self.lock_state(shard_id).map_err(BranchGetError::Branch)?;
            (state.buffer.get(key), state.current_root)
        };

        match lookup {
            LookupResult::BufferedValue(value) => Ok(Some(value)),
            LookupResult::BufferedDelete => Ok(None),
            LookupResult::NotBuffered => Cursor::new(store, root)
                .get_with_source(key)
                .map_err(BranchGetError::Tree),
        }
    }

    #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
    pub(crate) fn fork_points(&self) -> impl Iterator<Item = Hash> + '_ {
        std::iter::once(self.primary.fork_point).chain(
            self.additional_shards
                .values()
                .map(|shard| shard.fork_point),
        )
    }

    pub(crate) fn with_registry_guard(mut self, guard: BranchRegistryGuard) -> Self {
        self.registry_guard = Some(Arc::new(guard));
        self
    }

    /// The registry guard pinning this handle's roots, if any.
    ///
    /// `commit_branch` refuses handles without one (`Unregistered`): advancing
    /// a root that no guard pins is precisely the prune registration hazard.
    pub(crate) fn registry_guard(&self) -> Option<&BranchRegistryGuard> {
        self.registry_guard.as_deref()
    }

    /// Bind this handle to a durable named branch (`create_branch`/`open_branch`).
    pub(crate) fn with_binding(mut self, binding: RefBinding) -> Self {
        self.binding = Some(Arc::new(Mutex::new(binding)));
        self
    }

    /// Snapshots the durable identity carried by this handle.
    ///
    /// Kind and lineage are intentionally absent: policy entry points use the
    /// name and generation only to resolve those fields from the ref store,
    /// never from caller-held handle state.
    #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
    pub(crate) fn binding_snapshot(&self) -> Option<RefBinding> {
        self.binding.as_deref().map(|binding| match binding.lock() {
            Ok(bound) => bound.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        })
    }

    /// Lock every shard's `ShardState`, strictly ascending by shard id.
    ///
    /// Ascending shard id is tier 1 of the crate-wide commit lock order
    /// (BRANCH-COMMIT-PATH.md §16.1: shard states ascending → binding →
    /// registry counts). Every multi-shard acquirer — `commit_branch` step 1
    /// and `fork_from` — MUST take shard states through this method so two
    /// concurrent acquirers can never hold overlapping subsets in opposite
    /// orders and deadlock. The primary shard is not special: ordering is by
    /// id alone.
    pub(crate) fn lock_states_ascending(
        &self,
    ) -> Result<Vec<(ShardId, MutexGuard<'_, ShardState>)>, BranchError> {
        let mut shards: Vec<(ShardId, &BranchShard)> =
            std::iter::once((self.primary_shard, &self.primary))
                .chain(
                    self.additional_shards
                        .iter()
                        .map(|(shard_id, shard)| (*shard_id, shard)),
                )
                .collect();
        shards.sort_unstable_by_key(|(shard_id, _shard)| *shard_id);

        shards
            .into_iter()
            .map(|(shard_id, shard)| match shard.state.lock() {
                Ok(guard) => Ok((shard_id, guard)),
                Err(poisoned) => {
                    drop(poisoned.into_inner());
                    Err(BranchError::BufferPoisoned { shard_id })
                }
            })
            .collect()
    }

    fn shard(&self, shard_id: ShardId) -> Option<&BranchShard> {
        if shard_id == self.primary_shard {
            Some(&self.primary)
        } else {
            self.additional_shards.get(&shard_id)
        }
    }

    fn lock_state(&self, shard_id: ShardId) -> Result<MutexGuard<'_, ShardState>, BranchError> {
        let shard = self
            .shard(shard_id)
            .ok_or(BranchError::UnknownShard { shard_id })?;
        match shard.state.lock() {
            Ok(guard) => Ok(guard),
            Err(poisoned) => {
                drop(poisoned.into_inner());
                Err(BranchError::BufferPoisoned { shard_id })
            }
        }
    }
}

/// One shard's slice of a branch: the immutable divergence point plus the
/// advanceable `{current_root, buffer}` pair shared by every handle clone.
#[derive(Debug, Clone)]
struct BranchShard {
    /// Immutable for the handle's life — the fork anchor never advances.
    fork_point: Hash,
    /// Shared by all clones; one mutex covers the root AND the buffer so no
    /// clone can observe "new root + undrained buffer" or vice versa (§3).
    state: Arc<Mutex<ShardState>>,
}

/// The advanceable half of a branch shard, always swapped as one unit.
///
/// Fields are `pub(crate)` for exactly one out-of-module writer:
/// `commit_branch`'s step-6 swap, which runs under the guards handed out by
/// [`BranchHandle::lock_states_ascending`] and writes both fields together —
/// preserving the one-mutex-one-generation invariant of §3.
#[derive(Debug)]
pub(crate) struct ShardState {
    pub(crate) current_root: Hash,
    pub(crate) buffer: WalBuffer,
}

impl BranchShard {
    fn new(root: Hash) -> Self {
        Self::with_head(root, root)
    }

    fn with_head(fork_point: Hash, head: Hash) -> Self {
        Self {
            fork_point,
            state: Arc::new(Mutex::new(ShardState {
                current_root: head,
                buffer: WalBuffer::new(),
            })),
        }
    }

    /// Poison-tolerant root read: `current_root` is `Copy` and only ever
    /// written as a single field store under the guard, so a poisoned mutex
    /// still holds a valid, whole value. Recovering here keeps the root
    /// accessors' infallible signatures; mutating paths surface
    /// [`BranchError::BufferPoisoned`] instead.
    fn current_root(&self) -> Hash {
        match self.state.lock() {
            Ok(guard) => guard.current_root,
            Err(poisoned) => poisoned.into_inner().current_root,
        }
    }
}

#[cfg(test)]
#[path = "handle_tests.rs"]
mod tests;