use sp_std::{collections::btree_map::BTreeMap, vec::Vec, result::Result};
use rand_chacha::rand_core::{CryptoRng, RngCore};

use ghost_traits::exodus::{
    DistributedKeyGeneration, IdentifierConverter, MerkleTreeBuilder,
};

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

impl DistributedKeyGeneration<AuthIndex, ExodusError> for NetworkCurve {
    type RoundPackages = BTreeMap<AuthIndex, Vec<u8>>;

    type Part1Result = Result<(Vec<u8>, Vec<u8>), ExodusError>;
    type Part2Result = Result<(Vec<u8>, BTreeMap<AuthIndex, Vec<u8>>), ExodusError>;
    type Part3Result = Result<
        (
			Vec<u8>, Vec<u8>,
			Vec<<Self as MerkleTreeBuilder<AuthIndex, ExodusError>>::Hash>,
			<Self as MerkleTreeBuilder<AuthIndex, ExodusError>>::Hash,
        ),
        ExodusError
    >;

    fn dkg_narrow_round_max_signers(
        &self,
        secret_package_bytes: &[u8],
        new_max_signers: AuthIndex,
        round_number: u8,
    ) -> Result<Vec<u8>, ExodusError> {
        with_ciphersuite!(self, |f| {
            macro_rules! process_secret_package {
                (
                    $package_type:ty,
                    $field1:ident : $action1:tt,
                    $field2:ident : $action2:tt,
                ) => {{
                    let secret_package = <$package_type>::deserialize(secret_package_bytes)
                        .map_err(|_| ExodusError::DeserializationError)?;

                    let min_signers = *secret_package.min_signers();

                    frost_core::keys::validate_num_of_signers(min_signers, new_max_signers)
                        .map_err(|err| match err {
                            f::Error::InvalidMinSigners => {
                                ExodusError::InvalidMinSigners(min_signers, new_max_signers)
                            },
                            f::Error::InvalidMaxSigners => {
                                ExodusError::InvalidMaxSigners(new_max_signers)
                            },
                            _ => ExodusError::Unknown,
                        })?;

                    macro_rules! apply_action {
                        ($field:ident, clone) => { secret_package.$field().clone() };
                        ($field:ident, value) => { secret_package.$field() };
                    }

                    let new_secret_package = <$package_type>::new(
                        secret_package.identifier().clone(),
                        apply_action!($field1, $action1),
                        apply_action!($field2, $action2),
                        min_signers,
                        new_max_signers,
                    );

                    new_secret_package.serialize()
                        .map_err(|_| ExodusError::SerializationError)
                }};
            }

            match round_number {
                1 => process_secret_package!(
                    f::keys::dkg::round1::SecretPackage,
                    coefficients: value,
                    commitment: clone,
                ),
                2 => process_secret_package!(
                    f::keys::dkg::round2::SecretPackage,
                    commitment: clone,
                    secret_share: value,
                ),
                _ => Err(ExodusError::UnexpectedRound),
            }
        })
    }

    fn dkg_verify_proof_of_knowledge(
        &self,
        index: AuthIndex,
        round1_package_bytes: &[u8],
    ) -> Result<AuthIndex, ExodusError> {
        with_ciphersuite!(self, |f| {
            let identifier = Self::non_zero_index(index)?
                .try_into()
                .map_err(|_| ExodusError::InvalidParticipantId)?;

            let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let commitment = round1_package.commitment();

            frost_core::keys::dkg::verify_proof_of_knowledge(
                identifier,
                &commitment,
                &round1_package.proof_of_knowledge(),
            ).map_err(|_| ExodusError::InvalidProofOfKnowledge)?;

            let coefficients = commitment.coefficients();
            let coefficients_len = coefficients.len() as AuthIndex;

            Ok(coefficients_len)
        })
    }

    fn dkg_part1<R: RngCore + CryptoRng>(
        &self,
        index: AuthIndex,
        max_signers: AuthIndex,
        min_signers: AuthIndex,
        mut rng: R,
    ) -> Self::Part1Result {
        with_ciphersuite!(self, |f| {
            let identifier = Self::non_zero_index(index)?
                .try_into()
                .map_err(|_| ExodusError::InvalidParticipantId)?;

            let (round1_secret_package, round1_package) = f::keys::dkg::part1(
                identifier,
                max_signers,
                min_signers,
                &mut rng,
            ).map_err(|err| {
                match err {
                    f::Error::InvalidMinSigners => ExodusError::InvalidMinSigners(min_signers, max_signers),
                    f::Error::InvalidMaxSigners => ExodusError::InvalidMaxSigners(max_signers),
                    f::Error::DKGNotSupported => ExodusError::DKGNotSupported(*self),
                    _ => ExodusError::Unknown,
                }
            })?;

            let round1_secret_package_bytes = round1_secret_package.serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let round1_package_bytes = round1_package.serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            Ok((round1_secret_package_bytes, round1_package_bytes))
        })
    }

    fn dkg_part2(
        &self,
        round1_secret_package_bytes: &[u8],
        round1_packages_bytes: &Self::RoundPackages,
    ) -> Self::Part2Result {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;
            type InnerField = <InnerGroup as frost_core::Group>::Field;

            let round1_secret_package = f::keys::dkg::round1::SecretPackage::deserialize(round1_secret_package_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let round1_packages = round1_packages_bytes
                .iter()
                .map(|(&index, round1_package_bytes)| {
                    let identifier = Self::non_zero_index(index)?
                        .try_into()
                        .map_err(|_| ExodusError::InvalidParticipantId)?;

                    let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                        .map_err(|_| ExodusError::DeserializationError)?;

                    Ok((identifier, round1_package))
                })
            	.collect::<Result<BTreeMap<f::Identifier, f::keys::dkg::round1::Package>, ExodusError>>()?;

            let (round2_secret_package, round2_packages) = f::keys::dkg::part2(
                round1_secret_package,
                &round1_packages,
            ).map_err(|err| match err {
                f::Error::IncorrectNumberOfCommitments => ExodusError::IncorrectNumberOfCommitments,
                f::Error::IncorrectNumberOfPackages => ExodusError::IncorrectNumberOfPackages,
                f::Error::InvalidProofOfKnowledge { .. }  => ExodusError::InvalidProofOfKnowledge,
                _ => ExodusError::Unknown,
            })?;

            let round2_secret_package_bytes = round2_secret_package.serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let round2_packages_bytes = round2_packages
                .iter()
                .map(|(id, round2_package)| {
                    let id_scalar = id.to_scalar();
                    let id_bytes = <InnerField as frost_core::Field>::serialize(&id_scalar);

                    let index = Self::convert_identifier_to_index(id_bytes.as_ref())?;
                    let round2_package_bytes = round2_package.serialize()
                    	.map_err(|_| ExodusError::SerializationError)?;

                    Ok((index, round2_package_bytes))
                })
                .collect::<Result<BTreeMap<AuthIndex, Vec<u8>>, ExodusError>>()?;

            Ok((round2_secret_package_bytes, round2_packages_bytes))
        })
    }

    fn dkg_verify_private_package(
        &self,
        index: AuthIndex,
        round1_package_bytes: &[u8],
        round2_package_bytes: &[u8],
    ) -> Result<(), ExodusError> {
        with_ciphersuite!(self, |f| {
            let identifier = Self::non_zero_index(index)?
                .try_into()
                .map_err(|_| ExodusError::InvalidParticipantId)?;

            let commitment = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                .map(|round1_package| round1_package.commitment().clone())
                .map_err(|_| ExodusError::DeserializationError)?;

            let secret_share = f::keys::dkg::round2::Package::deserialize(round2_package_bytes)
                .map(|round2_package| {
                    let signing_share = round2_package.signing_share().clone();
                    f::keys::SecretShare::new(identifier, signing_share, commitment)
                })
                .map_err(|_| ExodusError::DeserializationError)?;

            let _ = secret_share.verify().map_err(|_| ExodusError::InvalidSecretShare)?;

            Ok(())
        })
    }

    fn dkg_part3(
        &self,
        authority_index: AuthIndex,
        round2_secret_package_bytes: &[u8],
        round1_packages_bytes: &Self::RoundPackages,
        round2_packages_bytes: &Self::RoundPackages,
    ) -> Self::Part3Result {
        with_ciphersuite!(self, |f| {
            let round2_secret_package = f::keys::dkg::round2::SecretPackage::deserialize(round2_secret_package_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let round1_packages = round1_packages_bytes
                .iter()
                .map(|(&index, round1_package_bytes)| {
                    let identifier = Self::non_zero_index(index)?
                        .try_into()
                        .map_err(|_| ExodusError::InvalidParticipantId)?;

                    let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                        .map_err(|_| ExodusError::DeserializationError)?;

                    Ok((identifier, round1_package))
                })
                .collect::<Result<BTreeMap<f::Identifier, f::keys::dkg::round1::Package>, ExodusError>>()?;

            let round2_packages = round2_packages_bytes
                .iter()
                .map(|(&index, round2_package_bytes)| {
                    let identifier = Self::non_zero_index(index)?
                        .try_into()
                        .map_err(|_| ExodusError::InvalidParticipantId)?;

                    let round2_package = f::keys::dkg::round2::Package::deserialize(round2_package_bytes)
                        .map_err(|_| ExodusError::DeserializationError)?;

                    Ok((identifier, round2_package))
                })
            	.collect::<Result<BTreeMap<f::Identifier, f::keys::dkg::round2::Package>, ExodusError>>()?;

            let (key_package, pubkey_package) = f::keys::dkg::part3(
                &round2_secret_package,
                &round1_packages,
                &round2_packages,
            ).map_err(|err| match err {
                f::Error::IncorrectNumberOfPackages => ExodusError::IncorrectNumberOfPackages,
                f::Error::IncorrectPackage => ExodusError::IncorrectPackage,
                f::Error::PackageNotFound => ExodusError::PackageNotFound,
                f::Error::InvalidSecretShare { .. } => ExodusError::InvalidSecretShare,
                f::Error::IncorrectNumberOfCommitments => ExodusError::IncorrectNumberOfCommitments,
                _ => ExodusError::Unknown,
            })?;

            let key_package = key_package.normalize_package();
            let pubkey_package = pubkey_package.normalize_package();

            let key_package_bytes = key_package.serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let verifying_shares_merkle_tree =
                Self::build_merkle_tree(
                    pubkey_package.verifying_shares(),
                    |verifying_share| verifying_share.serialize()
                        .map_err(|_| ExodusError::SerializationError)
                )?;

            let (_, merkle_proof) =
                Self::generate_merkle_proof_from_tree(
                    pubkey_package.verifying_shares(),
                    &verifying_shares_merkle_tree,
                    authority_index,
                    |verifying_share| verifying_share.serialize()
                    	.map_err(|_| ExodusError::SerializationError)
                )?;

            let merkle_root = verifying_shares_merkle_tree
                .last()
                .copied()
                .ok_or(ExodusError::IncorrectNumberOfPackages)?;

            let verifying_key_bytes = pubkey_package.verifying_key()
                .serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            Ok((
                key_package_bytes,
                verifying_key_bytes,
                merkle_proof,
                merkle_root,
            ))
        })
    }

    fn dkg_derive_verifying_share(
        &self,
        secret_key_package_bytes: &[u8],
    ) -> Result<Vec<u8>, ExodusError> {
        with_ciphersuite!(self, |f| {
            let secret_key_package =
            	f::keys::KeyPackage::deserialize(secret_key_package_bytes)
                    .map_err(|_| ExodusError::DeserializationError)?;

            secret_key_package
                .verifying_share()
                .serialize()
                .map_err(|_| ExodusError::SerializationError)
        })
    }
}