use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        circuits::x25519::get_shared_secret::get_shared_secret,
    },
    traits::{FromF25519, MxeRescueKey, MxeX25519PrivateKey, Reveal, ToMontgomery},
    utils::{
        crypto::{
            key::{RescueKey, X25519PrivateKey, X25519PublicKey, RESCUE_KEY_COUNT},
            rescue_desc::{RescueArg, RescueDesc},
            rescue_prime_hash::RescuePrimeHash,
        },
        curve_point::Curve,
        elliptic_curve::F25519,
        matrix::Matrix,
    },
};
use std::ops::Mul;

/// The Arcis Rescue cipher. This is an MPC implementation of the Marvellous Rescue cipher,
/// see <https://tosc.iacr.org/index.php/ToSC/article/view/8695/8287>. We use it in Counter (CTR)
/// mode, see <https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38a.pdf> (Section 6.5).
#[allow(dead_code)]
pub struct RescueCipher<F: ActuallyUsedField, T: RescueArg<F>> {
    desc: RescueDesc<F, T>,
}

impl<F: ActuallyUsedField, T: RescueArg<F>> RescueCipher<F, T> {
    pub fn new(key: RescueKey<T>) -> Self {
        RescueCipher {
            desc: RescueDesc::new_cipher_desc(Matrix::from(key)),
        }
    }

    /// Given a client public key:
    /// - perform the x25519 key exchange with the MXE private key
    /// - convert the output to `Vec<T>`
    /// - perform a key derivation, following [Section 4, Option 1.](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Cr2.pdf),
    ///   with the Arcis Rescue-Prime hash function.
    pub fn new_with_client<
        Base: F25519,
        S: Clone + Copy + MxeX25519PrivateKey + Mul<C, Output = C>,
        C: Curve + ToMontgomery<Output = Base>,
    >(
        public_key: X25519PublicKey<C>,
    ) -> Self
    where
        T: FromF25519<Base>,
    {
        Self::new_with_client_from_keys(X25519PrivateKey::<S>::mxe_private_key(), public_key)
    }

    /// Given a private key and a client public key:
    /// - perform the x25519 key exchange
    /// - convert the output to `Vec<T>`
    /// - perform a key derivation, following [Section 4, Option 1.](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Cr2.pdf),
    ///   with the Arcis Rescue-Prime hash function.
    pub fn new_with_client_from_keys<
        Base: F25519,
        S: Clone + Copy + Mul<C, Output = C>,
        C: Curve + ToMontgomery<Output = Base>,
    >(
        private_key: X25519PrivateKey<S>,
        client_public_key: X25519PublicKey<C>,
    ) -> Self
    where
        T: FromF25519<Base>,
    {
        let shared_secret = get_shared_secret(private_key, client_public_key);
        let converted = T::from_F25519(shared_secret);
        let hasher = RescuePrimeHash::new();
        // We follow [Section 4, Option 1.](https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Cr2.pdf).
        // For our choice of hash function, we have:
        // - H_outputBits = hasher.digest_len = RESCUE_KEY_COUNT
        // - max_H_inputBits = arbitrarily long, as the Rescue-Prime hash function is built upon the
        //   sponge construction
        // - L = RESCUE_KEY_COUNT.

        // Build the vector `counter || Z || FixedInfo` (we only have i = 1, since reps = 1).
        let mut counter = vec![T::from(F::from(1))];
        counter.extend(converted);
        // For the FixedInfo we simply take L.
        counter.push(T::from(F::from(RESCUE_KEY_COUNT as u64)));
        Self::new(RescueKey::new_from_inner(hasher.digest(counter)))
    }

    pub fn new_for_mxe() -> Self
    where
        T: MxeRescueKey,
    {
        Self::new(RescueKey::<T>::mxe_rescue_key())
    }

    fn get_counter(nonce: T, n_blocks: usize) -> Vec<T> {
        (0..n_blocks).fold(Vec::new(), |mut acc, i| {
            acc.extend([
                nonce,
                T::from(F::from(i as u64)),
                T::from(F::ZERO),
                T::from(F::ZERO),
                T::from(F::ZERO),
            ]);
            acc
        })
    }

    fn encrypt_batch(desc: &RescueDesc<F, T>, ptxt: &[T], counter: &[T], output: &mut Vec<T>)
    where
        T: Reveal,
    {
        let n_ptxt = ptxt.len();
        let encrypted_counter = desc.permute(&Matrix::from(counter.to_vec()));
        output.extend(
            (Matrix::from(ptxt)
                + Matrix::from(
                    encrypted_counter
                        .into_iter()
                        .take(n_ptxt)
                        .collect::<Vec<T>>(),
                ))
            .into_iter()
            .map(|c| c.reveal()),
        )
    }
    /// Encrypt the masked plaintext vector in Counter (CTR) mode.
    pub fn encrypt(&self, masked_plaintext: Vec<T>, nonce: T) -> Vec<T>
    where
        T: Reveal,
    {
        // According to https://crypto.stackexchange.com/questions/1666/can-i-safely-replace-xor-with-add-in-a-stream-cipher
        // we can mask `ptxt` by `desc.permute(counter)` via addition (as opposed to XORing)
        // without any loss of security.
        let plaintext = masked_plaintext;
        let n_blocks = plaintext.len().div_ceil(self.desc.m);
        let counter = RescueCipher::get_counter(nonce, n_blocks);

        let mut result = Vec::new();
        for i in 0..n_blocks {
            let cnt = self.desc.m * i;
            Self::encrypt_batch(
                &self.desc,
                &plaintext[cnt..(cnt + self.desc.m).min(plaintext.len())],
                &counter[cnt..cnt + self.desc.m],
                &mut result,
            )
        }
        result
    }

    fn decrypt_batch(desc: &RescueDesc<F, T>, ctxt: &[T], counter: &[T], output: &mut Vec<T>) {
        let n_ctxt = ctxt.len();
        let encrypted_counter = desc.permute(&Matrix::from(counter.to_vec()));
        output.extend(
            Matrix::from(ctxt)
                - Matrix::from(
                    encrypted_counter
                        .into_iter()
                        .take(n_ctxt)
                        .collect::<Vec<T>>(),
                ),
        )
    }

    /// Decrypt the ciphertext vector in Counter (CTR) mode.
    pub fn decrypt(&self, ciphertext: Vec<T>, nonce: T) -> Vec<T> {
        // According to https://crypto.stackexchange.com/questions/1666/can-i-safely-replace-xor-with-add-in-a-stream-cipher
        // we can unmask `ctxt` by `desc.permute(counter)` via subtraction (as opposed to
        // XORing) without any loss of security.
        let n_blocks = ciphertext.len().div_ceil(self.desc.m);
        let counter = RescueCipher::get_counter(nonce, n_blocks);

        let mut decrypted = Vec::new();
        for i in 0..n_blocks {
            let cnt = self.desc.m * i;
            Self::decrypt_batch(
                &self.desc,
                &ciphertext[cnt..(cnt + self.desc.m).min(ciphertext.len())],
                &counter[cnt..cnt + self.desc.m],
                &mut decrypted,
            );
        }
        decrypted
    }
}
