use sp_std::{collections::btree_map::BTreeMap, vec::Vec, result::Result};

use frost_core::{Ciphersuite, Group, Identifier};
use ghost_helpers::SubstrateBlake2Hasher;
use ghost_traits::{
    exodus::{MerkleTreeBuilder, IdentifierConverter},
    hashing::GhostHasher,
};

use crate::{AuthIndex, ExodusError, NetworkCurve};

impl MerkleTreeBuilder<AuthIndex, ExodusError> for NetworkCurve {
    type Hash = <SubstrateBlake2Hasher as GhostHasher>::Hash;

    fn build_merkle_tree<C, V, F>(
        size_hint: usize,
        values: &BTreeMap<Identifier<C>, V>,
        update_preimage_buffer: F,
    ) -> Result<Vec<Self::Hash>, ExodusError>
	where
		C: frost_core::Ciphersuite,
		<<C as Ciphersuite>::Group as Group>::Field: frost_core::Field,
		F: Fn(&V, &mut Vec<u8>) -> Result<(), ExodusError>
	{
        use ghost_helpers::merkle_tree::{generate_tree, MerkleTreeError};

        type CField<C> = <<C as Ciphersuite>::Group as Group>::Field;

        let max_index = values.keys()
            .next_back()
            .and_then(|id| {
                let id_scalar = id.to_scalar();
                let id_bytes = <CField<C> as frost_core::Field>::serialize(&id_scalar);
                Self::convert_identifier_to_index(id_bytes.as_ref()).ok()
            })
            .ok_or(ExodusError::IncorrectNumberOfIdentifiers)?;

        let maybe_tree = generate_tree::<SubstrateBlake2Hasher, _, _, _, _>(
            max_index,
            size_hint,
            values.iter(),
            |(identifier, value), buffer| {
                let id_scalar = identifier.to_scalar();
                let id_bytes = <CField<C> as frost_core::Field>::serialize(&id_scalar);

                let index = Self::convert_identifier_to_index(id_bytes.as_ref())?;

                buffer.extend_from_slice(id_bytes.as_ref());
                update_preimage_buffer(value, buffer)?;

                Ok(index as usize)
            });

        match maybe_tree {
            Ok(tree) => Ok(tree),
            Err(MerkleTreeError::Generate(exodus_err)) => Err(exodus_err),
            Err(MerkleTreeError::InvalidIndex) => Err(ExodusError::TreeGenerationFailed),
            Err(MerkleTreeError::Empty) => Err(ExodusError::IncorrectNumberOfIdentifiers),
        }
    }

    fn generate_merkle_proof_from_tree<C, V, F>(
        values: &BTreeMap<Identifier<C>, V>,
        merkle_tree: &[Self::Hash],
        authority_index: AuthIndex,
        serialize_fn: F,
    ) -> Result<(Vec<u8>, Vec<Self::Hash>), ExodusError>
	where
		C: frost_core::Ciphersuite,
		<<C as Ciphersuite>::Group as Group>::Field: frost_core::Field,
		F: Fn(&V) -> Result<Vec<u8>, ExodusError>
	{
        use ghost_helpers::merkle_tree::{generate_proof, MerkleTreeError};

        type CField<C> = <<C as Ciphersuite>::Group as Group>::Field;

        let non_zero_index = Self::non_zero_index(authority_index)?;
        let identifier = Identifier::try_from(non_zero_index)
            .map_err(|_| ExodusError::InvalidParticipantId)?;

        let value = values.get(&identifier).ok_or(ExodusError::InvalidParticipantId)?;
        let value_bytes = serialize_fn(&value)?;

        let max_index = values.keys()
            .next_back()
            .and_then(|id| {
                let id_scalar = id.to_scalar();
                let id_bytes = <CField<C> as frost_core::Field>::serialize(&id_scalar);
                Self::convert_identifier_to_index(id_bytes.as_ref()).ok()
            })
            .ok_or(ExodusError::IncorrectNumberOfIdentifiers)?;

        let proof = generate_proof::<SubstrateBlake2Hasher, _, _>(
            merkle_tree,
            max_index,
            authority_index,
        ).map_err(|err| match err {
            MerkleTreeError::Generate(exodus_err) => exodus_err,
            MerkleTreeError::InvalidIndex => ExodusError::TreeGenerationFailed,
            MerkleTreeError::Empty => ExodusError::IncorrectNumberOfIdentifiers,
        })?;

        Ok((value_bytes, proof))
    }

    fn verify_merkle_proof(
        &self,
        values: &[u8],
        merkle_proof: &[Self::Hash],
        merkle_root: Self::Hash,
        authority_index: AuthIndex,
    ) -> Result<(), ExodusError> {
        use ghost_helpers::merkle_tree::verify_tree_proof;

        let preimage = with_ciphersuite!(self, |f| {
            let non_zero_index = Self::non_zero_index(authority_index)?;
            let identifier = f::Identifier::try_from(non_zero_index)
                .map_err(|_| ExodusError::InvalidParticipantId)?;

            let mut preimage = identifier.serialize();
            preimage.extend_from_slice(values);

            Ok(preimage)
        })?;

        let is_verified = verify_tree_proof::<SubstrateBlake2Hasher, _>(
            &preimage,
            merkle_proof,
            merkle_root,
            authority_index,
        );

        is_verified.then(|| ()).ok_or(ExodusError::InvalidMerkleProof)
    }
}