//! Immutable Merkle Search Tree operations

use super::node::NodeEntry;
use super::util;
use crate::error::{RepoError, Result};
use crate::mst::util::validate_key;
use crate::storage::BlockStore;
use bytes::Bytes;
use cid::Cid as IpldCid;
use core::fmt;
use jacquard_common::BosStr;
use jacquard_common::types::recordkey::Rkey;
use jacquard_common::types::string::{Nsid, RecordKey};
use jacquard_common::types::value::RawData;
use n0_future::try_join_all;
use smol_str::SmolStr;
use std::collections::BTreeMap;
use std::fmt::{Display, Formatter};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Write operation for batch application
///
/// Represents a single operation to apply to an MST.
/// For firehose operations where `prev` may be optional (v3).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WriteOp {
    /// Create new record (error if exists)
    Create {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// Record CID
        cid: IpldCid,
    },

    /// Update existing record (error if not exists)
    ///
    /// `prev` is optional in v3 (required in v2)
    Update {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// New record CID
        cid: IpldCid,
        /// Previous CID (optional for validation)
        prev: Option<IpldCid>,
    },

    /// Delete record
    ///
    /// `prev` is optional in v3 (required in v2)
    Delete {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// Previous CID (optional for validation)
        prev: Option<IpldCid>,
    },
}

/// Record write operation with inline data
///
/// Used for high-level record operations where the actual record data
/// needs to be serialized and stored. The data is a generic IPLD map
/// (similar to rsky's `RepoRecord = BTreeMap<String, Lex>`).
#[derive(Debug, Clone, PartialEq)]
pub enum RecordWriteOp<'a, S: BosStr> {
    /// Create new record with data
    Create {
        /// Collection NSID
        collection: Nsid<S>,
        /// Record key
        rkey: RecordKey<Rkey<S>>,
        /// Record data (will be serialized to DAG-CBOR and CID computed)
        record: std::collections::BTreeMap<SmolStr, RawData<'a>>,
    },

    /// Update existing record with new data
    Update {
        /// Collection NSID
        collection: Nsid<S>,
        /// Record key
        rkey: RecordKey<Rkey<S>>,
        /// New record data
        record: std::collections::BTreeMap<SmolStr, RawData<'a>>,
        /// Previous CID (optional for validation)
        prev: Option<IpldCid>,
    },

    /// Delete record
    Delete {
        /// Collection NSID
        collection: Nsid<S>,
        /// Record key
        rkey: RecordKey<Rkey<S>>,
        /// Previous CID (optional for validation)
        prev: Option<IpldCid>,
    },
}

impl<'a, S: BosStr> RecordWriteOp<'a, S> {
    /// Get the collection NSID for this operation
    pub fn collection(&self) -> &Nsid<S> {
        match self {
            RecordWriteOp::Create { collection, .. } => collection,
            RecordWriteOp::Update { collection, .. } => collection,
            RecordWriteOp::Delete { collection, .. } => collection,
        }
    }

    /// Get the record key for this operation
    pub fn rkey(&self) -> &RecordKey<Rkey<S>> {
        match self {
            RecordWriteOp::Create { rkey, .. } => rkey,
            RecordWriteOp::Update { rkey, .. } => rkey,
            RecordWriteOp::Delete { rkey, .. } => rkey,
        }
    }
}

/// Verified write operation with required prev fields
///
/// Used for operations where prev CID has been verified against tree state.
/// Safer than `WriteOp` because it always validates prev values.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VerifiedWriteOp {
    /// Create new record (verified not to exist)
    Create {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// Record CID
        cid: IpldCid,
    },

    /// Update existing record (with verified prev CID)
    Update {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// New record CID
        cid: IpldCid,
        /// Previous CID (required, validated)
        prev: IpldCid,
    },

    /// Delete record (with verified current CID)
    Delete {
        /// Record key (collection/rkey)
        key: SmolStr,
        /// Previous CID (required, validated)
        prev: IpldCid,
    },
}

/// Immutable Merkle Search Tree
///
/// MST operations return new tree instances, leaving the original unchanged.
/// This enables versioning and safe concurrent access to different tree versions.
///
/// # Architecture
///
/// The tree uses a flat `Vec<NodeEntry>` representation in memory, where
/// `NodeEntry` is an enum of `Tree` (subtree) and `Leaf` (key-value pair).
///
/// Entries are interleaved: `[Tree, Leaf, Tree, Leaf, Leaf, Tree]` etc.
/// This representation makes operations simple (Vec slicing, splicing).
///
/// The wire format (CBOR) uses prefix compression and pointers (left/tree).
/// See `NodeData` and `TreeEntry` in node.rs for serialization format.
///
/// # Layer-based structure
///
/// Keys are hashed (SHA-256) and leading zero bits determine layer:
/// - More leading zeros = higher layer (deeper in tree)
/// - Layer = floor(leading_zeros / 2) for ~4 fanout
/// - Deterministic and insertion-order independent
#[derive(Clone)]
pub struct Mst<S> {
    /// Block storage for loading/saving nodes (shared via Arc)
    storage: Arc<S>,

    /// Flat list of entries (lazy-loaded, interior mutable)
    ///
    /// `None` means not yet loaded from storage.
    /// Empty `Vec` means tree has been loaded and has no entries.
    entries: Arc<RwLock<Option<Vec<NodeEntry<S>>>>>,

    /// CID pointer to this node in storage (interior mutable)
    pointer: Arc<RwLock<IpldCid>>,

    /// Whether pointer is stale (entries modified, interior mutable)
    ///
    /// When `true`, `pointer` doesn't match current `entries`.
    /// Call `get_pointer()` to recompute and update.
    outdated_pointer: Arc<RwLock<bool>>,

    /// Layer hint for this node
    ///
    /// `None` means layer unknown (will be computed from entries).
    /// Layer is the maximum layer of any key in this node.
    layer: Option<usize>,
}

impl<S: BlockStore + Sync + 'static> Mst<S> {
    /// Create empty MST
    pub fn new(storage: Arc<S>) -> Self {
        Self {
            storage,
            entries: Arc::new(RwLock::new(Some(Vec::new()))),
            pointer: Arc::new(RwLock::new(IpldCid::default())),
            outdated_pointer: Arc::new(RwLock::new(true)),
            layer: Some(0),
        }
    }

    /// Create MST with existing entries
    ///
    /// Used internally for tree operations.
    /// Computes CID from entries (doesn't persist to storage).
    pub(crate) async fn create(
        storage: Arc<S>,
        entries: Vec<NodeEntry<S>>,
        layer: Option<usize>,
    ) -> Result<Self> {
        // Serialize and compute CID (don't persist yet)
        let node_data = util::serialize_node_data(&entries).await?;
        let cbor = serde_ipld_dagcbor::to_vec(&node_data).map_err(|e| {
            RepoError::serialization(e).with_context("serializing MST node during creation")
        })?;
        let cid = util::compute_cid(&cbor)?;

        let mst = Self {
            storage,
            entries: Arc::new(RwLock::new(Some(entries))),
            pointer: Arc::new(RwLock::new(cid)),
            outdated_pointer: Arc::new(RwLock::new(false)),
            layer,
        };

        Ok(mst)
    }

    /// Load MST from CID (lazy)
    ///
    /// Doesn't actually load from storage until entries are accessed.
    pub fn load(storage: Arc<S>, cid: IpldCid, layer: Option<usize>) -> Self {
        Self {
            storage,
            entries: Arc::new(RwLock::new(None)), // Not loaded yet
            pointer: Arc::new(RwLock::new(cid)),
            outdated_pointer: Arc::new(RwLock::new(false)),
            layer,
        }
    }

    /// Create new tree with modified entries
    ///
    /// Returns a new Mst with updated entries. Marks pointer as outdated.
    async fn new_tree(&self, entries: Vec<NodeEntry<S>>) -> Result<Self> {
        Ok(Self {
            storage: self.storage.clone(),
            entries: Arc::new(RwLock::new(Some(entries))),
            pointer: Arc::new(RwLock::new(self.pointer.read().await.clone())),
            outdated_pointer: Arc::new(RwLock::new(true)),
            layer: self.layer,
        })
    }

    /// Get entries (lazy load if needed)
    pub(crate) async fn get_entries(&self) -> Result<Vec<NodeEntry<S>>> {
        {
            let entries_guard = self.entries.read().await;
            if let Some(ref entries) = *entries_guard {
                return Ok(entries.clone());
            }
        }

        // Load from storage
        let pointer = *self.pointer.read().await;
        let node_bytes = self
            .storage
            .get(&pointer)
            .await?
            .ok_or_else(|| {
                RepoError::not_found("MST node", &pointer)
                    .with_help("MST node missing from storage - ensure all blocks were properly persisted or that the tree CID is correct")
            })?;

        let node_data: super::node::NodeData = serde_ipld_dagcbor::from_slice(&node_bytes)
            .map_err(|e| {
                RepoError::serialization(e)
                    .with_context(format!("deserializing MST node from storage: {}", pointer))
            })?;

        let entries = util::deserialize_node_data(self.storage.clone(), &node_data, self.layer)?;

        // Cache the loaded entries
        {
            let mut entries_guard = self.entries.write().await;
            *entries_guard = Some(entries.clone());
        }

        Ok(entries)
    }

    /// Get CID pointer (recompute if outdated)
    ///
    /// Computes CID from current entries but doesn't persist to storage.
    /// Use `collect_blocks()` to gather blocks for persistence.
    pub fn get_pointer<'a>(&'a self) -> Pin<Box<dyn Future<Output = Result<IpldCid>> + Send + 'a>> {
        Box::pin(async move {
            let outdated = *self.outdated_pointer.read().await;
            if !outdated {
                return Ok(*self.pointer.read().await);
            }

            // Check for outdated children and recursively update them first
            let mut entries = self.get_entries().await?;
            let mut outdated_children = Vec::new();

            for entry in &entries {
                if let NodeEntry::Tree(mst) = entry {
                    if *mst.outdated_pointer.read().await {
                        let child = mst.clone();
                        outdated_children.push(n0_future::task::spawn(async move {
                            child.get_pointer().await
                        }));
                    }
                }
            }

            // Recursively update outdated children concurrently
            if !outdated_children.is_empty() {
                try_join_all(outdated_children)
                    .await
                    .map_err(|e| RepoError::task_failed(e))?;

                // Re-fetch entries with updated child CIDs
                entries = self.get_entries().await?;
            }

            // Now serialize and compute CID with fresh child CIDs
            let node_data = util::serialize_node_data(&entries).await?;
            let cbor = serde_ipld_dagcbor::to_vec(&node_data).map_err(|e| {
                RepoError::serialization(e).with_context("serializing MST node for CID computation")
            })?;
            let cid = util::compute_cid(&cbor)?;

            // Update pointer and mark as fresh
            {
                let mut pointer_guard = self.pointer.write().await;
                *pointer_guard = cid;
            }
            {
                let mut outdated_guard = self.outdated_pointer.write().await;
                *outdated_guard = false;
            }

            Ok(cid)
        })
    }

    /// Get root CID (alias for get_pointer)
    pub async fn root(&self) -> Result<IpldCid> {
        self.get_pointer().await
    }

    /// Get shared reference to the block storage
    ///
    /// Useful for CAR export and other operations that need direct storage access.
    pub fn storage(&self) -> &Arc<S> {
        &self.storage
    }

    /// Get the layer of this node
    ///
    /// Layer is the maximum layer of any leaf key in this node.
    /// For nodes with no leaves, recursively checks subtrees.
    pub(crate) fn get_layer<'a>(
        &'a self,
    ) -> Pin<Box<dyn Future<Output = Result<usize>> + Send + 'a>> {
        Box::pin(async move {
            if let Some(layer) = self.layer {
                return Ok(layer);
            }

            // Compute layer from entries
            let entries = self.get_entries().await?;

            // Find first leaf and get its layer
            for entry in &entries {
                if let NodeEntry::Leaf { key, .. } = entry {
                    let layer = util::layer_for_key(key.as_str());
                    return Ok(layer);
                }
            }

            // No leaves found - check first subtree
            for entry in &entries {
                if let NodeEntry::Tree(subtree) = entry {
                    let child_layer = subtree.get_layer().await?;
                    return Ok(child_layer + 1);
                }
            }

            // Empty tree
            Ok(0)
        })
    }

    /// Find index of first leaf >= key
    ///
    /// Returns `entries.len()` if all leaves are < key.
    fn find_gt_or_equal_leaf_index_in(entries: &[NodeEntry<S>], key: &str) -> usize {
        for (i, entry) in entries.iter().enumerate() {
            if let NodeEntry::Leaf { key: leaf_key, .. } = entry {
                if leaf_key.as_str() >= key {
                    return i;
                }
            }
        }

        entries.len()
    }

    /// Get a value by key
    pub fn get<'a>(
        &'a self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<Option<IpldCid>>> + Send + 'a>> {
        Box::pin(async move {
            validate_key(key)?;

            let entries = self.get_entries().await?;
            let index = Self::find_gt_or_equal_leaf_index_in(&entries, key);

            // Check if we found exact match
            if index < entries.len() {
                if let NodeEntry::Leaf {
                    key: leaf_key,
                    value,
                } = &entries[index]
                {
                    if leaf_key.as_str() == key {
                        return Ok(Some(*value));
                    }
                }
            }

            // Not found at this level - check subtree before this index
            if index > 0 {
                if let NodeEntry::Tree(subtree) = &entries[index - 1] {
                    return subtree.get(key).await;
                }
            }

            Ok(None)
        })
    }

    /// Add a key-value pair, mutating the current tree
    pub async fn add_mut<'a>(&'a mut self, key: &'a str, cid: IpldCid) -> Result<()> {
        *self = self.add(key, cid).await?;
        Ok(())
    }

    /// Add a key-value pair (returns new tree)
    pub fn add<'a>(
        &'a self,
        key: &'a str,
        cid: IpldCid,
    ) -> Pin<Box<dyn Future<Output = Result<Mst<S>>> + Send + 'a>> {
        Box::pin(async move {
            validate_key(key)?;

            let key_layer = util::layer_for_key(key);
            let node_layer = self.get_layer().await?;
            let entries = self.get_entries().await?;

            if key_layer == node_layer {
                // Key belongs at this layer - insert here
                let index = Self::find_gt_or_equal_leaf_index_in(&entries, key);

                // Check if key already exists
                if index < entries.len() {
                    if let NodeEntry::Leaf { key: leaf_key, .. } = &entries[index] {
                        if leaf_key.as_str() == key {
                            // Key exists - replace by just inserting at same position
                            let mut new_entries = entries.clone();
                            new_entries[index] = NodeEntry::Leaf {
                                key: smol_str::SmolStr::new(key),
                                value: cid,
                            };
                            return self.new_tree(new_entries).await;
                        }
                    }
                }

                // Check entry before insertion point
                if index > 0 {
                    match &entries[index - 1] {
                        NodeEntry::Leaf { .. } => {
                            // Prev is Leaf - just splice in
                            self.splice_in(
                                NodeEntry::Leaf {
                                    key: smol_str::SmolStr::new(key),
                                    value: cid,
                                },
                                index,
                            )
                            .await
                        }
                        NodeEntry::Tree(subtree) => {
                            // Prev is Tree - split it around key
                            let (left, right) = subtree.split_around(key).await?;
                            self.replace_with_split(
                                index - 1,
                                left,
                                NodeEntry::Leaf {
                                    key: smol_str::SmolStr::new(key),
                                    value: cid,
                                },
                                right,
                            )
                            .await
                        }
                    }
                } else {
                    // At far left - splice in
                    self.splice_in(
                        NodeEntry::Leaf {
                            key: smol_str::SmolStr::new(key),
                            value: cid,
                        },
                        index,
                    )
                    .await
                }
            } else if key_layer < node_layer {
                // Key belongs on lower layer - recurse into subtree
                let index = Self::find_gt_or_equal_leaf_index_in(&entries, key);

                if index > 0 {
                    if let NodeEntry::Tree(prev_tree) = &entries[index - 1] {
                        // Prev is Tree - add to it
                        let new_subtree = prev_tree.add(key, cid).await?;
                        return self
                            .update_entry(index - 1, NodeEntry::Tree(new_subtree))
                            .await;
                    }
                }

                // No prev tree - create child and add to it
                let child = self.create_child().await?;
                let new_subtree = child.add(key, cid).await?;
                self.splice_in(NodeEntry::Tree(new_subtree), index).await
            } else {
                // Key belongs on higher layer - create parent layers
                let extra_layers = key_layer - node_layer;

                let (mut left, mut right) = self.split_around(key).await?;

                // Create intermediate layers if gap > 1
                for _ in 1..extra_layers {
                    if let Some(l) = left {
                        left = Some(l.create_parent().await?);
                    }
                    if let Some(r) = right {
                        right = Some(r.create_parent().await?);
                    }
                }

                // Build new root
                let mut new_entries = Vec::new();
                if let Some(l) = left {
                    new_entries.push(NodeEntry::Tree(l));
                }
                new_entries.push(NodeEntry::Leaf {
                    key: smol_str::SmolStr::new(key),
                    value: cid,
                });
                if let Some(r) = right {
                    new_entries.push(NodeEntry::Tree(r));
                }

                Mst::create(self.storage.clone(), new_entries, Some(key_layer)).await
            }
        })
    }

    /// invert an update function, returning the previous cid
    pub async fn check_update(&mut self, key: &str, cid: IpldCid) -> Result<IpldCid> {
        validate_key(key)?;

        // Check key exists
        let Ok(Some(prev)) = self.get(key).await else {
            return Err(RepoError::not_found("key", key));
        };

        if prev == cid {
            return Ok(prev);
        }

        // Update is just add (which replaces)
        *self = self.add(key, cid).await?;
        Ok(prev)
    }

    /// Delete a key (returns new tree)
    pub fn delete<'a>(
        &'a self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<Mst<S>>> + Send + 'a>> {
        Box::pin(async move {
            validate_key(key)?;

            let altered = self.delete_recurse(key).await?;
            altered.trim_top().await
        })
    }

    /// mutates a tree in place to delete, returns the CID of what was deleted
    ///
    /// Used to invert tree operations for verification
    pub fn delete_cid<'a>(
        &'a mut self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<IpldCid>> + Send + 'a>> {
        Box::pin(async move {
            let cid = self
                .get(key)
                .await?
                .ok_or(RepoError::not_found("cid for key", key))?;
            *self = self.delete(key).await?;
            Ok(cid)
        })
    }

    /// Recursively delete a key
    fn delete_recurse<'a>(
        &'a self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<Mst<S>>> + Send + 'a>> {
        Box::pin(async move {
            let entries = self.get_entries().await?;
            let index = Self::find_gt_or_equal_leaf_index_in(&entries, key);

            // Check if found at this level
            if index < entries.len() {
                if let NodeEntry::Leaf { key: leaf_key, .. } = &entries[index] {
                    if leaf_key.as_str() == key {
                        // Found it - delete this entry
                        let prev = if index > 0 {
                            Some(&entries[index - 1])
                        } else {
                            None
                        };
                        let next = entries.get(index + 1);

                        // Check if we need to merge Trees
                        if let (
                            Some(NodeEntry::Tree(prev_tree)),
                            Some(NodeEntry::Tree(next_tree)),
                        ) = (prev, next)
                        {
                            // Merge the two Trees
                            let merged = prev_tree.append_merge(next_tree).await?;

                            // Build: [0..index-1] + [merged] + [index+2..]
                            let mut new_entries = entries[..index - 1].to_vec();
                            new_entries.push(NodeEntry::Tree(merged));
                            new_entries.extend_from_slice(&entries[index + 2..]);

                            return self.new_tree(new_entries).await;
                        }

                        // Simple case: just remove the entry
                        return self.remove_entry(index).await;
                    }
                }
            }

            // Not found at this level - recurse into prev Tree
            if index > 0 {
                if let NodeEntry::Tree(prev_tree) = &entries[index - 1] {
                    let subtree = prev_tree.delete_recurse(key).await?;
                    let subtree_entries = subtree.get_entries().await?;

                    if subtree_entries.is_empty() {
                        // Subtree is now empty - remove it
                        return self.remove_entry(index - 1).await;
                    } else {
                        // Update with new subtree
                        return self.update_entry(index - 1, NodeEntry::Tree(subtree)).await;
                    }
                }
            }

            // Key not found
            Err(RepoError::not_found("key", key))
        })
    }

    /// Update an existing key (returns new tree)
    pub async fn update(&self, key: &str, cid: IpldCid) -> Result<Mst<S>> {
        validate_key(key)?;

        // Check key exists
        if self.get(key).await?.is_none() {
            return Err(RepoError::not_found("key", key));
        }

        // Update is just add (which replaces)
        self.add(key, cid).await
    }

    /// Update entry at index
    async fn update_entry(&self, index: usize, entry: NodeEntry<S>) -> Result<Mst<S>> {
        let mut entries = self.get_entries().await?;
        entries[index] = entry;
        self.new_tree(entries).await
    }

    /// Remove entry at index
    async fn remove_entry(&self, index: usize) -> Result<Mst<S>> {
        let mut entries = self.get_entries().await?;
        entries.remove(index);
        self.new_tree(entries).await
    }

    /// Append entry to end
    async fn append(&self, entry: NodeEntry<S>) -> Result<Mst<S>> {
        let mut entries = self.get_entries().await?;
        entries.push(entry);
        self.new_tree(entries).await
    }

    /// Prepend entry to start
    async fn prepend(&self, entry: NodeEntry<S>) -> Result<Mst<S>> {
        let mut entries = self.get_entries().await?;
        entries.insert(0, entry);
        self.new_tree(entries).await
    }

    /// Splice in entry at index
    async fn splice_in(&self, entry: NodeEntry<S>, index: usize) -> Result<Mst<S>> {
        let mut entries = self.get_entries().await?;
        entries.insert(index, entry);
        self.new_tree(entries).await
    }

    /// Get slice of entries
    pub async fn slice(&self, start: usize, end: usize) -> Result<Vec<NodeEntry<S>>> {
        let entries = self.get_entries().await?;
        Ok(entries[start..end].to_vec())
    }

    /// Trim top node if it only contains one subtree
    fn trim_top(self) -> Pin<Box<dyn Future<Output = Result<Mst<S>>> + Send>> {
        Box::pin(async move {
            let entries = self.get_entries().await?;

            if entries.len() == 1 {
                if let NodeEntry::Tree(subtree) = &entries[0] {
                    return subtree.clone().trim_top().await;
                }
            }

            Ok(self)
        })
    }

    /// Split tree around a key into left and right subtrees
    ///
    /// Returns (left, right) where:
    /// - left contains all entries < key
    /// - right contains all entries >= key
    ///
    /// Either side can be None if empty.
    pub fn split_around<'a>(
        &'a self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<(Option<Mst<S>>, Option<Mst<S>>)>> + Send + 'a>> {
        Box::pin(async move {
            let entries = self.get_entries().await?;
            let index = Self::find_gt_or_equal_leaf_index_in(&entries, key);

            // Split at index
            let left_data = entries[..index].to_vec();
            let right_data = entries[index..].to_vec();

            let mut left = self.new_tree(left_data.clone()).await?;
            let mut right = self.new_tree(right_data).await?;

            // If last entry in left is a Tree, recursively split it
            if let Some(NodeEntry::Tree(last_tree)) = left_data.last() {
                let left_len = left_data.len();
                left = left.remove_entry(left_len - 1).await?;

                let (split_left, split_right) = last_tree.split_around(key).await?;

                if let Some(sl) = split_left {
                    left = left.append(NodeEntry::Tree(sl)).await?;
                }
                if let Some(sr) = split_right {
                    right = right.prepend(NodeEntry::Tree(sr)).await?;
                }
            }

            // Return None for empty sides
            let left_out = if left.get_entries().await?.is_empty() {
                None
            } else {
                Some(left)
            };

            let right_out = if right.get_entries().await?.is_empty() {
                None
            } else {
                Some(right)
            };

            Ok((left_out, right_out))
        })
    }

    /// Merge two adjacent subtrees
    ///
    /// All keys in `to_merge` must be > all keys in `self`.
    /// Used primarily for delete operations.
    pub fn append_merge<'a>(
        &'a self,
        to_merge: &'a Mst<S>,
    ) -> Pin<Box<dyn Future<Output = Result<Mst<S>>> + Send + 'a>> {
        Box::pin(async move {
            // Check same layer
            let self_layer = self.get_layer().await?;
            let merge_layer = to_merge.get_layer().await?;

            if self_layer != merge_layer {
                return Err(RepoError::invalid_mst(
                    "Cannot merge MST nodes from different layers",
                ));
            }

            let mut self_entries = self.get_entries().await?;
            let merge_entries = to_merge.get_entries().await?;

            // Check if we need to merge adjacent Trees
            let last_is_tree = matches!(self_entries.last(), Some(NodeEntry::Tree(_)));
            let first_is_tree = matches!(merge_entries.first(), Some(NodeEntry::Tree(_)));

            if last_is_tree && first_is_tree {
                // Both are Trees - recursively merge them
                if let (Some(NodeEntry::Tree(left_tree)), Some(NodeEntry::Tree(right_tree))) =
                    (self_entries.last(), merge_entries.first())
                {
                    let merged = left_tree.append_merge(right_tree).await?;

                    // Build new entries: self[..-1] + merged + merge[1..]
                    let mut new_entries = self_entries[..self_entries.len() - 1].to_vec();
                    new_entries.push(NodeEntry::Tree(merged));
                    new_entries.extend_from_slice(&merge_entries[1..]);

                    return self.new_tree(new_entries).await;
                }
            }

            // Simple case: just concatenate
            self_entries.extend(merge_entries);
            self.new_tree(self_entries).await
        })
    }

    /// Create empty child tree at layer-1
    pub async fn create_child(&self) -> Result<Mst<S>> {
        let layer = self.get_layer().await?;
        let child_layer = if layer > 0 { Some(layer - 1) } else { Some(0) };

        Mst::create(self.storage.clone(), Vec::new(), child_layer).await
    }

    /// Create parent tree at layer+1 containing self
    pub async fn create_parent(self) -> Result<Mst<S>> {
        let layer = self.get_layer().await?;

        Mst::create(
            self.storage.clone(),
            vec![NodeEntry::Tree(self)],
            Some(layer + 1),
        )
        .await
    }

    /// Replace entry at index with [left?, leaf, right?]
    async fn replace_with_split(
        &self,
        index: usize,
        left: Option<Mst<S>>,
        leaf: NodeEntry<S>,
        right: Option<Mst<S>>,
    ) -> Result<Mst<S>> {
        let entries = self.get_entries().await?;

        // Build: [0..index] + [left?] + [leaf] + [right?] + [index+1..]
        let mut new_entries = entries[..index].to_vec();

        if let Some(l) = left {
            new_entries.push(NodeEntry::Tree(l));
        }
        new_entries.push(leaf);
        if let Some(r) = right {
            new_entries.push(NodeEntry::Tree(r));
        }

        new_entries.extend_from_slice(&entries[index + 1..]);

        self.new_tree(new_entries).await
    }

    /// Get all leaf entries (key-CID pairs) in lexicographic order
    ///
    /// Recursively traverses the tree to collect all leaves.
    /// Used for diff calculation and tree listing.
    ///
    /// Uses parallel traversal to collect leaves from independent subtrees concurrently.
    pub fn leaves<'a>(
        &'a self,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<(smol_str::SmolStr, IpldCid)>>> + Send + 'a>> {
        Box::pin(async move { collect_leaves_parallel(self.clone()).await })
    }

    /// Get all leaf entries sequentially (for benchmarking)
    pub fn leaves_sequential<'a>(
        &'a self,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<(smol_str::SmolStr, IpldCid)>>> + Send + 'a>> {
        Box::pin(async move {
            let mut result = Vec::new();
            self.collect_leaves_sequential(&mut result).await?;
            Ok(result)
        })
    }

    /// Recursively collect all leaves into the result vector (sequential)
    fn collect_leaves_sequential<'a>(
        &'a self,
        result: &'a mut Vec<(smol_str::SmolStr, IpldCid)>,
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
        Box::pin(async move {
            let entries = self.get_entries().await?;

            for entry in entries {
                match entry {
                    NodeEntry::Tree(subtree) => {
                        // Recurse into subtree
                        subtree.collect_leaves_sequential(result).await?;
                    }
                    NodeEntry::Leaf { key, value } => {
                        // Add leaf to result
                        result.push((key, value));
                    }
                }
            }

            Ok(())
        })
    }

    /// Copy tree with same entries (marking pointer as outdated)
    ///
    /// Internal helper for creating modified tree copies.
    pub async fn copy_tree(&self) -> Result<Self> {
        let entries = self.get_entries().await?;
        self.new_tree(entries).await
    }

    /// invert a tree operation in-place for validation
    pub async fn invert_op(&mut self, op: VerifiedWriteOp) -> Result<bool> {
        //println!("tree before op inversion:\n{}", self);
        match op {
            VerifiedWriteOp::Create { key, cid: expected } => {
                let Ok(found) = self.delete_cid(&key).await else {
                    //println!("tree at failure:\n{}", self);
                    return Ok(false);
                };
                if found == expected {
                    Ok(true)
                } else {
                    //println!("tree at failure:\n{}", self);
                    Ok(false)
                }
            }
            VerifiedWriteOp::Update {
                key,
                cid: expected,
                prev,
            } => {
                let Ok(found) = self.check_update(&key, prev).await else {
                    //println!("tree at failure:\n{}", self);
                    return Ok(false);
                };
                if found == expected {
                    Ok(true)
                } else {
                    //println!("tree at failure:\n{}", self);
                    Ok(false)
                }
            }
            VerifiedWriteOp::Delete { key, prev } => {
                if let Ok(Some(_)) = self.get(&key).await {
                    Ok(false)
                } else {
                    self.add_mut(&key, prev).await?;
                    Ok(true)
                }
            }
        }
    }

    /// Apply batch of verified write operations (returns new tree)
    ///
    /// More efficient than individual operations as it only rebuilds
    /// the tree structure once per operation. Operations are applied in order.
    ///
    /// # Validation
    ///
    /// - Create: errors if key already exists
    /// - Update: errors if key doesn't exist OR prev CID doesn't match
    /// - Delete: errors if key doesn't exist OR prev CID doesn't match
    ///
    /// All operations validate prev CIDs against current tree state.
    pub async fn batch(&self, ops: &[VerifiedWriteOp]) -> Result<Mst<S>> {
        let mut tree = self.clone();

        for op in ops {
            tree = match op {
                VerifiedWriteOp::Create { key, cid } => {
                    // Check doesn't exist
                    if tree.get(key.as_str()).await?.is_some() {
                        return Err(RepoError::invalid_mst(format!(
                            "Cannot create: key already exists: {}",
                            key
                        )));
                    }
                    tree.add(key.as_str(), *cid).await?
                }

                VerifiedWriteOp::Update { key, cid, prev } => {
                    // Check exists and validate prev
                    let current = tree
                        .get(key.as_str())
                        .await?
                        .ok_or_else(|| RepoError::not_found("key", key.as_str()))?;

                    if &current != prev {
                        return Err(RepoError::cid_mismatch(format!(
                            "Update prev CID mismatch for key {}: expected {}, got {}",
                            key, prev, current
                        )));
                    }

                    tree.add(key.as_str(), *cid).await?
                }

                VerifiedWriteOp::Delete { key, prev } => {
                    // Check exists and validate prev
                    let current = tree
                        .get(key.as_str())
                        .await?
                        .ok_or_else(|| RepoError::not_found("key", key.as_str()))?;

                    if &current != prev {
                        return Err(RepoError::cid_mismatch(format!(
                            "Delete prev CID mismatch for key {}: expected {}, got {}",
                            key, prev, current
                        )));
                    }

                    tree.delete(key.as_str()).await?
                }
            };
        }

        Ok(tree)
    }

    /// Collect all blocks that need persisting
    ///
    /// Recursively walks the tree, serializing nodes and collecting blocks
    /// that aren't already in storage. Skips nodes that are already persisted.
    ///
    /// Returns (root_cid, blocks) where blocks is a map of CID → bytes.
    ///
    /// Uses parallel traversal to collect blocks from independent subtrees concurrently.
    pub fn collect_blocks<'a>(
        &'a self,
    ) -> Pin<
        Box<
            dyn Future<
                    Output = Result<(IpldCid, std::collections::BTreeMap<IpldCid, bytes::Bytes>)>,
                > + Send
                + 'a,
        >,
    > {
        Box::pin(async move { collect_blocks_parallel(self.clone()).await })
    }

    /// Collect all blocks sequentially (for benchmarking)
    pub fn collect_blocks_sequential<'a>(
        &'a self,
    ) -> Pin<
        Box<
            dyn Future<
                    Output = Result<(IpldCid, std::collections::BTreeMap<IpldCid, bytes::Bytes>)>,
                > + Send
                + 'a,
        >,
    > {
        Box::pin(async move {
            use bytes::Bytes;
            use std::collections::BTreeMap;

            let mut blocks = BTreeMap::new();
            let pointer = self.get_pointer().await?;

            // Check if already in storage
            if self.storage.has(&pointer).await? {
                return Ok((pointer, blocks));
            }

            // Serialize this node
            let entries = self.get_entries().await?;
            let node_data = util::serialize_node_data(&entries).await?;
            let cbor = serde_ipld_dagcbor::to_vec(&node_data).map_err(|e| {
                RepoError::serialization(e)
                    .with_context("serializing MST node for block collection")
            })?;
            blocks.insert(pointer, Bytes::from(cbor));

            // Recursively collect from subtrees
            for entry in &entries {
                if let NodeEntry::Tree(subtree) = entry {
                    let (_, subtree_blocks) = subtree.collect_blocks_sequential().await?;
                    blocks.extend(subtree_blocks);
                }
            }

            Ok((pointer, blocks))
        })
    }

    /// Persist all unstored blocks to storage
    ///
    /// Convenience method that calls `collect_blocks()` and `put_many()`.
    /// Returns the root CID after persisting.
    pub async fn persist(&self) -> Result<IpldCid> {
        let (root_cid, blocks) = self.collect_blocks().await?;

        if !blocks.is_empty() {
            self.storage.put_many(blocks).await?;
        }

        Ok(root_cid)
    }

    /// Collect all MST node CIDs in this tree
    ///
    /// Returns all CIDs for MST nodes (internal nodes), not leaves.
    /// Used for diff calculation to determine which MST blocks are removed.
    ///
    /// Uses parallel traversal to collect CIDs from independent subtrees concurrently.
    pub async fn collect_node_cids(&self) -> Result<Vec<IpldCid>> {
        collect_node_cids_parallel(self.clone()).await
    }

    /// Collect all MST node CIDs sequentially (for benchmarking)
    pub async fn collect_node_cids_sequential(&self) -> Result<Vec<IpldCid>> {
        let mut cids = Vec::new();
        let pointer = self.get_pointer().await?;
        cids.push(pointer);

        let entries = self.get_entries().await?;
        for entry in &entries {
            if let NodeEntry::Tree(subtree) = entry {
                let subtree_cids = subtree.collect_node_cids_sequential().await?;
                cids.extend(subtree_cids);
            }
        }
        Ok(cids)
    }

    /// Get all CIDs in the merkle path to a key
    ///
    /// Returns a list of CIDs representing the proof path from root to the target key:
    /// - Always includes the root CID (this node's pointer)
    /// - If key exists, includes the record CID
    /// - Includes all intermediate MST node CIDs in the path
    ///
    /// This is used for generating merkle proofs for record existence/non-existence.
    ///
    /// # Example
    ///
    /// For a key that exists:
    /// - Returns: `[root_cid, intermediate_node_cid?, ..., record_cid]`
    ///
    /// For a key that doesn't exist:
    /// - Returns: `[root_cid, intermediate_node_cid?, ...]` (proves absence)
    pub fn cids_for_path<'a>(
        &'a self,
        key: &'a str,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<IpldCid>>> + Send + 'a>> {
        Box::pin(async move {
            validate_key(key)?;

            let mut cids = vec![self.get_pointer().await?];
            let entries = self.get_entries().await?;
            let index = Self::find_gt_or_equal_leaf_index_in(&entries, key) as isize;

            let found = self.at_index(index).await?;

            if let Some(NodeEntry::Leaf {
                key: leaf_key,
                value,
            }) = found
            {
                if leaf_key.as_str() == key {
                    cids.push(value);
                    return Ok(cids);
                }
            }

            // Not found at this level - check subtree before this index
            if let Some(NodeEntry::Tree(subtree)) = self.at_index(index - 1).await? {
                let mut subtree_cids = subtree.cids_for_path(key).await?;
                cids.append(&mut subtree_cids);
                return Ok(cids);
            }

            // Key not found in tree
            Ok(cids)
        })
    }

    /// serialize the tree as car bytes, returning the cid of the car and the car bytes
    pub fn serialize_tree<'a>(
        &'a self,
    ) -> Pin<Box<dyn Future<Output = Result<(IpldCid, Bytes)>> + Send + 'a>> {
        Box::pin(async move {
            let mut entries = self.get_entries().await?;
            let mut outdated: Vec<Self> = Vec::new();
            for entry in &entries {
                if let NodeEntry::Tree(mst) = entry {
                    let is_outdated = *mst.outdated_pointer.read().await;
                    if is_outdated {
                        outdated.push(mst.clone());
                    }
                }
            }

            if outdated.len() > 0 {
                for outdated_entry in &outdated {
                    let _ = outdated_entry.get_pointer().await?;
                }
                entries = self.get_entries().await?
            }
            let data = util::serialize_node_data(entries.as_slice()).await?;
            let bytes = serde_ipld_dagcbor::to_vec(&data).map_err(|e| {
                // Provide best available context for debugging serialization failures
                let context = if let Some(left_cid) = &data.left {
                    format!("serializing MST node (left pointer: {})", left_cid)
                } else {
                    // No left pointer - use last leaf key as identifier
                    let last_key = entries.iter().rev().find_map(|e| match e {
                        NodeEntry::Leaf { key, .. } => Some(key.as_str()),
                        _ => None,
                    });
                    match last_key {
                        Some(k) => format!("serializing MST node (last key: {})", k),
                        None => "serializing MST node (empty or tree-only)".to_string(),
                    }
                };
                RepoError::serialization(e).with_context(context)
            })?;
            let cid = util::compute_cid(&bytes)?;

            Ok((cid, Bytes::from_owner(bytes)))
        })
    }

    /// Find the node at the given index if any
    pub async fn at_index(&self, index: isize) -> Result<Option<NodeEntry<S>>> {
        let entries = self.get_entries().await?;
        if index < 0 || index as usize >= entries.len() {
            return Ok(None);
        }
        Ok(entries
            .into_iter()
            .nth(index as usize)
            .map(|entry| entry.clone()))
    }

    /// Add any relevant blocks along the path to the given key to the map
    pub fn blocks_for_path<'a>(
        &'a self,
        key: &'a str,
        blocks: &'a mut BTreeMap<IpldCid, bytes::Bytes>,
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
        Box::pin(async move {
            validate_key(key)?;
            let (cid, bytes) = self.serialize_tree().await?;
            blocks.insert(cid, bytes);

            let entries = self.get_entries().await?;
            let index = Self::find_gt_or_equal_leaf_index_in(&entries, key) as isize;
            let found = self.at_index(index).await?;

            if let Some(NodeEntry::Leaf { key: leaf_key, .. }) = found {
                if leaf_key.as_str() == key {
                    return Ok(());
                }
            }
            if let Some(NodeEntry::Tree(subtree)) = self.at_index(index - 1).await? {
                subtree.blocks_for_path(key, blocks).await?;
                return Ok(());
            }

            // Key not found in tree
            Ok(())
        })
    }

    /// Write all MST and record blocks to CAR writer
    ///
    /// Streams blocks directly to the writer as the tree is walked:
    /// - All MST node blocks (read from storage)
    /// - All leaf record blocks (read from storage)
    ///
    /// This is suitable for CAR export and avoids loading all blocks into memory.
    ///
    /// TODO: get rid of tokio dependency here if possible
    pub async fn write_blocks_to_car<W: tokio::io::AsyncWrite + Send + Unpin>(
        &self,
        writer: &mut iroh_car::CarWriter<W>,
    ) -> Result<()> {
        let mut leaf_cids = Vec::new();

        // Walk tree, writing MST nodes and collecting leaf CIDs
        self.write_mst_nodes_to_car(writer, &mut leaf_cids).await?;

        // Fetch and write all leaf record blocks
        let leaf_blocks = self.storage.get_many(&leaf_cids).await?;
        for (cid, maybe_data) in leaf_cids.iter().zip(leaf_blocks) {
            if let Some(data) = maybe_data {
                writer
                    .write(*cid, &data)
                    .await
                    .map_err(|e| RepoError::car_write(cid, e))?;
            }
        }

        Ok(())
    }

    /// Recursively write MST nodes to CAR and collect leaf CIDs
    ///
    /// TODO: get rid of tokio dependency here if possible
    fn write_mst_nodes_to_car<'a, W: tokio::io::AsyncWrite + Send + Unpin>(
        &'a self,
        writer: &'a mut iroh_car::CarWriter<W>,
        leaf_cids: &'a mut Vec<IpldCid>,
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
        Box::pin(async move {
            let pointer = self.get_pointer().await?;

            // Read MST node from storage and write to CAR
            let node_bytes = self
                .storage
                .get(&pointer)
                .await?
                .ok_or_else(|| RepoError::not_found("MST node", &pointer))?;

            writer
                .write(pointer, &node_bytes)
                .await
                .map_err(|e| RepoError::car_write(&pointer, e))?;

            // Parse to get entries
            let entries = self.get_entries().await?;

            // Collect leaf CIDs and recurse into subtrees
            for entry in &entries {
                match entry {
                    NodeEntry::Leaf { value, .. } => {
                        leaf_cids.push(*value);
                    }
                    NodeEntry::Tree(subtree) => {
                        subtree.write_mst_nodes_to_car(writer, leaf_cids).await?;
                    }
                }
            }

            Ok(())
        })
    }
}

/// Recursively collect MST node CIDs in parallel
///
/// Spawns concurrent tasks for each subtree branch, then merges results.
fn collect_node_cids_parallel<S: BlockStore + Sync + Send + 'static>(
    tree: Mst<S>,
) -> Pin<Box<dyn Future<Output = Result<Vec<IpldCid>>> + Send>> {
    Box::pin(async move {
        let pointer = tree.get_pointer().await?;
        let entries = tree.get_entries().await?;

        // Spawn tasks for each subtree
        let tasks: Vec<_> = entries
            .into_iter()
            .filter_map(|entry| {
                if let NodeEntry::Tree(subtree) = entry {
                    Some(n0_future::task::spawn(async move {
                        collect_node_cids_parallel(subtree).await
                    }))
                } else {
                    None
                }
            })
            .collect();

        // Await all tasks concurrently
        let results = try_join_all(tasks).await.map_err(RepoError::task_failed)?;

        // Flatten results
        let mut cids = vec![pointer];
        for subtree_cids in results {
            cids.extend(subtree_cids?);
        }

        Ok(cids)
    })
}

/// Recursively collect leaves in parallel
///
/// Spawns concurrent tasks for each subtree branch, preserving lexicographic order.
fn collect_leaves_parallel<S: BlockStore + Sync + Send + 'static>(
    tree: Mst<S>,
) -> Pin<Box<dyn Future<Output = Result<Vec<(smol_str::SmolStr, IpldCid)>>> + Send>> {
    Box::pin(async move {
        let entries = tree.get_entries().await?;
        let mut result = Vec::new();

        // Collect tasks and immediate leaves in order
        let mut tasks = Vec::new();
        let mut task_positions = Vec::new();

        for (i, entry) in entries.into_iter().enumerate() {
            match entry {
                NodeEntry::Tree(subtree) => {
                    task_positions.push(i);
                    tasks.push(n0_future::task::spawn(async move {
                        collect_leaves_parallel(subtree).await
                    }));
                }
                NodeEntry::Leaf { key, value } => {
                    result.push((i, vec![(key, value)]));
                }
            }
        }

        // Await all tasks concurrently
        if !tasks.is_empty() {
            let subtree_results = try_join_all(tasks).await.map_err(RepoError::task_failed)?;

            for (pos, leaves) in task_positions.into_iter().zip(subtree_results) {
                result.push((pos, leaves?));
            }
        }

        // Sort by position and flatten
        result.sort_by_key(|(pos, _)| *pos);
        Ok(result.into_iter().flat_map(|(_, leaves)| leaves).collect())
    })
}

/// Recursively collect blocks in parallel
///
/// Spawns concurrent tasks for each subtree branch, then merges results.
fn collect_blocks_parallel<S: BlockStore + Sync + Send + 'static>(
    tree: Mst<S>,
) -> Pin<
    Box<
        dyn Future<Output = Result<(IpldCid, std::collections::BTreeMap<IpldCid, bytes::Bytes>)>>
            + Send,
    >,
> {
    Box::pin(async move {
        use bytes::Bytes;
        use std::collections::BTreeMap;

        let pointer = tree.get_pointer().await?;
        let mut blocks = BTreeMap::new();

        // Check if already in storage
        if tree.storage.has(&pointer).await? {
            return Ok((pointer, blocks));
        }

        // Serialize this node
        let entries = tree.get_entries().await?;
        let node_data = util::serialize_node_data(&entries).await?;
        let cbor = serde_ipld_dagcbor::to_vec(&node_data).map_err(|e| {
            RepoError::serialization(e)
                .with_context("serializing MST node for parallel block collection")
        })?;
        blocks.insert(pointer, Bytes::from(cbor));

        // Spawn tasks for each subtree
        let tasks: Vec<_> = entries
            .into_iter()
            .filter_map(|entry| {
                if let NodeEntry::Tree(subtree) = entry {
                    Some(n0_future::task::spawn(async move {
                        collect_blocks_parallel(subtree).await
                    }))
                } else {
                    None
                }
            })
            .collect();

        // Await all tasks concurrently
        if !tasks.is_empty() {
            let results = try_join_all(tasks).await.map_err(RepoError::task_failed)?;

            for subtree_result in results {
                let (_, subtree_blocks) = subtree_result?;
                blocks.extend(subtree_blocks);
            }
        }

        Ok((pointer, blocks))
    })
}

impl<S: BlockStore> std::fmt::Debug for Mst<S> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("MST")
            .field("entries", &self.entries.try_read().unwrap())
            .field("layer", &self.layer)
            .field("pointer", &self.pointer.try_read().unwrap().to_string())
            .field(
                "outdated_pointer",
                &self.outdated_pointer.try_read().unwrap(),
            )
            .finish()
    }
}

/// Format a CID for display (shortens long CIDs)
///
/// Truncates long CIDs to first 7 and last 8 characters with `...` in between.
pub fn short_cid(cid: &IpldCid) -> String {
    let cid_string = cid.to_string();
    let len = cid_string.len();
    if len > 15 {
        let first = &cid_string[0..7];
        let last = &cid_string[len - 8..];
        format!("{}...{}", first, last)
    } else {
        cid_string
    }
}

impl<S: BlockStore> Display for Mst<S> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        fn pointer_str<S: BlockStore>(mst: &Mst<S>) -> String {
            let cid_guard = mst.pointer.try_read().unwrap();
            format!("*({})", short_cid(&*cid_guard))
        }

        fn fmt_mst<S: BlockStore>(
            mst: &Mst<S>,
            f: &mut Formatter<'_>,
            prefix: &str,
            is_last: bool,
        ) -> fmt::Result {
            // Print MST pointer using our helper
            writeln!(
                f,
                "{}{}── {}",
                prefix,
                if is_last { "└" } else { "├" },
                pointer_str(mst),
            )?;

            // Prepare the child prefix
            let child_prefix = format!("{}{}", prefix, if is_last { "   " } else { "│  " });

            let entries_guard = mst.entries.try_read().unwrap();
            let entries = match &*entries_guard {
                Some(e) => e,
                None => {
                    writeln!(f, "{}(virtual node)", child_prefix)?;
                    return Ok(());
                }
            };

            for (i, entry) in entries.iter().enumerate() {
                let last_child = i == entries.len() - 1;
                match entry {
                    NodeEntry::Leaf { key, value } => {
                        // Print leaf key and (short) leaf value
                        writeln!(
                            f,
                            "{}{}── {} -> {}",
                            child_prefix,
                            if last_child { "└" } else { "├" },
                            key,
                            short_cid(&value)
                        )?;
                    }
                    NodeEntry::Tree(child_mst) => {
                        // Recurse
                        fmt_mst(child_mst, f, &child_prefix, last_child)?;
                    }
                }
            }
            Ok(())
        }

        // Start with empty prefix for the root
        fmt_mst(self, f, "", true)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{DAG_CBOR_CID_CODEC, storage::memory::MemoryBlockStore};
    use jacquard_common::types::crypto::SHA2_256;
    use smol_str::SmolStr;

    fn test_cid(n: u8) -> IpldCid {
        let data = vec![n; 32];
        let mh = multihash::Multihash::wrap(SHA2_256, &data).unwrap();
        IpldCid::new_v1(DAG_CBOR_CID_CODEC, mh)
    }

    #[tokio::test]
    async fn test_empty_tree() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let entries = mst.get_entries().await.unwrap();
        assert_eq!(entries.len(), 0);
    }

    #[tokio::test]
    async fn test_get_from_empty() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let result = mst.get("com.example.test/key").await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_manually_constructed_tree() {
        // Test with manually constructed entries (no CBOR)
        let storage = Arc::new(MemoryBlockStore::new());

        let entries = vec![
            NodeEntry::Leaf {
                key: SmolStr::new("com.example.test/a"),
                value: test_cid(1),
            },
            NodeEntry::Leaf {
                key: SmolStr::new("com.example.test/b"),
                value: test_cid(2),
            },
            NodeEntry::Leaf {
                key: SmolStr::new("com.example.test/c"),
                value: test_cid(3),
            },
        ];

        let mst = Mst::create(storage, entries, Some(0)).await.unwrap();

        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
        assert_eq!(
            mst.get("com.example.test/b").await.unwrap(),
            Some(test_cid(2))
        );
        assert_eq!(
            mst.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        );
        assert_eq!(mst.get("com.example.test/d").await.unwrap(), None);
    }

    #[tokio::test]
    async fn test_allowed_keys() -> Result<()> {
        let cid1str = "bafyreie5cvv4h45feadgeuwhbcutmh6t2ceseocckahdoe6uat64zmz454";
        let cid1 = IpldCid::try_from(cid1str).unwrap();

        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);
        // Rejects empty key
        let result = mst.add(&"".to_string(), cid1).await;
        assert!(result.is_err());

        // Rejects a key with no collection
        let result = mst.add(&"asdf".to_string(), cid1).await;
        assert!(result.is_err());

        // Rejects a key with a nested collection
        let result = mst.add(&"nested/collection/asdf".to_string(), cid1).await;
        assert!(result.is_err());

        // Rejects on empty coll or rkey
        let result = mst.add(&"coll/".to_string(), cid1).await;
        assert!(result.is_err());
        let result = mst.add(&"/rkey".to_string(), cid1).await;
        assert!(result.is_err());

        // Rejects non-ascii chars
        let result = mst.add(&"coll/jalapeñoA".to_string(), cid1).await;
        assert!(result.is_err());
        let result = mst.add(&"coll/coöperative".to_string(), cid1).await;
        assert!(result.is_err());
        let result = mst.add(&"coll/abc💩".to_string(), cid1).await;
        assert!(result.is_err());

        // Rejects ascii that we don't support
        let invalid_chars = vec!["$", "%", "(", ")", "+", "="];
        for ch in invalid_chars {
            let key = format!("coll/key{}", ch);
            let result = mst.add(&key, cid1).await;
            assert!(result.is_err(), "Key '{}' should be invalid", key);
        }

        // Rejects keys over 256 chars
        let long_key: String = "a".repeat(253);
        let key = format!("coll/{}", long_key);
        let result = mst.add(&key, cid1).await;
        assert!(result.is_err());

        // Allows long key under 256 chars
        let long_key: String = "a".repeat(250);
        let key = format!("coll/{}", long_key);
        let result = mst.add(&key, cid1).await;
        assert!(result.is_ok());

        // Allows URL-safe chars
        let valid_keys = vec![
            "coll/key0",
            "coll/key_",
            "coll/key:",
            "coll/key.",
            "coll/key-",
        ];
        for key in valid_keys {
            let result = mst.add(&key.to_string(), cid1).await;
            assert!(result.is_ok(), "Key '{}' should be valid", key);
        }

        Ok(())
    }

    #[tokio::test]
    async fn test_add_single_key() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let updated = mst.add("com.example.test/key", test_cid(1)).await.unwrap();

        assert_eq!(
            updated.get("com.example.test/key").await.unwrap(),
            Some(test_cid(1))
        );
    }

    #[tokio::test]
    async fn test_add_multiple_keys() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
        assert_eq!(
            mst.get("com.example.test/b").await.unwrap(),
            Some(test_cid(2))
        );
        assert_eq!(
            mst.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        );
    }

    #[tokio::test]
    async fn test_add_replace_key() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/key1", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/key1", test_cid(2)).await.unwrap();

        assert_eq!(
            mst.get("com.example.test/key1").await.unwrap(),
            Some(test_cid(2))
        );
    }

    #[tokio::test]
    async fn test_delete_single_key() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/key1", test_cid(1)).await.unwrap();
        let mst = mst.delete("com.example.test/key1").await.unwrap();

        assert_eq!(mst.get("com.example.test/key1").await.unwrap(), None);
        assert_eq!(mst.get_entries().await.unwrap().len(), 0);
    }

    #[tokio::test]
    async fn test_delete_from_multi_key_tree() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        let mst = mst.delete("com.example.test/b").await.unwrap();

        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
        assert_eq!(mst.get("com.example.test/b").await.unwrap(), None);
        assert_eq!(
            mst.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        );
    }

    #[tokio::test]
    async fn test_delete_nonexistent_key() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();

        let result = mst.delete("com.example.test/b").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_serialization_roundtrip() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage.clone());

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        // Persist to storage
        let cid = mst.persist().await.unwrap();

        // Load from storage
        let reloaded = Mst::load(storage, cid, Some(0));

        // Verify all keys are present
        assert_eq!(
            reloaded.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
        assert_eq!(
            reloaded.get("com.example.test/b").await.unwrap(),
            Some(test_cid(2))
        );
        assert_eq!(
            reloaded.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        );
    }

    #[tokio::test]
    async fn test_cid_determinism() {
        // Same keys inserted in same order should produce same CID
        let storage1 = Arc::new(MemoryBlockStore::new());
        let mst1 = Mst::new(storage1);
        let mst1 = mst1.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst1 = mst1.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst1 = mst1.add("com.example.test/c", test_cid(3)).await.unwrap();
        let cid1 = mst1.get_pointer().await.unwrap();

        let storage2 = Arc::new(MemoryBlockStore::new());
        let mst2 = Mst::new(storage2);
        let mst2 = mst2.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst2 = mst2.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst2 = mst2.add("com.example.test/c", test_cid(3)).await.unwrap();
        let cid2 = mst2.get_pointer().await.unwrap();

        assert_eq!(cid1, cid2);
    }

    #[tokio::test]
    async fn test_insertion_order_determinism() {
        // Different insertion orders should produce same CID
        let storage1 = Arc::new(MemoryBlockStore::new());
        let mst1 = Mst::new(storage1);
        let mst1 = mst1.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst1 = mst1.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst1 = mst1.add("com.example.test/c", test_cid(3)).await.unwrap();
        let cid1 = mst1.get_pointer().await.unwrap();

        let storage2 = Arc::new(MemoryBlockStore::new());
        let mst2 = Mst::new(storage2);
        let mst2 = mst2.add("com.example.test/c", test_cid(3)).await.unwrap();
        let mst2 = mst2.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst2 = mst2.add("com.example.test/b", test_cid(2)).await.unwrap();
        let cid2 = mst2.get_pointer().await.unwrap();

        assert_eq!(cid1, cid2);
    }

    #[tokio::test]
    async fn test_batch_creates() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let ops = vec![
            VerifiedWriteOp::Create {
                key: SmolStr::new("com.example.test/a"),
                cid: test_cid(1),
            },
            VerifiedWriteOp::Create {
                key: SmolStr::new("com.example.test/b"),
                cid: test_cid(2),
            },
            VerifiedWriteOp::Create {
                key: SmolStr::new("com.example.test/c"),
                cid: test_cid(3),
            },
        ];

        let mst = mst.batch(&ops).await.unwrap();

        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
        assert_eq!(
            mst.get("com.example.test/b").await.unwrap(),
            Some(test_cid(2))
        );
        assert_eq!(
            mst.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        );
    }

    #[tokio::test]
    async fn test_batch_mixed_operations() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        // Start with some keys
        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        let ops = vec![
            VerifiedWriteOp::Create {
                key: SmolStr::new("com.example.test/d"),
                cid: test_cid(4),
            },
            VerifiedWriteOp::Update {
                key: SmolStr::new("com.example.test/a"),
                cid: test_cid(10),
                prev: test_cid(1),
            },
            VerifiedWriteOp::Delete {
                key: SmolStr::new("com.example.test/b"),
                prev: test_cid(2),
            },
        ];

        let mst = mst.batch(&ops).await.unwrap();

        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(10))
        ); // Updated
        assert_eq!(mst.get("com.example.test/b").await.unwrap(), None); // Deleted
        assert_eq!(
            mst.get("com.example.test/c").await.unwrap(),
            Some(test_cid(3))
        ); // Unchanged
        assert_eq!(
            mst.get("com.example.test/d").await.unwrap(),
            Some(test_cid(4))
        ); // Created
    }

    #[tokio::test]
    async fn test_batch_with_prev_validation() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);
        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();

        // Update with correct prev - should succeed
        let ops = vec![VerifiedWriteOp::Update {
            key: SmolStr::new("com.example.test/a"),
            cid: test_cid(2),
            prev: test_cid(1),
        }];
        let mst = mst.batch(&ops).await.unwrap();
        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(2))
        );

        // Update with wrong prev - should fail
        let ops = vec![VerifiedWriteOp::Update {
            key: SmolStr::new("com.example.test/a"),
            cid: test_cid(3),
            prev: test_cid(99), // Wrong CID
        }];
        assert!(mst.batch(&ops).await.is_err());

        // Delete with correct prev - should succeed
        let ops = vec![VerifiedWriteOp::Delete {
            key: SmolStr::new("com.example.test/a"),
            prev: test_cid(2),
        }];
        let mst = mst.batch(&ops).await.unwrap();
        assert_eq!(mst.get("com.example.test/a").await.unwrap(), None);
    }

    #[tokio::test]
    async fn test_batch_create_duplicate_error() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);
        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();

        let ops = vec![VerifiedWriteOp::Create {
            key: SmolStr::new("com.example.test/a"),
            cid: test_cid(2),
        }];

        // Should error because key already exists
        assert!(mst.batch(&ops).await.is_err());
    }

    #[tokio::test]
    async fn test_batch_update_nonexistent_error() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let ops = vec![VerifiedWriteOp::Update {
            key: SmolStr::new("com.example.test/a"),
            cid: test_cid(1),
            prev: test_cid(99), // Doesn't matter since key doesn't exist
        }];

        // Should error because key doesn't exist
        assert!(mst.batch(&ops).await.is_err());
    }

    #[tokio::test]
    async fn test_batch_delete_nonexistent_error() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let ops = vec![VerifiedWriteOp::Delete {
            key: SmolStr::new("com.example.test/a"),
            prev: test_cid(99), // Doesn't matter since key doesn't exist
        }];

        // Should error because key doesn't exist
        assert!(mst.batch(&ops).await.is_err());
    }

    #[tokio::test]
    async fn test_batch_empty() {
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);
        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();

        let ops = vec![];
        let mst = mst.batch(&ops).await.unwrap();

        // Should be unchanged
        assert_eq!(
            mst.get("com.example.test/a").await.unwrap(),
            Some(test_cid(1))
        );
    }

    #[tokio::test]
    async fn test_cids_for_path_simple() {
        // Test cids_for_path with a simple flat tree
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/b", test_cid(2)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        // Get proof path for key "b"
        let cids = mst.cids_for_path("com.example.test/b").await.unwrap();

        // Should contain: root CID, record CID
        assert_eq!(cids.len(), 2);
        assert_eq!(cids[0], mst.get_pointer().await.unwrap());
        assert_eq!(cids[1], test_cid(2));
    }

    #[tokio::test]
    async fn test_cids_for_path_nonexistent() {
        // Test cids_for_path with a key that doesn't exist
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        let mst = mst.add("com.example.test/a", test_cid(1)).await.unwrap();
        let mst = mst.add("com.example.test/c", test_cid(3)).await.unwrap();

        // Get proof path for nonexistent key "b"
        let cids = mst.cids_for_path("com.example.test/b").await.unwrap();

        // Should contain root CID first, and NOT contain the record CID (proves absence)
        assert!(cids.len() >= 1, "Should have at least the root CID");
        assert_eq!(
            cids[0],
            mst.get_pointer().await.unwrap(),
            "First CID should be root"
        );
        assert!(
            !cids.contains(&test_cid(2)),
            "Should not contain nonexistent record"
        );
    }

    #[tokio::test]
    async fn test_cids_for_path_multiple_records() {
        // Test with multiple records across different collection paths (like rsky)
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        // Simulate records from different collections (app.bsky.feed.post, app.bsky.feed.like)
        let keys = vec![
            "app.bsky.feed.post/3l4qpz7ajrc2a",
            "app.bsky.feed.post/3l4qpz7bjrc2b",
            "app.bsky.feed.like/3l4qpz7cjrc2c",
            "app.bsky.feed.like/3l4qpz7djrc2d",
            "app.bsky.graph.follow/3l4qpz7ejrc2e",
        ];

        let mut mst = mst;
        for (i, key) in keys.iter().enumerate() {
            mst = mst.add(key, test_cid((i + 1) as u8)).await.unwrap();
        }

        // Get proof for each record
        for (i, key) in keys.iter().enumerate() {
            let cids = mst.cids_for_path(key).await.unwrap();

            // Should have root CID first
            assert_eq!(cids[0], mst.get_pointer().await.unwrap());

            // Should have record CID last (since record exists)
            assert_eq!(*cids.last().unwrap(), test_cid((i + 1) as u8));

            // Should have at least root + record
            assert!(cids.len() >= 2);
        }
    }

    #[tokio::test]
    async fn test_cids_for_path_proves_nonexistence() {
        // Test that we can prove a record doesn't exist in a tree with many records
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        // Add several records
        let existing_keys = vec![
            "com.example.posts/key1",
            "com.example.posts/key2",
            "com.example.posts/key4",
            "com.example.posts/key5",
        ];

        let mut mst = mst;
        for (i, key) in existing_keys.iter().enumerate() {
            mst = mst.add(key, test_cid((i + 1) as u8)).await.unwrap();
        }

        // Prove key3 doesn't exist (between key2 and key4)
        let nonexistent_key = "com.example.posts/key3";
        let cids = mst.cids_for_path(nonexistent_key).await.unwrap();

        // Should have root CID
        assert_eq!(cids[0], mst.get_pointer().await.unwrap());

        // Should NOT contain a record CID for key3
        assert!(!cids.contains(&test_cid(3)));

        // Proof is just the path showing where key3 WOULD be (proves absence)
        assert!(cids.len() >= 1);
    }

    #[tokio::test]
    async fn test_cids_for_path_collection_structure() {
        // Test proof generation for realistic collection/rkey structure
        let storage = Arc::new(MemoryBlockStore::new());
        let mst = Mst::new(storage);

        // Simulate a repo with multiple collections and records
        let records = vec![
            ("com.atproto.repo.strongRef", "abc123", test_cid(1)),
            ("app.bsky.feed.post", "post1", test_cid(2)),
            ("app.bsky.feed.post", "post2", test_cid(3)),
            ("app.bsky.feed.like", "like1", test_cid(4)),
            ("app.bsky.graph.follow", "follow1", test_cid(5)),
        ];

        let mut mst = mst;
        for (collection, rkey, cid) in &records {
            let key = format!("{}/{}", collection, rkey);
            mst = mst.add(&key, *cid).await.unwrap();
        }

        // Persist to storage so we have real MST blocks
        let root_cid = mst.persist().await.unwrap();
        assert_eq!(root_cid, mst.get_pointer().await.unwrap());

        // Get proofs for each record
        for (collection, rkey, expected_cid) in &records {
            let key = format!("{}/{}", collection, rkey);
            let cids = mst.cids_for_path(&key).await.unwrap();

            // Verify structure
            assert_eq!(cids[0], root_cid, "First CID should be root");
            assert_eq!(
                *cids.last().unwrap(),
                *expected_cid,
                "Last CID should be record"
            );
        }
    }
}
