use ghost_traits::exodus::{
    IdentifierConverter, MerkleTreeBuilder,
    FlexibleRoundOptimizedSchnorrThresholdSignature,
};
use sp_std::{
    vec::Vec, result::Result,
    collections::{btree_set::BTreeSet, btree_map::BTreeMap},
};

use rand_chacha::rand_core::{CryptoRng, RngCore};
use ghost_helpers::networks::NetworkCurve;
use crate::{AuthIndex, ExodusHash, ExodusError};

impl FlexibleRoundOptimizedSchnorrThresholdSignature<AuthIndex, ExodusError> for NetworkCurve {
    type Packages<'a> = BTreeMap<AuthIndex, &'a [u8]>;
    type NoncePackages<'a> = BTreeMap<AuthIndex, (&'a [u8], &'a [u8])>;

    fn header_bytes_len(&self) -> usize { 5 }

    fn signature_bytes_len(&self) -> usize {
        match self {
            NetworkCurve::Secp256k1 => 64,
            NetworkCurve::Ed25519 => 64,
        }
    }

    fn element_bytes_len(&self) -> usize {
        match self {
            NetworkCurve::Secp256k1 => 33,
            NetworkCurve::Ed25519 => 32,
        }
    }

    fn scalar_bytes_len(&self) -> usize {
        match self {
            NetworkCurve::Secp256k1 => 32,
            NetworkCurve::Ed25519 => 32,
        }
    }

    fn verify_signature_share(
        &self,
        index: AuthIndex,
        participants_iter: impl Iterator<Item = AuthIndex>,
        signature_share_bytes: &[u8],
        binding_factor_bytes: &[u8],
        nonce_hiding_bytes: &[u8],
        nonce_binding_bytes: &[u8],
        group_commitment_bytes: &[u8],
        verifying_share_bytes: &[u8],
        verifying_key_bytes: &[u8],
        message: &[u8],
    ) -> Result<(), ExodusError> {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;
            type InnerField = <InnerGroup as frost_core::Group>::Field;

            let identifier = Self::non_zero_index(index)?
                .try_into()
                .map_err(|_| ExodusError::InvalidParticipantId)?;

            let identifier_set = participants_iter
                .filter_map(|auth_index| {
                    let non_zero_index = Self::non_zero_index(auth_index).ok()?;
                    non_zero_index.try_into().ok()
                })
                .collect::<BTreeSet<f::Identifier>>();

            let lambda_i =
                frost_core::compute_lagrange_coefficient(&identifier_set, None, identifier)
                    .map_err(|err| match err {
                        f::Error::IncorrectNumberOfIdentifiers => ExodusError::IncorrectNumberOfIdentifiers,
                        f::Error::UnknownIdentifier => ExodusError::UnknownIdentifier,
                        f::Error::DuplicatedIdentifier => ExodusError::DuplicatedIdentifier,
                        _ => ExodusError::Unknown,
                    })?;

            let signature_share =
                f::round2::SignatureShare::deserialize(signature_share_bytes)
                    .map_err(|_| ExodusError::DeserializationError)?;

            let group_commitment_serialization =
            	<InnerGroup as frost_core::Group>::Serialization::try_from(group_commitment_bytes.as_ref())
            		.map_err(|_| ExodusError::DeserializationError)?;

            let mut group_commitment_element =
            	<InnerGroup as frost_core::Group>::deserialize(&group_commitment_serialization)
                    .map_err(|_| ExodusError::DeserializationError)?;

            let binding_factor_serialization =
            	<InnerField as frost_core::Field>::Serialization::try_from(binding_factor_bytes.as_ref())
            		.map_err(|_| ExodusError::DeserializationError)?;

            let binding_factor_scalar =
            	<InnerField as frost_core::Field>::deserialize(&binding_factor_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let verifying_share =
                f::keys::VerifyingShare::deserialize(verifying_share_bytes)
                    .map_err(|_| ExodusError::DeserializationError)?;

            let verifying_key = f::VerifyingKey::deserialize(verifying_key_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let hiding_element_serialization =
            	<InnerGroup as frost_core::Group>::Serialization::try_from(nonce_hiding_bytes.as_ref())
            		.map_err(|_| ExodusError::DeserializationError)?;

            let hiding_element =
            	<InnerGroup as frost_core::Group>::deserialize(&hiding_element_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let binding_element_serialization =
            	<InnerGroup as frost_core::Group>::Serialization::try_from(nonce_binding_bytes.as_ref())
            		.map_err(|_| ExodusError::DeserializationError)?;

            let binding_element =
            	<InnerGroup as frost_core::Group>::deserialize(&binding_element_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let mut r_i_element = hiding_element + (binding_element * binding_factor_scalar);

            let group_commitment_wrapper =
                frost_core::GroupCommitment::<Ciphersuite>::from_element(group_commitment_element);

            if !group_commitment_wrapper.is_normalized() {
                r_i_element = -r_i_element;
                group_commitment_element = -group_commitment_element;
            }

            let challenge = <Ciphersuite as frost_core::Ciphersuite>::challenge(
                &group_commitment_element,
                &verifying_key,
                message,
            ).map_err(|err| match err {
                f::Error::SerializationError => ExodusError::SerializationError,
                _ => ExodusError::Unknown,
            })?;

            let commitment_share =
                frost_core::round1::GroupCommitmentShare::from_element(r_i_element);

            signature_share.verify(
                identifier,
                &commitment_share,
                &verifying_share,
                lambda_i,
                &challenge,
            ).map_err(|err| match err {
                f::Error::InvalidSignatureShare { .. } => {
                    ExodusError::InvalidSignatureShare
                },
                _ => ExodusError::Unknown,
            })
        })
    }

    fn accumulate_signature_scalar(
        &self,
        signature_scalar_bytes: &[u8],
        signature_share_bytes: &[u8],
    ) -> Result<Vec<u8>, ExodusError> {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;
            type InnerField = <InnerGroup as frost_core::Group>::Field;

            let signature_scalar_serialization =
            	<InnerField as frost_core::Field>::Serialization::try_from(signature_scalar_bytes)
            		.map_err(|_| ExodusError::DeserializationError)?;

            let mut signature_scalar =
            	<InnerField as frost_core::Field>::deserialize(&signature_scalar_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let signature_share_serialization =
            	<InnerField as frost_core::Field>::Serialization::try_from(signature_share_bytes)
            		.map_err(|_| ExodusError::DeserializationError)?;

            let signature_share =
            	<InnerField as frost_core::Field>::deserialize(&signature_share_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            signature_scalar += signature_share;

            Ok(<InnerField as frost_core::Field>::serialize(&signature_scalar).to_vec())
        })
    }

    fn is_signature_valid(
        &self,
        verifying_key_bytes: &[u8],
        r_element_bytes: &[u8],
        z_scalar_bytes: &[u8],
        message: &[u8],
    ) -> Result<(), ExodusError> {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;
            type InnerField = <InnerGroup as frost_core::Group>::Field;

            let verifying_key = f::VerifyingKey::deserialize(verifying_key_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let r_element_serialization =
            	<InnerGroup as frost_core::Group>::Serialization::try_from(r_element_bytes)
            		.map_err(|_| ExodusError::DeserializationError)?;

            let r_element =
            	<InnerGroup as frost_core::Group>::deserialize(&r_element_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let z_scalar_serialization =
            	<InnerField as frost_core::Field>::Serialization::try_from(z_scalar_bytes)
            		.map_err(|_| ExodusError::DeserializationError)?;

            let z_scalar =
            	<InnerField as frost_core::Field>::deserialize(&z_scalar_serialization)
                	.map_err(|_| ExodusError::DeserializationError)?;

            let signature = f::Signature::new(r_element, z_scalar);

            verifying_key.verify(message, &signature)
                .map_err(|_| ExodusError::InvalidSignature)
        })
    }

    fn generate_nonce_commitment<R: CryptoRng + RngCore>(
        &self,
        secret_key_package_bytes: &[u8],
        mut rng: R
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), ExodusError> {
        with_ciphersuite!(self, |f| {
            let signing_share =
                f::keys::KeyPackage::deserialize(secret_key_package_bytes)
                    .map(|key_package| *key_package.signing_share())
                    .map_err(|_| ExodusError::DeserializationError)?;

            let (signing_nonce, signing_commitment) =
                f::round1::commit(&signing_share, &mut rng);

            let signing_nonce_bytes = signing_nonce.serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let hiding_nonce_bytes = signing_commitment.hiding()
            	.serialize()
            	.map_err(|_| ExodusError::SerializationError)?;

            let binding_nonce_bytes = signing_commitment.binding()
            	.serialize()
            	.map_err(|_| ExodusError::SerializationError)?;

            Ok((signing_nonce_bytes, hiding_nonce_bytes, binding_nonce_bytes))
        })
    }

    fn generate_group_commitment(
        &self,
        authority_index: AuthIndex,
        nonce_commitments_bytes: &Self::NoncePackages<'_>,
        verifying_key_bytes: &[u8],
        message: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<ExodusHash>, ExodusHash), ExodusError> {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;
            type InnerField = <InnerGroup as frost_core::Group>::Field;

            let verifying_key = f::VerifyingKey::deserialize(verifying_key_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let signing_commitments = nonce_commitments_bytes
                .iter()
                .filter_map(|(&index, &(hiding_bytes, binding_bytes))| {
                    let non_zero_index = Self::non_zero_index(index).ok()?;
                    let identifier = non_zero_index.try_into().ok()?;

                    let hiding_serialization = <InnerGroup as frost_core::Group>::Serialization::try_from(hiding_bytes.as_ref()).ok()?;
                    let hiding_nonce = f::round1::NonceCommitment::deserialize(&hiding_serialization).ok()?;

                    let binding_serialization = <InnerGroup as frost_core::Group>::Serialization::try_from((*binding_bytes).as_ref()).ok()?;
                    let binding_nonce = f::round1::NonceCommitment::deserialize(&binding_serialization).ok()?;

                    let commitments = f::round1::SigningCommitments::new(hiding_nonce, binding_nonce);

                    Some((identifier, commitments))
                })
                .collect::<BTreeMap<_, _>>();

            if nonce_commitments_bytes.len() != signing_commitments.len() {
                return Err(ExodusError::InvalidParticipantId);
            }

            let signing_package = f::SigningPackage::new(signing_commitments, message);

            let binding_factors_list_wrapped =
                frost_core::compute_binding_factor_list::<Ciphersuite>(&signing_package, &verifying_key, &[])
                    .map_err(|err| match err {
                        f::Error::SerializationError => ExodusError::SerializationError,
                        _ => ExodusError::Unknown,
                    })?;

            let group_commitment =
                frost_core::compute_group_commitment::<Ciphersuite>(&signing_package, &binding_factors_list_wrapped)
                    .map_err(|err| match err {
                        f::Error::IdentityCommitment => ExodusError::IdentityCommitment,
                        f::Error::UnknownIdentifier => ExodusError::UnknownIdentifier,
                        _ => ExodusError::Unknown,
                    })?;

            let binding_factors_list = nonce_commitments_bytes
                .iter()
                .filter_map(|(&index, _)| {
                    let non_zero_index = Self::non_zero_index(index).ok()?;
                    let identifier = non_zero_index.try_into().ok()?;

                    let binding_factor = binding_factors_list_wrapped
                        .get(&identifier)
                        .map(|binding_factor| {
                            let fixed_bytes =
                            	<InnerField as frost_core::Field>::Serialization::try_from(
                                	binding_factor.serialize().as_slice()
                            	).ok()?;
                            <InnerField as frost_core::Field>::deserialize(&fixed_bytes).ok()
                        })
                    	.flatten()?;

                    Some((identifier, binding_factor))
                })
            	.collect::<BTreeMap<_, _>>();

            if nonce_commitments_bytes.len() != binding_factors_list.len() {
                return Err(ExodusError::InvalidParticipantId);
            }

            let binding_factor_tree = Self::build_merkle_tree(
                &binding_factors_list,
                |binding_factor| {
                    let serialized = <InnerField as frost_core::Field>::serialize(binding_factor);
	                Ok(serialized.as_ref().to_vec())
                }
            )?;

            let (self_binding_factor_bytes, binding_factors_proof) =
            	Self::generate_merkle_proof_from_tree(
                	&binding_factors_list,
                	&binding_factor_tree,
                	authority_index,
                    |binding_factor| {
	                    let serialized = <InnerField as frost_core::Field>::serialize(binding_factor);
	                    Ok(serialized.as_ref().to_vec())
	                }
            	)?;

            let group_commitment_bytes =
            	<InnerGroup as frost_core::Group>::serialize(&group_commitment.to_element())
                	.map(|serialization| serialization.to_vec())
            		.map_err(|_| ExodusError::SerializationError)?;

            let binding_factors_root = binding_factor_tree
            	.last()
            	.cloned()
            	.ok_or(ExodusError::IncorrectNumberOfPackages)?;

            Ok((
                group_commitment_bytes,
                self_binding_factor_bytes,
                binding_factors_proof,
                binding_factors_root,
            ))
        })
    }

    fn init_signing_package(
        &self,
        nonce_commitments_bytes: &Self::NoncePackages<'_>,
        message: &[u8],
    ) -> Result<Vec<u8>, ExodusError> {
        with_ciphersuite!(self, |f| {
            type InnerGroup = <Ciphersuite as frost_core::Ciphersuite>::Group;

            let signing_commitments = nonce_commitments_bytes
                .iter()
                .map(|(&index, &(hiding_bytes, binding_bytes))| {
                    let identifier = Self::non_zero_index(index)?
                        .try_into()
                        .map_err(|_| ExodusError::InvalidParticipantId)?;

                    let hiding_serialization =
                    	<InnerGroup as frost_core::Group>::Serialization::try_from(hiding_bytes.as_ref())
                            .map_err(|_| ExodusError::DeserializationError)?;

                    let hiding_nonce =
                        f::round1::NonceCommitment::deserialize(&hiding_serialization)
                            .map_err(|_| ExodusError::DeserializationError)?;

                    let binding_serialization =
                    	<InnerGroup as frost_core::Group>::Serialization::try_from((*binding_bytes).as_ref())
                            .map_err(|_| ExodusError::DeserializationError)?;

                    let binding_nonce =
                        f::round1::NonceCommitment::deserialize(&binding_serialization)
                            .map_err(|_| ExodusError::DeserializationError)?;

                    let commitments = f::round1::SigningCommitments::new(hiding_nonce, binding_nonce);

                    Ok((identifier, commitments))
                })
                .collect::<Result<BTreeMap<f::Identifier, f::round1::SigningCommitments>, ExodusError>>()?;

            f::SigningPackage::new(signing_commitments, message)
                .serialize()
                .map_err(|_| ExodusError::SerializationError)
        })
    }

    fn partial_sign_message(
        &self,
        signing_package_bytes: &[u8],
        signer_nonces_bytes: &[u8],
        secret_share_bytes: &[u8],
    ) -> Result<Vec<u8>, ExodusError> {
        with_ciphersuite!(self, |f| {
            let signing_package = f::SigningPackage::deserialize(signing_package_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let signer_nonces = f::round1::SigningNonces::deserialize(signer_nonces_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let secret_share = f::keys::KeyPackage::deserialize(secret_share_bytes)
                .map_err(|_| ExodusError::DeserializationError)?;

            let signature_shares =
                f::round2::sign(&signing_package, &signer_nonces, &secret_share)
                    .map_err(|err| match err {
                        f::Error::IncorrectNumberOfCommitments => ExodusError::IncorrectNumberOfCommitments,
                        f::Error::MissingCommitment => ExodusError::MissingCommitment,
                        f::Error::IncorrectCommitment => ExodusError::IncorrectCommitment,
                        f::Error::SerializationError => ExodusError::SerializationError,
                        f::Error::DuplicatedIdentifier => ExodusError::DuplicatedIdentifier,
                        f::Error::UnknownIdentifier => ExodusError::UnknownIdentifier,
                        f::Error::IncorrectNumberOfIdentifiers => ExodusError::IncorrectNumberOfIdentifiers,
                        _ => ExodusError::Unknown,
                    })?;

            Ok(signature_shares.serialize())
        })
    }
}