use alloc::{collections::BTreeMap, vec::Vec};

use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
use crate::{
    hash::rpo::{Rpo256, RpoDigest},
    Felt, Word, EMPTY_WORD,
};

mod full;
pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};

mod simple;
pub use simple::SimpleSmt;

// CONSTANTS
// ================================================================================================

/// Minimum supported depth.
pub const SMT_MIN_DEPTH: u8 = 1;

/// Maximum supported depth.
pub const SMT_MAX_DEPTH: u8 = 64;

// SPARSE MERKLE TREE
// ================================================================================================

/// An abstract description of a sparse Merkle tree.
///
/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
/// value was not explicitly set, then its value is the default value. Typically, the vast majority
/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
/// representation of the tree will only keep track of the leaves that have a different value from
/// the default.
///
/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
/// tree is just a single value, and is probably a programming mistake.
///
/// Every key maps to one leaf. If there are as many keys as there are leaves, then
/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
/// must accomodate all keys that map to the same leaf.
///
/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
pub trait SparseMerkleTree<const DEPTH: u8> {
    /// The type for a key
    type Key: Clone + Ord;
    /// The type for a value
    type Value: Clone + PartialEq;
    /// The type for a leaf
    type Leaf: Clone;
    /// The type for an opening (i.e. a "proof") of a leaf
    type Opening;

    /// The default value used to compute the hash of empty leaves
    const EMPTY_VALUE: Self::Value;

    // PROVIDED METHODS
    // ---------------------------------------------------------------------------------------------

    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
    /// path to the leaf, as well as the leaf itself.
    fn open(&self, key: &Self::Key) -> Self::Opening {
        let leaf = self.get_leaf(key);

        let mut index: NodeIndex = {
            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
            leaf_index.into()
        };

        let merkle_path = {
            let mut path = Vec::with_capacity(index.depth() as usize);
            for _ in 0..index.depth() {
                let is_right = index.is_value_odd();
                index.move_up();
                let InnerNode { left, right } = self.get_inner_node(index);
                let value = if is_right { left } else { right };
                path.push(value);
            }

            MerklePath::new(path)
        };

        Self::path_and_leaf_to_opening(merkle_path, leaf)
    }

    /// Inserts a value at the specified key, returning the previous value associated with that key.
    /// Recall that by definition, any key that hasn't been updated is associated with
    /// [`Self::EMPTY_VALUE`].
    ///
    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
    /// updating the root itself.
    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
        let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);

        // if the old value and new value are the same, there is nothing to update
        if value == old_value {
            return value;
        }

        let leaf = self.get_leaf(&key);
        let node_index = {
            let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
            leaf_index.into()
        };

        self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));

        old_value
    }

    /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
    /// `node_hash_at_index` is the hash of the node stored at index.
    fn recompute_nodes_from_index_to_root(
        &mut self,
        mut index: NodeIndex,
        node_hash_at_index: RpoDigest,
    ) {
        let mut node_hash = node_hash_at_index;
        for node_depth in (0..index.depth()).rev() {
            let is_right = index.is_value_odd();
            index.move_up();
            let InnerNode { left, right } = self.get_inner_node(index);
            let (left, right) = if is_right {
                (left, node_hash)
            } else {
                (node_hash, right)
            };
            node_hash = Rpo256::merge(&[left, right]);

            if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
                // If a subtree is empty, when can remove the inner node, since it's equal to the
                // default value
                self.remove_inner_node(index)
            } else {
                self.insert_inner_node(index, InnerNode { left, right });
            }
        }
        self.set_root(node_hash);
    }

    /// Start a prospective mutation transaction, which can be queried, discarded, or applied.
    ///
    /// This method returns a [`MutationSet`], which you can call [`MutationSet::insert()`] on to
    /// perform prospective mutations to this Merkle tree, and check the computed root hash given
    /// those mutations with [`MutationSet::root()`]. When you are done, call
    /// [`SparseMerkleTree::commit()`] with the return value of [`MutationSet::done()`] to apply the
    /// prospective mutations to this tree. Or to discard the changes, simply [`drop`] the
    /// [`MutationSet`].
    ///
    /// ## Example
    /// ```
    /// # use miden_crypto::{hash::rpo::RpoDigest, Felt};
    /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
    /// # let mut smt = Smt::default();
    /// let mut mutations = smt.mutate();
    /// mutations.insert(Default::default(), Default::default());
    /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
    /// smt.commit(mutations.done());
    /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
    /// ```
    fn mutate(&self) -> MutationSet<DEPTH, Self>
    where
        Self: Sized,
    {
        MutationSet {
            tree: self,
            node_mutations: Default::default(),
            new_pairs: Default::default(),
            new_root: self.root(),
        }
    }

    /// Apply prospective mutations started with [`SparseMerkleTree::mutate()`] to this tree.
    ///
    /// This method takes the return value of [`MutationSet::done()`], and applies the changes
    /// represented in that [`MutationSet`] to this Merkle tree. See [`SparseMerkleTree::mutate()`]
    /// for more details.
    fn commit(&mut self, mutations: DoneMutationSet<DEPTH, Self::Key, Self::Value>)
    where
        Self: Sized,
    {
        use NodeMutation::*;
        let DoneMutationSet { node_mutations, new_pairs, new_root } = mutations;

        for (index, mutation) in node_mutations {
            match mutation {
                Removal => self.remove_inner_node(index),
                Addition(node) => self.insert_inner_node(index, node),
            }
        }

        for (key, value) in new_pairs {
            self.insert_value(key, value);
        }

        self.set_root(new_root);
    }

    // REQUIRED METHODS
    // ---------------------------------------------------------------------------------------------

    /// The root of the tree
    fn root(&self) -> RpoDigest;

    /// Sets the root of the tree
    fn set_root(&mut self, root: RpoDigest);

    /// Retrieves an inner node at the given index
    fn get_inner_node(&self, index: NodeIndex) -> InnerNode;

    /// Inserts an inner node at the given index
    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode);

    /// Removes an inner node at the given index
    fn remove_inner_node(&mut self, index: NodeIndex);

    /// Inserts a leaf node, and returns the value at the key if already exists
    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;

    /// Returns the value at the specified key. Recall that by definition, any key that hasn't been
    /// updated is associated with [`Self::EMPTY_VALUE`].
    fn get_value(&self, key: &Self::Key) -> Self::Value;

    /// Returns the leaf at the specified index.
    fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;

    /// Returns the hash of a leaf
    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;

    /// Returns what `existing_leaf` would look like if `key` and `value` WERE inserted into the
    /// tree, without mutating the tree itself.
    ///
    /// `existing_leaf` must have the same index as the key, or the result will be meaningless. To
    /// get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)` as
    /// the argument for `existing_leaf`. The return value from this function can be chained back
    /// into this function as the first argument to continue making prospective changes.
    fn get_prospective_leaf(
        &self,
        existing_leaf: Self::Leaf,
        key: &Self::Key,
        value: &Self::Value,
    ) -> Self::Leaf;

    /// Maps a key to a leaf index
    fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;

    /// Maps a (MerklePath, Self::Leaf) to an opening.
    ///
    /// The length `path` is guaranteed to be equal to `DEPTH`
    fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
}

// INNER NODE
// ================================================================================================

#[derive(Debug, Default, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct InnerNode {
    pub left: RpoDigest,
    pub right: RpoDigest,
}

impl InnerNode {
    pub fn hash(&self) -> RpoDigest {
        Rpo256::merge(&[self.left, self.right])
    }
}

// LEAF INDEX
// ================================================================================================

/// The index of a leaf, at a depth known at compile-time.
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct LeafIndex<const DEPTH: u8> {
    index: NodeIndex,
}

impl<const DEPTH: u8> LeafIndex<DEPTH> {
    pub fn new(value: u64) -> Result<Self, MerkleError> {
        if DEPTH < SMT_MIN_DEPTH {
            return Err(MerkleError::DepthTooSmall(DEPTH));
        }

        Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
    }

    pub fn value(&self) -> u64 {
        self.index.value()
    }
}

impl LeafIndex<SMT_MAX_DEPTH> {
    pub const fn new_max_depth(value: u64) -> Self {
        LeafIndex {
            index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
        }
    }
}

impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
    fn from(value: LeafIndex<DEPTH>) -> Self {
        value.index
    }
}

impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
    type Error = MerkleError;

    fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
        if node_index.depth() != DEPTH {
            return Err(MerkleError::InvalidDepth {
                expected: DEPTH,
                provided: node_index.depth(),
            });
        }

        Self::new(node_index.value())
    }
}

// MUTATIONS
// ================================================================================================

/// A change to an inner node of a [`SparseMerkleTree`] that hasn't been applied yet.
/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
/// need to occur at what node indices.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum NodeMutation {
    Removal,
    Addition(InnerNode),
}

/// Represents prospective mutations to a [`SparseMerkleTree`], created by
/// [`SparseMerkleTree::mutate()`], and applied with [`SparseMerkleTree::commit()`].
///
/// `T` is the Merkle tree type this is prospectively mutating.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MutationSet<'t, const DEPTH: u8, T>
where
    T: SparseMerkleTree<DEPTH>,
{
    /// A reference to the tree we're mutating. We need to be able to query the tree's existing
    /// nodes and leaves to calculate prospective nodes, leaves, and roots.
    tree: &'t T,
    /// The set of nodes that need to be removed or added. The "effective" node at an index is the
    /// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that
    /// index overlayed, if any.
    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
    /// The set of top-level key-value pairs we're prospectively adding to the tree, including
    /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling
    /// back to the existing value in the Merkle tree.
    new_pairs: BTreeMap<T::Key, T::Value>,
    /// The currently-calculated root for the Merkle tree, given these mutations. When this struct
    /// is constructed, `new_root` is set to the existing root, since there aren't any changes yet.
    new_root: RpoDigest,
}

impl<'t, const DEPTH: u8, T> MutationSet<'t, DEPTH, T>
where
    T: SparseMerkleTree<DEPTH>,
{
    /// Pass the return value of this method to [`SparseMerkleTree::commit()`] to apply these
    /// mutations to the Merkle tree.
    pub fn done(self) -> DoneMutationSet<DEPTH, T::Key, T::Value> {
        // Get rid of the `tree` reference, so we can pass the mutation info to a `&mut self` method
        // on `SparseMerkleTree`.
        let Self { node_mutations, new_pairs, new_root, .. } = self;
        DoneMutationSet { node_mutations, new_pairs, new_root }
    }

    /// Get the prospective root hash for if these mutations were applied to the Merkle tree.
    pub fn root(&self) -> RpoDigest {
        self.new_root
    }

    /// Add a new key-value pair to the set of prospective changes to the Merkle tree.
    ///
    /// After `insert()`, you may call [`MutationSet::root()`] to query what the new calculated root
    /// hash of the Merkle tree would be after all requested insertions. Recall that semantically
    /// removing a value is done by inserting [`SparseMerkleTree::EMPTY_VALUE`].
    pub fn insert(&mut self, key: T::Key, value: T::Value) {
        // This function's calculations are eager, so multiple inserts that affect the same nodes
        // will calculate those node's hashes each time. Future work could lazily calculate node
        // hashes all at once instead.

        // If the old value and the new value are the same, there is nothing to update.
        let old_value = self.get_effective_value(&key);
        if value == old_value {
            return;
        }

        let mut node_index: NodeIndex = T::key_to_leaf_index(&key).into();

        let old_leaf = self.get_effective_leaf(&key);
        let new_leaf = self.tree.get_prospective_leaf(old_leaf, &key, &value);

        let mut new_child_hash = T::hash_leaf(&new_leaf);

        for node_depth in (0..node_index.depth()).rev() {
            // Whether the node we're replacing is the right child or left child.
            let is_right = node_index.is_value_odd();
            node_index.move_up();

            let old_node = self.get_effective_node(node_index);

            let new_node = if is_right {
                InnerNode {
                    left: old_node.left,
                    right: new_child_hash,
                }
            } else {
                InnerNode {
                    left: new_child_hash,
                    right: old_node.right,
                }
            };

            // The next iteration will operate on this new node's hash.
            new_child_hash = new_node.hash();

            let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
            if new_child_hash == equivalent_empty_hash {
                self.remove_inner_node(node_index);
            } else {
                self.insert_inner_node(node_index, new_node);
            }
        }

        // Once we're at depth 0, the last node we made is the new root.
        self.new_root = new_child_hash;
        self.new_pairs.insert(key, value);
    }

    // Implementation details.

    fn get_effective_value(&self, key: &T::Key) -> T::Value {
        self.new_pairs.get(key).cloned().unwrap_or_else(|| self.tree.get_value(key))
    }

    fn get_effective_leaf(&self, key: &T::Key) -> T::Leaf {
        let leaf_key = key;
        let leaf_index = T::key_to_leaf_index(leaf_key);
        let pairs_at_index = self
            .new_pairs
            .iter()
            .filter(|&(new_key, _)| T::key_to_leaf_index(new_key) == leaf_index);

        let initial_acc = self.tree.get_leaf(key);
        pairs_at_index.fold(initial_acc, |acc, (k, v)| {
            // In practice this should only run once (if that), most of the time.
            let existing_leaf = acc.clone();
            self.tree.get_prospective_leaf(existing_leaf, k, v)
        })
    }

    fn get_effective_node(&self, index: NodeIndex) -> InnerNode {
        use NodeMutation::*;

        self.node_mutations
            .get(&index)
            .map(|mutation| match mutation {
                Addition(node) => node.clone(),
                Removal => {
                    let &child = EmptySubtreeRoots::entry(DEPTH, index.depth() + 1);
                    InnerNode { left: child, right: child }
                },
            })
            .unwrap_or_else(|| self.tree.get_inner_node(index))
    }

    fn remove_inner_node(&mut self, index: NodeIndex) {
        use alloc::collections::btree_map::Entry::*;
        use NodeMutation::*;

        match self.node_mutations.entry(index) {
            Vacant(entry) => {
                entry.insert(Removal);
            },
            Occupied(mut entry) => match entry.get_mut() {
                // If we have an addition with this index, we don't care about what value it has, we
                // just need to convert it to a removal instead.
                Addition(_) => {
                    entry.insert(Removal);
                },
                Removal => (),
            },
        }
    }

    fn insert_inner_node(&mut self, index: NodeIndex, new_node: InnerNode) {
        use alloc::collections::btree_map::Entry::*;
        use NodeMutation::*;

        match self.node_mutations.entry(index) {
            Vacant(entry) => {
                entry.insert(Addition(new_node));
            },
            Occupied(mut entry) => match entry.get_mut() {
                Addition(existing) => {
                    // If there's an existing addition with this key, then overwrite it to be an
                    // addition of this new node instead.
                    *existing = new_node;
                },
                Removal => {
                    // Likewise a removal of this key gets overwritten with an addition.
                    entry.insert(Addition(new_node));
                },
            },
        }
    }
}

/// Helper type that represents a [`MutationSet`] that's ready to be applied to a [`SparseMerkleTree`].
///
/// It is created by [`MutationSet::done()`], which should be directly passed to
/// [`SparseMerkleTree::commit()`]:
/// ```
/// # use miden_crypto::{hash::rpo::RpoDigest, merkle::Smt, Felt};
/// # let mut smt = Smt::default();
/// # let key = RpoDigest::default();
/// # let value: [Felt; 4] = Default::default();
/// let mut mutations = smt.mutate();
/// mutations.insert(key, value);
/// smt.commit(mutations.done());
/// ```
// This type exists for the sake of the borrow checker -- SparseMerkleTree::commit() needs a
// mutable reference to `self`, but MutationSet stores a shared reference to the SparseMerkleTree,
// and those can't both exist at the same time. By going through this type first, which stores all
// the same data as MutationSet except for the shared reference to SparseMerkleTree, the shared
// reference is dropped just before SparseMerkleTree::commit() is called, and the borrow checker is
// happy. Interior mutability would also work, but then we would lose static lifetime checking.
pub struct DoneMutationSet<const DEPTH: u8, K, V> {
    node_mutations: BTreeMap<NodeIndex, NodeMutation>,
    new_pairs: BTreeMap<K, V>,
    new_root: RpoDigest,
}