use crate::node::{null_node, Node, NodeRef};
use crate::tally::TallyList;
use crate::utilstring::to_hex;
use crate::VoteReference;
use std::collections::VecDeque;
use std::sync::Arc;

/// Build a 'level' in the merkle tally tree with given nodes. Returns the new parent nodes.
fn build_level(mut queue: VecDeque<Node>) -> Vec<Node> {
    assert!(queue.len() % 2 == 0);
    let mut nodes: Vec<Node> = vec![];
    loop {
        let left = queue.pop_front();
        if left.is_none() {
            return nodes;
        }
        let right = queue.pop_front();
        nodes.push(Node {
            vote: None,
            left: Some(Arc::new(left.unwrap())),
            right: Some(Arc::new(right.unwrap())),
        });
    }
}

/// Generate a merkle tally tree from votes.
///
/// Example:
/// ```
///# use tallytree::generate::generate_tree;
/// // Generate a tree where vote reference where:
/// // - 0xaa votes for option 0,
/// // - 0xbb votes for option 1,
/// // - 0xcc votes for option 0
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
///     ([0xcc; 32], vec![1, 0]),
/// ]);
/// ```
///
/// The above example results in a tree like this:
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       aa  bb cc
/// ```
pub fn generate_tree(mut votes: Vec<(VoteReference, TallyList)>) -> Result<NodeRef, String> {
    if votes.is_empty() {
        return Ok(None);
    }

    votes.sort_by(|a, b| a.partial_cmp(b).unwrap());

    // Check for duplicates
    let mut last: Option<&VoteReference> = None;
    for (v, _) in &votes {
        if last == Some(v) {
            return Err(format!(
                "Cannot have duplicate vote references (found duplciate of '{}')",
                to_hex(v)?
            ));
        }
        last = Some(v);
    }

    let leafs = votes.into_iter().map(|v| Node {
        vote: Some(v),
        left: None,
        right: None,
    });

    // Wrap the leaf to secure against the issue bitcoin was vulnerable to.
    let nodes: Vec<Node> = leafs
        .map(|l| Node {
            vote: None,
            left: Some(Arc::new(l)),
            right: None,
        })
        .collect();

    let mut queue = VecDeque::from(nodes);

    // Special case: If there is only 1 vote leaf, we need to pair
    // it with a Ø-node.
    //
    // Without this, the node would be the merkle root and we would
    // not be able to generate "proof of number of participants".
    if queue.len() == 1 {
        queue.push_back(null_node());
    }

    loop {
        if queue.len() == 1 {
            return Ok(Some(Arc::new(queue.pop_front().unwrap())));
        }
        if queue.len() % 2 != 0 {
            // Push a NULL node to balance level.
            queue.push_back(null_node());
        }
        queue = VecDeque::from(build_level(queue));
    }
}

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

    #[test]
    fn test_generate_tree_one_vote() {
        //  Should look like this:
        //
        // ```text
        //          N
        //        /   \
        //       N     Ø
        //       |
        //       V
        // ```
        //
        //   Where V is the vote, N is a node and Ø is a null node.
        //
        let voter = [0xaa; 32];
        let vote = vec![1, 0];
        let root = generate_tree(vec![(voter, vote.clone())]).unwrap();

        let root = root.unwrap();
        assert!(is_null_node_ref(&root.right));
        let (found_voter, found_vote) = root
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();

        assert_eq!(&voter, found_voter);
        assert_eq!(&vote, found_vote);
    }

    #[test]
    fn test_generate_tree_is_sorted() {
        //  Should look like this:
        //
        // ```text
        //          N
        //        /   \
        //       N     N
        //      / \   / \
        //     N  N  N   Ø
        //     |  |  |
        //     a  b  c
        // ```

        let root = generate_tree(vec![
            ([0xcc; 32], vec![1, 0]),
            ([0xaa; 32], vec![1, 0]),
            ([0xbb; 32], vec![1, 0]),
        ])
        .unwrap()
        .unwrap();
        let (found_voter, _) = root
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();
        assert_eq!(&[0xaa; 32], found_voter);

        let (found_voter, _) = root
            .right
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();
        assert_eq!(&[0xcc; 32], found_voter);
    }

    #[test]
    fn test_generate_tree_rejects_dupes() {
        let error = generate_tree(vec![
            ([0xaa; 32], vec![1, 0]),
            ([0xbb; 32], vec![1, 0]),
            ([0xaa; 32], vec![0, 1]), // duplicate
        ]);

        let error_msg = "Cannot have duplicate vote references (found duplciate of 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa')";
        assert_eq!(error.err().unwrap(), error_msg);
    }
}
