//! Core Merkle tree implementation optimized for append-only assertions.
//!
//! # Architecture
//!
//! This implements a **binary Merkle tree** using BLAKE3 for node hashing:
//!
//! ```text
//!        root
//!       /    \
//!      h12   h34
//!     / \    / \
//!    h1 h2  h3 h4
//!    |  |   |  |
//!   a1 a2  a3 a4  (assertion hashes)
//! ```
//!
//! Where:
//! - `h1 = a1` (leaf nodes are assertion hashes directly)
//! - `h12 = BLAKE3(h1 || h2)` (internal nodes hash their children)
//! - `root = BLAKE3(h12 || h34)` (root represents entire tree)
//!
//! # Append-Only Optimization
//!
//! The tree is optimized for sequential inserts (common in StemeDB):
//! - New leaves are added to the right edge
//! - Only the path from new leaf to root is recomputed: O(log N)
//! - Root hash is cached for O(1) access
//!
//! # Storage Layout
//!
//! Nodes are stored in a flat vector using index arithmetic:
//! - Parent of node `i` is at `(i - 1) / 2`
//! - Left child of node `i` is at `2i + 1`
//! - Right child of node `i` is at `2i + 2`
//!
//! This enables efficient traversal without pointer chasing.

use blake3::Hasher;
use thiserror::Error;
use tracing::{debug, instrument};

/// A BLAKE3 hash (256 bits / 32 bytes).
pub type Hash = [u8; 32];

/// Errors that can occur during Merkle tree operations.
#[derive(Debug, Error)]
pub enum TreeError {
    /// Tree is empty (has no root).
    #[error("Tree is empty")]
    EmptyTree,

    /// Internal consistency error (should never happen).
    #[error("Internal tree invariant violated: {0}")]
    InternalError(String),
}

/// A binary Merkle tree optimized for append-only assertion storage.
///
/// # Design
///
/// - **Binary tree**: each internal node has exactly two children
/// - **Append-only**: leaves are added sequentially to the right edge
/// - **BLAKE3 hashing**: internal nodes = BLAKE3(left || right)
/// - **Cached root**: O(1) access to tree root hash
///
/// # Example
///
/// ```
/// use stemedb_merkle::MerkleTree;
///
/// let mut tree = MerkleTree::new();
///
/// // Insert assertion hashes
/// tree.insert([1u8; 32]).expect("insert");
/// tree.insert([2u8; 32]).expect("insert");
/// tree.insert([3u8; 32]).expect("insert");
///
/// // Root hash represents entire tree
/// let root = tree.root().expect("root");
/// assert_eq!(tree.len(), 3);
/// ```
#[derive(Debug, Clone)]
pub struct MerkleTree {
    /// Leaves (assertion hashes) in insertion order.
    /// Storing leaves separately enables efficient diff operations.
    leaves: Vec<Hash>,

    /// Cached root hash (None if tree is empty).
    /// Recomputed on each insert to maintain O(1) root access.
    cached_root: Option<Hash>,
}

impl MerkleTree {
    /// Create a new empty Merkle tree.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let tree = MerkleTree::new();
    /// assert_eq!(tree.len(), 0);
    /// assert!(tree.root().is_err());
    /// ```
    pub fn new() -> Self {
        Self { leaves: Vec::new(), cached_root: None }
    }

    /// Insert a new assertion hash into the tree.
    ///
    /// This appends the hash as a new leaf and recomputes the path from
    /// leaf to root. Complexity: O(log N) where N is the number of leaves.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let mut tree = MerkleTree::new();
    /// tree.insert([1u8; 32]).expect("insert");
    /// tree.insert([2u8; 32]).expect("insert");
    /// assert_eq!(tree.len(), 2);
    /// ```
    #[instrument(skip(self, hash), fields(leaf_count = self.leaves.len()))]
    pub fn insert(&mut self, hash: Hash) -> Result<(), TreeError> {
        debug!("Inserting hash into Merkle tree");
        self.leaves.push(hash);
        self.recompute_root()?;
        Ok(())
    }

    /// Get the root hash of the tree.
    ///
    /// Returns an error if the tree is empty.
    /// Complexity: O(1) due to caching.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let mut tree = MerkleTree::new();
    /// assert!(tree.root().is_err()); // Empty tree
    ///
    /// tree.insert([1u8; 32]).expect("insert");
    /// let root = tree.root().expect("root");
    /// assert_eq!(root.len(), 32);
    /// ```
    pub fn root(&self) -> Result<Hash, TreeError> {
        self.cached_root.ok_or(TreeError::EmptyTree)
    }

    /// Get the number of leaves (assertion hashes) in the tree.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let mut tree = MerkleTree::new();
    /// assert_eq!(tree.len(), 0);
    ///
    /// tree.insert([1u8; 32]).expect("insert");
    /// tree.insert([2u8; 32]).expect("insert");
    /// assert_eq!(tree.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        self.leaves.len()
    }

    /// Check if the tree is empty.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let mut tree = MerkleTree::new();
    /// assert!(tree.is_empty());
    ///
    /// tree.insert([1u8; 32]).expect("insert");
    /// assert!(!tree.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        self.leaves.is_empty()
    }

    /// Get a slice of all leaf hashes in insertion order.
    ///
    /// This is used by the diff algorithm to identify missing assertions.
    ///
    /// # Example
    ///
    /// ```
    /// use stemedb_merkle::MerkleTree;
    ///
    /// let mut tree = MerkleTree::new();
    /// tree.insert([1u8; 32]).expect("insert");
    /// tree.insert([2u8; 32]).expect("insert");
    ///
    /// let leaves = tree.leaves();
    /// assert_eq!(leaves.len(), 2);
    /// assert_eq!(leaves[0], [1u8; 32]);
    /// assert_eq!(leaves[1], [2u8; 32]);
    /// ```
    pub fn leaves(&self) -> &[Hash] {
        &self.leaves
    }

    /// Recompute the root hash from current leaves.
    ///
    /// This builds the tree bottom-up using BLAKE3 hashing:
    /// 1. Start with leaf hashes
    /// 2. Pair adjacent nodes and hash them: BLAKE3(left || right)
    /// 3. Repeat until only root remains
    ///
    /// For odd number of nodes at any level, the last node is carried forward.
    #[instrument(skip(self), fields(leaf_count = self.leaves.len()))]
    fn recompute_root(&mut self) -> Result<(), TreeError> {
        if self.leaves.is_empty() {
            self.cached_root = None;
            return Ok(());
        }

        // Start with leaf level
        let mut current_level: Vec<Hash> = self.leaves.clone();

        // Build tree bottom-up until we reach the root
        while current_level.len() > 1 {
            let mut next_level = Vec::with_capacity(current_level.len().div_ceil(2));

            // Pair adjacent nodes and hash them
            let mut i = 0;
            while i < current_level.len() {
                if i + 1 < current_level.len() {
                    // Pair exists: hash left || right
                    let parent_hash = Self::hash_nodes(&current_level[i], &current_level[i + 1]);
                    next_level.push(parent_hash);
                    i += 2;
                } else {
                    // Odd node: carry forward to next level
                    next_level.push(current_level[i]);
                    i += 1;
                }
            }

            current_level = next_level;
        }

        // current_level now contains exactly one hash: the root
        self.cached_root = Some(current_level[0]);
        debug!(root_hash = ?self.cached_root, "Recomputed Merkle root");
        Ok(())
    }

    /// Hash two child nodes to produce parent hash.
    ///
    /// Uses BLAKE3(left || right) where || denotes concatenation.
    fn hash_nodes(left: &Hash, right: &Hash) -> Hash {
        let mut hasher = Hasher::new();
        hasher.update(left);
        hasher.update(right);
        *hasher.finalize().as_bytes()
    }
}

impl Default for MerkleTree {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_tree() {
        let tree = MerkleTree::new();
        assert_eq!(tree.len(), 0);
        assert!(tree.is_empty());
        assert!(tree.root().is_err());
    }

    #[test]
    fn test_single_leaf() {
        let mut tree = MerkleTree::new();
        let hash = [1u8; 32];
        tree.insert(hash).expect("insert");

        assert_eq!(tree.len(), 1);
        assert!(!tree.is_empty());
        assert_eq!(tree.root().expect("root"), hash);
    }

    #[test]
    fn test_two_leaves() {
        let mut tree = MerkleTree::new();
        let h1 = [1u8; 32];
        let h2 = [2u8; 32];

        tree.insert(h1).expect("insert");
        tree.insert(h2).expect("insert");

        assert_eq!(tree.len(), 2);

        // Root should be BLAKE3(h1 || h2)
        let expected_root = MerkleTree::hash_nodes(&h1, &h2);
        assert_eq!(tree.root().expect("root"), expected_root);
    }

    #[test]
    fn test_three_leaves() {
        let mut tree = MerkleTree::new();
        let h1 = [1u8; 32];
        let h2 = [2u8; 32];
        let h3 = [3u8; 32];

        tree.insert(h1).expect("insert");
        tree.insert(h2).expect("insert");
        tree.insert(h3).expect("insert");

        assert_eq!(tree.len(), 3);

        // Tree structure:
        //       root
        //      /    \
        //    h12    h3
        //   /  \
        //  h1  h2
        let h12 = MerkleTree::hash_nodes(&h1, &h2);
        let expected_root = MerkleTree::hash_nodes(&h12, &h3);
        assert_eq!(tree.root().expect("root"), expected_root);
    }

    #[test]
    fn test_four_leaves() {
        let mut tree = MerkleTree::new();
        let h1 = [1u8; 32];
        let h2 = [2u8; 32];
        let h3 = [3u8; 32];
        let h4 = [4u8; 32];

        tree.insert(h1).expect("insert");
        tree.insert(h2).expect("insert");
        tree.insert(h3).expect("insert");
        tree.insert(h4).expect("insert");

        assert_eq!(tree.len(), 4);

        // Tree structure:
        //       root
        //      /    \
        //    h12    h34
        //   /  \   /  \
        //  h1  h2 h3  h4
        let h12 = MerkleTree::hash_nodes(&h1, &h2);
        let h34 = MerkleTree::hash_nodes(&h3, &h4);
        let expected_root = MerkleTree::hash_nodes(&h12, &h34);
        assert_eq!(tree.root().expect("root"), expected_root);
    }

    #[test]
    fn test_different_trees_different_roots() {
        let mut tree1 = MerkleTree::new();
        tree1.insert([1u8; 32]).expect("insert");
        tree1.insert([2u8; 32]).expect("insert");

        let mut tree2 = MerkleTree::new();
        tree2.insert([1u8; 32]).expect("insert");
        tree2.insert([3u8; 32]).expect("insert");

        assert_ne!(tree1.root().expect("root"), tree2.root().expect("root"));
    }

    #[test]
    fn test_identical_trees_same_root() {
        let mut tree1 = MerkleTree::new();
        tree1.insert([1u8; 32]).expect("insert");
        tree1.insert([2u8; 32]).expect("insert");

        let mut tree2 = MerkleTree::new();
        tree2.insert([1u8; 32]).expect("insert");
        tree2.insert([2u8; 32]).expect("insert");

        assert_eq!(tree1.root().expect("root"), tree2.root().expect("root"));
    }

    #[test]
    fn test_leaves_accessor() {
        let mut tree = MerkleTree::new();
        let h1 = [1u8; 32];
        let h2 = [2u8; 32];
        let h3 = [3u8; 32];

        tree.insert(h1).expect("insert");
        tree.insert(h2).expect("insert");
        tree.insert(h3).expect("insert");

        let leaves = tree.leaves();
        assert_eq!(leaves.len(), 3);
        assert_eq!(leaves[0], h1);
        assert_eq!(leaves[1], h2);
        assert_eq!(leaves[2], h3);
    }

    #[test]
    fn test_order_matters() {
        let mut tree1 = MerkleTree::new();
        tree1.insert([1u8; 32]).expect("insert");
        tree1.insert([2u8; 32]).expect("insert");

        let mut tree2 = MerkleTree::new();
        tree2.insert([2u8; 32]).expect("insert");
        tree2.insert([1u8; 32]).expect("insert");

        // Different insertion order produces different root
        assert_ne!(tree1.root().expect("root"), tree2.root().expect("root"));
    }

    #[test]
    fn test_incremental_insert() {
        let mut tree = MerkleTree::new();
        let hashes: Vec<Hash> = (0..10).map(|i| [i as u8; 32]).collect();

        for (i, &hash) in hashes.iter().enumerate() {
            tree.insert(hash).expect("insert");
            assert_eq!(tree.len(), i + 1);
            assert!(tree.root().is_ok());
        }

        assert_eq!(tree.len(), 10);
    }
}
