use sp_std::{
    vec::Vec,
    result::Result,
    marker::PhantomData,
};
use rand_chacha::rand_core::{CryptoRng, RngCore};
use ghost_traits::exodus::EllipticCurveDiffieHellman;

use hkdf::Hkdf as KeyDerivationFunction;
use sha2::Sha256 as Hashing;

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

use crate::pallet::Config;
use crate::{AuthIndex, EncryptionData, ExodusError, NetworkCurve, DkgIndex};

impl<T: Config> EllipticCurveDiffieHellman<AuthIndex, EncryptionCipher, EncryptionData<T>, DkgIndex, ExodusError> for NetworkCurve {
    fn ecdh_encrypt_package<R: RngCore + CryptoRng>(
        &self,
        cipher: &EncryptionCipher,
        aad: &[u8],
        msg: &[u8],
        mut rng: R,
    ) -> Result<EncryptionData<T>, ExodusError> {
        let payload = aead::Payload { msg, aad };

        let mut nonce_bytes = [0u8; crate::ENCRYPTION_NONCE_MAX_BYTES as usize];
        rng.fill_bytes(&mut nonce_bytes);
        let nonce = EncryptionNonce::from_slice(&nonce_bytes);

        let ciphertext = cipher.encrypt(&nonce, payload)
            .map_err(|_| ExodusError::EncryptionFailed)?;

        EncryptionData::try_new(ciphertext, nonce.as_slice())
    }

    fn ecdh_decrypt_package(
        &self,
        cipher: &EncryptionCipher,
        encrypted_data: &EncryptionData<T>,
        aad: &[u8],
    ) -> Result<Vec<u8>, ExodusError> {
        let nonce = EncryptionNonce::from_iter(
            encrypted_data.nonce.iter().cloned()
        );
        let payload = aead::Payload { msg: &encrypted_data.ciphertext, aad };
        cipher.decrypt(&nonce, payload).map_err(|_| ExodusError::DecryptionFailed)
    }

    fn ecdh_prepare_additional_info(
        &self,
        sender_index: AuthIndex,
        receiver_index: AuthIndex,
        dkg_index: DkgIndex,
        _marker: PhantomData<EncryptionData<T>>,
    ) -> Vec<u8> {
        let mut additional_info = match self {
            NetworkCurve::Secp256k1 => b"EXODUS-SECP256K1-DKG-V1".to_vec(),
            NetworkCurve::Ed25519 => b"EXODUS-ED25519-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(&dkg_index.to_be_bytes());

        additional_info
    }

    fn ecdh_get_cipher_from_keys(
        &self,
        public: &[u8],
        secret: &[u8],
        additional_info: &[u8],
        _marker: PhantomData<EncryptionData<T>>,
    ) -> Result<EncryptionCipher, ExodusError> {
        with_ciphersuite!(self, |f| {
            let secret_package = f::keys::dkg::round1::SecretPackage::deserialize(secret)
                .map_err(|_| ExodusError::DeserializationError)?;
            let public_package = f::keys::dkg::round1::Package::deserialize(public)
                .map_err(|_| ExodusError::DeserializationError)?;

            let coefficients = secret_package.coefficients();
            let local_sk = coefficients.first()
                .ok_or(ExodusError::IncorrectNumberOfCoefficients)?;
            let remote_pk = public_package.commitment().coefficients().first()
                .ok_or(ExodusError::MissingCommitment)?.value();

            let shared_secret_bytes = f::VerifyingKey::new(remote_pk * local_sk)
                .serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let local_pk: f::VerifyingKey = f::SigningKey::from_scalar(*local_sk)
                .map_err(|_| ExodusError::MalformedSigningKey)?
                .into();

            let local_pk_bytes = local_pk.serialize()
                .map_err(|_| ExodusError::SerializationError)?;
            let remote_pk_bytes = f::VerifyingKey::new(remote_pk)
                .serialize()
                .map_err(|_| ExodusError::SerializationError)?;

            let mut keys = [local_pk_bytes, remote_pk_bytes];
            keys.sort();

            let salt = keys.concat();
            let hk = KeyDerivationFunction::<Hashing>::new(Some(&salt), &shared_secret_bytes);

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

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

            Ok(cipher)
        })
    }
}