//! A fully materialized Merkle mountain range (MMR).
//!
//! A MMR is a forest structure, i.e. it is an ordered set of disjoint rooted trees. The trees are
//! ordered by size, from the most to least number of leaves. Every tree is a perfect binary tree,
//! meaning a tree has all its leaves at the same depth, and every inner node has a branch-factor
//! of 2 with both children set.
//!
//! Additionally the structure only supports adding leaves to the right-most tree, the one with the
//! least number of leaves. The structure preserves the invariant that each tree has different
//! depths, i.e. as part of adding adding a new element to the forest the trees with same depth are
//! merged, creating a new tree with depth d+1, this process is continued until the property is
//! restabilished.
use super::{
    super::{InnerNodeInfo, MerklePath, RpoDigest, Vec},
    bit::TrueBitPositionIterator,
    MmrPeaks, MmrProof, Rpo256,
};
use core::fmt::{Display, Formatter};

#[cfg(feature = "std")]
use std::error::Error;

// MMR
// ===============================================================================================

/// A fully materialized Merkle Mountain Range, with every tree in the forest and all their
/// elements.
///
/// Since this is a full representation of the MMR, elements are never removed and the MMR will
/// grow roughly `O(2n)` in number of leaf elements.
pub struct Mmr {
    /// Refer to the `forest` method documentation for details of the semantics of this value.
    pub(super) forest: usize,

    /// Contains every element of the forest.
    ///
    /// The trees are in postorder sequential representation. This representation allows for all
    /// the elements of every tree in the forest to be stored in the same sequential buffer. It
    /// also means new elements can be added to the forest, and merging of trees is very cheap with
    /// no need to copy elements.
    pub(super) nodes: Vec<RpoDigest>,
}

#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum MmrError {
    InvalidPosition(usize),
}

impl Display for MmrError {
    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
        match self {
            MmrError::InvalidPosition(pos) => write!(fmt, "Mmr does not contain position {pos}"),
        }
    }
}

#[cfg(feature = "std")]
impl Error for MmrError {}

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

impl Mmr {
    // CONSTRUCTORS
    // ============================================================================================

    /// Constructor for an empty `Mmr`.
    pub fn new() -> Mmr {
        Mmr {
            forest: 0,
            nodes: Vec::new(),
        }
    }

    // ACCESSORS
    // ============================================================================================

    /// Returns the MMR forest representation.
    ///
    /// The forest value has the following interpretations:
    /// - its value is the number of elements in the forest
    /// - bit count corresponds to the number of trees in the forest
    /// - each true bit position determines the depth of a tree in the forest
    pub const fn forest(&self) -> usize {
        self.forest
    }

    // FUNCTIONALITY
    // ============================================================================================

    /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
    /// is greater-or-equal than the tree size an error is returned.
    ///
    /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
    /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
    /// has position 0, the second position 1, and so on.
    pub fn open(&self, pos: usize) -> Result<MmrProof, MmrError> {
        // find the target tree responsible for the MMR position
        let tree_bit =
            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
        let forest_target = 1usize << tree_bit;

        // isolate the trees before the target
        let forest_before = self.forest & high_bitmask(tree_bit + 1);
        let index_offset = nodes_in_forest(forest_before);

        // find the root
        let index = nodes_in_forest(forest_target) - 1;

        // update the value position from global to the target tree
        let relative_pos = pos - forest_before;

        // collect the path and the final index of the target value
        let (_, path) =
            self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);

        Ok(MmrProof {
            forest: self.forest,
            position: pos,
            merkle_path: MerklePath::new(path),
        })
    }

    /// Returns the leaf value at position `pos`.
    ///
    /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
    /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
    /// has position 0, the second position 1, and so on.
    pub fn get(&self, pos: usize) -> Result<RpoDigest, MmrError> {
        // find the target tree responsible for the MMR position
        let tree_bit =
            leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
        let forest_target = 1usize << tree_bit;

        // isolate the trees before the target
        let forest_before = self.forest & high_bitmask(tree_bit + 1);
        let index_offset = nodes_in_forest(forest_before);

        // find the root
        let index = nodes_in_forest(forest_target) - 1;

        // update the value position from global to the target tree
        let relative_pos = pos - forest_before;

        // collect the path and the final index of the target value
        let (value, _) =
            self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);

        Ok(value)
    }

    /// Adds a new element to the MMR.
    pub fn add(&mut self, el: RpoDigest) {
        // Note: every node is also a tree of size 1, adding an element to the forest creates a new
        // rooted-tree of size 1. This may temporarily break the invariant that every tree in the
        // forest has different sizes, the loop below will eagerly merge trees of same size and
        // restore the invariant.
        self.nodes.push(el);

        let mut left_offset = self.nodes.len().saturating_sub(2);
        let mut right = el;
        let mut left_tree = 1;
        while self.forest & left_tree != 0 {
            right = Rpo256::merge(&[self.nodes[left_offset], right]);
            self.nodes.push(right);

            left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree));
            left_tree <<= 1;
        }

        self.forest += 1;
    }

    /// Returns an accumulator representing the current state of the MMR.
    pub fn accumulator(&self) -> MmrPeaks {
        let peaks: Vec<RpoDigest> = TrueBitPositionIterator::new(self.forest)
            .rev()
            .map(|bit| nodes_in_forest(1 << bit))
            .scan(0, |offset, el| {
                *offset += el;
                Some(*offset)
            })
            .map(|offset| self.nodes[offset - 1])
            .collect();

        MmrPeaks {
            num_leaves: self.forest,
            peaks,
        }
    }

    /// An iterator over inner nodes in the MMR. The order of iteration is unspecified.
    pub fn inner_nodes(&self) -> MmrNodes {
        MmrNodes {
            mmr: self,
            forest: 0,
            last_right: 0,
            index: 0,
        }
    }

    // UTILITIES
    // ============================================================================================

    /// Internal function used to collect the Merkle path of a value.
    fn collect_merkle_path_and_value(
        &self,
        tree_bit: u32,
        relative_pos: usize,
        index_offset: usize,
        mut index: usize,
    ) -> (RpoDigest, Vec<RpoDigest>) {
        // collect the Merkle path
        let mut tree_depth = tree_bit as usize;
        let mut path = Vec::with_capacity(tree_depth + 1);
        while tree_depth > 0 {
            let bit = relative_pos & tree_depth;
            let right_offset = index - 1;
            let left_offset = right_offset - nodes_in_forest(tree_depth);

            // Elements to the right have a higher position because they were
            // added later. Therefore when the bit is true the node's path is
            // to the right, and its sibling to the left.
            let sibling = if bit != 0 {
                index = right_offset;
                self.nodes[index_offset + left_offset]
            } else {
                index = left_offset;
                self.nodes[index_offset + right_offset]
            };

            tree_depth >>= 1;
            path.push(sibling);
        }

        // the rest of the codebase has the elements going from leaf to root, adjust it here for
        // easy of use/consistency sake
        path.reverse();

        let value = self.nodes[index_offset + index];
        (value, path)
    }
}

impl<T> From<T> for Mmr
where
    T: IntoIterator<Item = RpoDigest>,
{
    fn from(values: T) -> Self {
        let mut mmr = Mmr::new();
        for v in values {
            mmr.add(v)
        }
        mmr
    }
}

// ITERATOR
// ===============================================================================================

/// Yields inner nodes of the [Mmr].
pub struct MmrNodes<'a> {
    /// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished.
    mmr: &'a Mmr,
    /// Keeps track of the left nodes yielded so far waiting for a right pair, this matches the
    /// semantics of the [Mmr]'s forest attribute, since that too works as a buffer of left nodes
    /// waiting for a pair to be hashed together.
    forest: usize,
    /// Keeps track of the last right node yielded, after this value is set, the next iteration
    /// will be its parent with its corresponding left node that has been yield already.
    last_right: usize,
    /// The current index in the `nodes` vector.
    index: usize,
}

impl<'a> Iterator for MmrNodes<'a> {
    type Item = InnerNodeInfo;

    fn next(&mut self) -> Option<Self::Item> {
        debug_assert!(self.last_right.count_ones() <= 1, "last_right tracks zero or one element");

        // only parent nodes are emitted, remove the single node tree from the forest
        let target = self.mmr.forest & (usize::MAX << 1);

        if self.forest < target {
            if self.last_right == 0 {
                // yield the left leaf
                debug_assert!(self.last_right == 0, "left must be before right");
                self.forest |= 1;
                self.index += 1;

                // yield the right leaf
                debug_assert!((self.forest & 1) == 1, "right must be after left");
                self.last_right |= 1;
                self.index += 1;
            };

            debug_assert!(
                self.forest & self.last_right != 0,
                "parent requires both a left and right",
            );

            // compute the number of nodes in the right tree, this is the offset to the
            // previous left parent
            let right_nodes = nodes_in_forest(self.last_right);
            // the next parent position is one above the position of the pair
            let parent = self.last_right << 1;

            // the left node has been paired and the current parent yielded, removed it from the forest
            self.forest ^= self.last_right;
            if self.forest & parent == 0 {
                // this iteration yielded the left parent node
                debug_assert!(self.forest & 1 == 0, "next iteration yields a left leaf");
                self.last_right = 0;
                self.forest ^= parent;
            } else {
                // the left node of the parent level has been yielded already, this iteration
                // was the right parent. Next iteration yields their parent.
                self.last_right = parent;
            }

            // yields a parent
            let value = self.mmr.nodes[self.index];
            let right = self.mmr.nodes[self.index - 1];
            let left = self.mmr.nodes[self.index - 1 - right_nodes];
            self.index += 1;
            let node = InnerNodeInfo { value, left, right };

            Some(node)
        } else {
            None
        }
    }
}

// UTILITIES
// ===============================================================================================

/// Given a 0-indexed leaf position and the current forest, return the tree number responsible for
/// the position.
///
/// Note:
/// The result is a tree position `p`, it has the following interpretations. $p+1$ is the depth of
/// the tree, which corresponds to the size of a Merkle proof for that tree. $2^p$ is equal to the
/// number of leaves in this particular tree. and $2^(p+1)-1$ corresponds to size of the tree.
pub(crate) const fn leaf_to_corresponding_tree(pos: usize, forest: usize) -> Option<u32> {
    if pos >= forest {
        None
    } else {
        // - each bit in the forest is a unique tree and the bit position its power-of-two size
        // - each tree owns a consecutive range of positions equal to its size from left-to-right
        // - this means the first tree owns from `0` up to the `2^k_0` first positions, where `k_0`
        // is the highest true bit position, the second tree from `2^k_0 + 1` up to `2^k_1` where
        // `k_1` is the second higest bit, so on.
        // - this means the highest bits work as a category marker, and the position is owned by
        // the first tree which doesn't share a high bit with the position
        let before = forest & pos;
        let after = forest ^ before;
        let tree = after.ilog2();

        Some(tree)
    }
}

/// Return a bitmask for the bits including and above the given position.
pub(crate) const fn high_bitmask(bit: u32) -> usize {
    if bit > usize::BITS - 1 {
        0
    } else {
        usize::MAX << bit
    }
}

/// Return the total number of nodes of a given forest
///
/// Panics:
///
/// This will panic if the forest has size greater than `usize::MAX / 2`
pub(crate) const fn nodes_in_forest(forest: usize) -> usize {
    // - the size of a perfect binary tree is $2^{k+1}-1$ or $2*2^k-1$
    // - the forest represents the sum of $2^k$ so a single multiplication is necessary
    // - the number of `-1` is the same as the number of trees, which is the same as the number
    // bits set
    let tree_count = forest.count_ones() as usize;
    forest * 2 - tree_count
}