use crate::{EncryptedMessage, FrostError};
use ghost_networks::NetworkCurve;

use chacha20poly1305::{
    aead::{self, Aead, AeadCore, KeyInit},
    Key as EncryptionKey,
    Nonce as EncryptionNonce,
    ChaCha20Poly1305 as EncryptionCipher,
};

use hkdf::Hkdf as KeyDerivation;
use sha2::{Digest, Sha256 as Hashing};

use rand_chacha::rand_core::{CryptoRng, RngCore};
use sp_std::collections::btree_map::BTreeMap;

macro_rules! with_ciphersuite {
    ($curve:expr, |$alias:ident| $body:block) => {
        match $curve {
            ExodusCurve::Secp256k1 => {
                use frost_secp256k1 as $alias;
                $body
            },
            NetworkCurve::Ed25519 => {
                use frost_ed25519 as $alias;
                $body

            }
        }
    }
}

impl ExodusCurve {
    pub fn dkg_verify_proof_of_knowledge(
        &self,
        index: u16,
        round1_package_bytes: &[u8],
    ) -> Result<(), FrostError> {
        with_ciphersuite!(self, |f| {
            let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
            let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            frost_core::keys::dkg::verify_proof_of_knowledge(
                identifier,
                &round1_package.commitment(),
                &round1_package.proof_of_knowledge(),
            ).map_err(|_| FrostError::InvalidProofOfKnowledge)
        })
    }

    pub fn verify_signature_share(
        &self,
        index: u16,
        verifying_share_bytes: &[u8],
        signature_share_bytes: &[u8],
        signing_package_bytes: &[u8],
        verifying_key_bytes: &[u8],
    ) -> Result<(), FrostError> {
        with_ciphersuite!(self, |f| {
            let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
            let verifying_share = f::keys::VerifyingShare::deserialize(verifying_share_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

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

            let signing_package = f::SigningPackage::deserialize(signing_package_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

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

            frost_core::verify_signature_share(
                identifier,
                &verifying_share,
                &signature_share,
                &signing_package,
                &verifying_key,
            ).map_err(|err| match err {
                f::Error::IdentityCommitment => FrostError::IdentityCommitment,
                f::Error::UnknownIdentifier => FrostError::UnknownIdentifier,
                f::Error::DuplicatedIdentifier => FrostError::DuplicatedIdentifier,
                f::Error::InvalidSignatureShare { .. } => FrostError::InvalidSignatureShare,
                _ => FrostError::Unknown,
            })
        })
    }

    pub fn aggregate_signature(
        &self,
        signing_package_bytes: &[u8],
        signature_shares_bytes: &BTreeMap<u16, Vec<u8>>,
        pubkeys_bytes: &[u8],
    ) -> Result<Vec<u8>, FrostError> {
        with_ciphersuite!(self, |f| {
            let signing_package = f::SigningPackage::deserialize(signing_package_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            let mut signature_shares = BTreeMap::new();
            for (&index, signature_share_bytes) in signature_shares_bytes.iter() {
                let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
                let signature_share = f::round2::SignatureShare::deserialize(signature_share_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;
                signature_shares.insert(identifier, signature_share);
            }

            let pubkeys = f::keys::PublicKeyPackage::deserialize(pubkeys_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            let signature = f::aggregate(&signing_package, &signature_shares, &pubkeys)
                .map_err(|err| match err {
                    f::Error::UnknownIdentifier => FrostError::UnknownIdentifier,
                    f::Error::SerializationError => FrostError::SerializationError,
                    f::Error::IdentityCommitment => FrostError::IdentityCommitment,
                    f::Error::IncorrectNumberOfIdentifiers => FrostError::IncorrectNumberOfIdentifiers,
                    f::Error::DuplicatedIdentifier => FrostError::DuplicatedIdentifier,
                    f::Error::InvalidSecretShare { .. } => FrostError::InvalidSecretShare,
                    f::Error::InvalidSignature { .. } => FrostError::InvalidSignature,
                    _ => FrostError::Unknown,
                })?;

            let signature_bytes = signature.serialize()
                .map_err(|_| FrostError::SerializationError)?;

            Ok(signature_bytes)
        })
    }

    pub fn is_signature_valid(
        &self,
        pubkey_bytes: &[u8],
        signature_bytes: &[u8],
        message: &[u8],
    ) -> bool {
        with_ciphersuite!(self, |f| {
            let get_verification_result = || -> Result<(), FrostError> {
                let pubkeys = f::keys::PublicKeyPackage::deserialize(pubkey_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;

                let signature = f::Signature::deserialize(signature_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;

                pubkeys.verifying_key()
                    .verify(message, &signature)
                    .map_err(|_| FrostError::InvalidSignature)?;

                Ok(())
            };

            get_verification_result().is_ok()
        })
    }

    pub fn dkg_part1<R: RngCore + CryptoRng>(
        &self,
        index: u16,
        max_signers: u16,
        min_signers: u16,
        mut rng: R,
    ) -> Result<(Vec<u8>, Vec<u8>), FrostError> {
        with_ciphersuite!(self, |f| {
            let identifier = index.try_into().map_err(|_| FrostError::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 => FrostError::InvalidMinSigners(min_signers, max_signers),
                    f::Error::InvalidMaxSigners => FrostError::InvalidMaxSigners(max_signers),
                    f::Error::DKGNotSupported => FrostError::DKGNotSupported(*self),
                    _ => FrostError::Unknown,
                }
            })?;

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


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

            Ok((round1_secret_package_bytes, round1_package_bytes))
        })
    }

    fn convert_identifier_to_index(
        identifier_bytes: Vec<u8>
    ) -> Result<u16, FrostError> {
        const SIZE: usize = core::mem::size_of::<u16>();

        let start = identifier_bytes.len().checked_sub(SIZE)
            .ok_or(FrostError::InvalidParticipantId)?;

        let bytes_array = identifier_bytes.get(start..)
            .and_then(|slice| slice.try_into().ok())
            .ok_or(FrostError::InvalidParticipantId)?;

        Ok(u16::from_be_bytes(bytes_array))
    }

    pub fn dkg_part2(
        &self,
        round1_secret_package_bytes: &[u8],
        round1_packages_bytes: &BTreeMap<u16, Vec<u8>>,
    ) -> Result<(Vec<u8>, BTreeMap<u16, Vec<u8>>), FrostError> {
        with_ciphersuite!(self, |f| {
            let round1_secret_package = f::keys::dkg::round1::SecretPackage::deserialize(round1_secret_package_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            let mut round1_packages = BTreeMap::new();
            for (&index, round1_package_bytes) in round1_packages_bytes.iter() {
                let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
                let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;

                round1_packages.insert(identifier, round1_package);
            }

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

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

            let mut round2_packages_bytes = BTreeMap::new();
            for (participant_identifier, round2_package) in round2_packages.iter() {
                let index = Self::convert_identifier_to_index(
                    participant_identifier.serialize(),
                )?;

                let round2_package_bytes = round2_package.serialize()
                .map_err(|_| FrostError::SerializationError)?;

                round2_packages_bytes.insert(index as u16, round2_package_bytes);
            }

            Ok((round2_secret_package_bytes, round2_packages_bytes))
        })
    }

    pub fn dkg_encrypt_round2_package<R: RngCore + CryptoRng>(
        &self,
        cipher: &EncryptionCipher,
        associated_data: &[u8],
        round2_package_bytes: &[u8],
        mut rng: R,
    ) -> Result<(Vec<u8>, Vec<u8>), FrostError> {
        let payload = aead::Payload {
            msg: round2_package_bytes,
            aad: associated_data,
        };

        let nonce = EncryptionCipher::generate_nonce(&mut rng);
        let ciphertext = cipher.encrypt(&nonce, payload)
            .map_err(|_| FrostError::EncryptionFailed)?;

        Ok((ciphertext, nonce.to_vec()))
    }

    pub fn dkg_decrypt_round2_package(
        &self,
        cipher: &EncryptionCipher,
        encrypted_message: &EncryptedMessage,
        associated_data: &[u8],
    ) -> Result<Vec<u8>, FrostError> {
        let nonce = EncryptionNonce::from_iter(encrypted_message.nonce.iter().cloned());
        let payload = aead::Payload {
            msg: &encrypted_message.ciphertext,
            aad: associated_data,
        };

        cipher.decrypt(&nonce, payload)
            .map_err(|_| FrostError::DecryptionFailed)
    }

    pub fn prepare_additional_info(
        &self,
        sender_index: u16,
        receiver_index: u16,
        session_index: u16,
    ) -> Vec<u8> {
        let mut additional_info = match self {
            ExodusCurve::Secp256k1 => b"EXODUS-SECP256K1-DKG-V1".to_vec(),
        };

        additional_info.extend_from_slice(&sender_index.to_be_bytes());
        additional_info.extend_from_slice(&receiver_index.to_be_bytes());
        additional_info.extend_from_slice(&session_index.to_be_bytes());

        additional_info
    }

    pub fn prepare_cipher_from_keys(
        &self,
        round1_receiver_package_bytes: &[u8],
        round1_secret_package_bytes: &[u8],
        additional_info: &[u8],
    ) -> Result<EncryptionCipher, FrostError> {
        with_ciphersuite!(self, |f| {
            let coefficients =
                f::keys::dkg::round1::SecretPackage::deserialize(round1_secret_package_bytes)
                    .map_err(|_| FrostError::DeserializationError)?
                    .coefficients();

            let sender_secret_key = coefficients
                .first()
                .ok_or(FrostError::IncorrectNumberOfCoefficients)?;

            let receiver_public_key =
                f::keys::dkg::round1::Package::deserialize(round1_receiver_package_bytes)
                    .map_err(|_| FrostError::DeserializationError)?
                    .commitment()
                    .coefficients()
                    .first()
                    .ok_or(FrostError::MissingCommitment)?
                    .value();

            let shared_secret_bytes = f::VerifyingKey::new(receiver_public_key * sender_secret_key)
                .serialize()
                .map_err(|_| FrostError::SerializationError)?;

            let salt = Hashing::digest(&additional_info);
            let hk = KeyDerivation::<Hashing>::new(Some(&salt), &shared_secret_bytes);

            let mut encryption_key = EncryptionKey::default();
            hk.expand(additional_info, &mut encryption_key)
                .map_err(|_| FrostError::HKDFFailed)?;

            let cipher = EncryptionCipher::new_from_slice(&encryption_key)
                .expect("could not happen; length is already correct qed");

            Ok(cipher)
        })
    }

    pub fn dkg_verify_private_package(
        &self,
        index: u16,
        round1_package_bytes: &[u8],
        round2_package_bytes: &[u8],
    ) -> Result<(), FrostError> {
        with_ciphersuite!(self, |f| {
            let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
            let commitment = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                .map(|round1_package| round1_package.commitment().clone())
                .map_err(|_| FrostError::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(|_| FrostError::DeserializationError)?;

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

            Ok(())
        })
    }

    pub fn dkg_part3(
        &self,
        round2_secret_package_bytes: &[u8],
        round1_packages_bytes: &BTreeMap<u16, Vec<u8>>,
        round2_packages_bytes: &BTreeMap<u16, Vec<u8>>,
    ) -> Result<(Vec<u8>, Vec<u8>), FrostError> {
        with_ciphersuite!(self, |f| {
            let round2_secret_package = f::keys::dkg::round2::SecretPackage::deserialize(round2_secret_package_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            let mut round1_packages = BTreeMap::new();
            for (&index, round1_package_bytes) in round1_packages_bytes.iter() {
                let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
                let round1_package = f::keys::dkg::round1::Package::deserialize(round1_package_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;
                round1_packages.insert(identifier, round1_package);
            }

            let mut round2_packages = BTreeMap::new();
            for (&index, round2_package_bytes) in round2_packages_bytes.iter() {
                let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
                let round2_package = f::keys::dkg::round2::Package::deserialize(round2_package_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;
                round2_packages.insert(identifier, round2_package);
            }

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

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

            let pubkey_package_bytes = pubkey_package.serialize()
                .map_err(|_| FrostError::SerializationError)?;

            Ok((key_package_bytes, pubkey_package_bytes))
        })
    }

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

            let (signing_nonces, signing_commitments) =
                f::round1::commit(&signing_share, &mut rng);

            let signing_nonces_bytes = signing_nonces.serialize()
                .map_err(|_| FrostError::SerializationError)?;

            let signing_commitments_bytes = signing_commitments.serialize()
                .map_err(|_| FrostError::SerializationError)?;

            Ok((signing_nonces_bytes, signing_commitments_bytes))
        })
    }

    pub fn generate_signing_package(
        &self,
        commitments_map_bytes: &BTreeMap<u16, Vec<u8>>,
        message: &[u8],
    ) -> Result<Vec<u8>, FrostError> {
        with_ciphersuite!(self, |f| {
            let mut commitments_map = BTreeMap::new();
            for (&index, commitment_bytes) in commitments_map_bytes.iter() {
                let identifier = index.try_into().map_err(|_| FrostError::InvalidParticipantId)?;
                let commitment = f::round1::SigningCommitments::deserialize(commitment_bytes)
                    .map_err(|_| FrostError::DeserializationError)?;
                commitments_map.insert(identifier, commitment);
            }

            f::SigningPackage::new(commitments_map, message)
                .serialize()
                .map_err(|_| FrostError::SerializationError)
        })
    }

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

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

            let secret_key_share = f::keys::KeyPackage::deserialize(secret_key_share_bytes)
                .map_err(|_| FrostError::DeserializationError)?;

            let signature_share_bytes = f::round2::sign(&signing_package, &signer_nonces, &secret_key_share)
                .map(|signature_share| signature_share.serialize())
                .map_err(|err| match err {
                    f::Error::IncorrectNumberOfCommitments => FrostError::IncorrectNumberOfCommitments,
                    f::Error::MissingCommitment => FrostError::MissingCommitment,
                    f::Error::IncorrectCommitment => FrostError::IncorrectCommitment,
                    f::Error::SerializationError => FrostError::SerializationError,
                    f::Error::UnknownIdentifier => FrostError::UnknownIdentifier,
                    f::Error::IdentityCommitment => FrostError::IdentityCommitment,
                    f::Error::DuplicatedIdentifier => FrostError::DuplicatedIdentifier,
                    f::Error::IncorrectNumberOfIdentifiers => FrostError::IncorrectNumberOfIdentifiers,
                    _ => FrostError::Unknown,
                })?;

            Ok(signature_share_bytes)
        })
    }
}