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>(
        values: &BTreeMap<Identifier<C>, V>,
        serialize_fn: F,
    ) -> Result<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_tree;

        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)?;

        generate_tree::<SubstrateBlake2Hasher, _, _, _, _>(
            max_index,
            values.iter(),
            |(identifier, value)| {
                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())?;
                let value_bytes = serialize_fn(value).map_err(|_| ExodusError::SerializationError)?;

                // NOTE: revisit
                // Is it possible to use [u8; CONSTANT] where constant is
                // purely dependant on provided generic.
                let mut preimage = Vec::<u8>::with_capacity(id_bytes.as_ref().len() + value_bytes.len());
                preimage.extend_from_slice(id_bytes.as_ref());
                preimage.extend_from_slice(value_bytes.as_ref());

                Ok((index as usize, preimage))
            })
    }

    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;

        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).map_err(|_| ExodusError::SerializationError)?;

        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,
        );

        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)
    }
}