use crate::node::{is_null_node_ref, Node, NodeRef};
use crate::proof::ProofHash;
use crate::tally::{tally_node, TallyList};
use crate::utilstring::to_hex;
use crate::{Validation, VoteReference};
use sha2::{Digest, Sha256};
use std::fmt;

/// The hash of a node in the merkle tally tree.
#[derive(PartialEq, Clone, Eq, Hash, Copy)]
pub struct NodeHash(pub [u8; 32]);
/// The hash of a Ø-node.
pub const NULL_HASH: NodeHash = NodeHash([0x0; 32]);
/// This "hash value" can be used to indicate an error in user facing messages.
pub const ERROR_HASH: NodeHash = NodeHash([0xff; 32]);

impl NodeHash {
    /// Get hash as a slice
    pub fn as_slice(&self) -> &[u8] {
        &self.0
    }

    /// Get hash as array
    pub fn as_array(&self) -> &[u8; 32] {
        &self.0
    }
}

impl From<ProofHash> for NodeHash {
    fn from(h: ProofHash) -> Self {
        NodeHash(h)
    }
}

impl fmt::Display for NodeHash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "NodeHash {{ {} }}", to_hex(self.as_slice()).unwrap())
    }
}

impl fmt::Debug for NodeHash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "NodeHash {{ {} }}", to_hex(self.as_slice()).unwrap())
    }
}

/**
 * Get the hash of a leaf node.
 *
 * The hash is made from its vote, so we don't need to pass the full Node struct.
 */
pub fn hash_leaf(vote_reference: &VoteReference, vote: &[u32]) -> NodeHash {
    let mut hasher = Sha256::new();
    hasher.update(vote_reference);
    vote.iter()
        .map(|t| t.to_be_bytes())
        .for_each(|b| hasher.update(b));
    NodeHash(
        hasher
            .finalize()
            .as_slice()
            .try_into()
            .expect("Expected hash to be 32 bytes"),
    )
}

/// Get the hash of a node in a merkle tally tree.
///
/// The hash value is generated from this nodes tally and if applicable the
/// hash of the nodes child nodes or vote reference.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::hash::hash_node;
/// use tallytree::Validation;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
///     ([0xcc; 32], vec![1, 0]),
/// ]).unwrap();
/// let hash = hash_node(&tree.unwrap(), &Validation::Strict);
/// ```
pub fn hash_node(node: &Node, v: &Validation) -> Result<(NodeHash, Option<TallyList>), String> {
    let mut hasher = Sha256::new();
    if let Some((vote_reference, _)) = node.vote {
        assert!(node.left.is_none());
        assert!(node.right.is_none());
        hasher.update(vote_reference);
    }
    let tally = tally_node(node, v)?.ok_or_else(|| "Expected node to provide tally".to_string())?;
    tally
        .iter()
        .map(|t| t.to_be_bytes())
        .for_each(|b| hasher.update(b));
    if let Some((h, _)) = hash_node_ref(&node.left, v)? {
        hasher.update(h.as_slice());
    }
    if let Some((h, _)) = hash_node_ref(&node.right, v)? {
        assert!(node.left.is_some());
        hasher.update(h.as_slice());
    }
    Ok((
        NodeHash(
            hasher
                .finalize()
                .as_slice()
                .try_into()
                .expect("Expected hash to be 32 bytes"),
        ),
        Some(tally),
    ))
}

/// Get the hash of a node in a merkle tally tree.
///
/// The hash value is generated from this nodes tally and if applicable the
/// hash of the nodes child nodes or vote reference.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::hash::hash_node_ref;
/// use tallytree::Validation;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
///     ([0xcc; 32], vec![1, 0]),
/// ]).unwrap();
/// let hash = hash_node_ref(&tree, &Validation::Strict);
/// let hash_child = hash_node_ref(&tree.unwrap().left, &Validation::Strict);
/// ```
pub fn hash_node_ref(
    node: &NodeRef,
    v: &Validation,
) -> Result<Option<(NodeHash, Option<TallyList>)>, String> {
    if is_null_node_ref(node) {
        return Ok(Some((NULL_HASH, None)));
    }
    match node {
        Some(n) => Ok(Some(hash_node(n, v)?)),
        None => Ok(None),
    }
}
