use crate::{
    core::actually_used_field::ActuallyUsedField,
    utils::{
        crypto::{
            key::RESCUE_KEY_COUNT,
            rescue_desc::{RescueArg, RescueDesc},
        },
        matrix::Matrix,
    },
};
use std::ops::IndexMut;

/// The Arcis Rescue-Prime hash function. This is an MPC implementation of the Rescue-Prime hash
/// function, see <https://eprint.iacr.org/2020/1143.pdf>, and offers 256 bits of security
/// against collision, preimage and second-preimage attacks for any field of size at least 102 bits.
/// We use the sponge construction with fixed rate = 7 and capacity = 5 (i.e., m = 12), and truncate
/// the output to RESCUE_KEY_COUNT field elements.
#[derive(Debug)]
pub struct RescuePrimeHash<F: ActuallyUsedField, T: RescueArg<F>> {
    desc: RescueDesc<F, T>,
    pub rate: usize,
    pub digest_len: usize,
}

impl<F: ActuallyUsedField, T: RescueArg<F>> RescuePrimeHash<F, T> {
    pub fn new() -> Self {
        RescuePrimeHash {
            desc: RescueDesc::new_hash_desc(12, 5),
            rate: 7,
            digest_len: RESCUE_KEY_COUNT,
        }
    }

    // This is Algorithm 1 from https://eprint.iacr.org/2020/1143.pdf, though with the padding (see Algorithm 2).
    // The hash is truncated to digest_len elements.
    // According to Section 2.2, this offers min(F::NUM_BITS / 2 * min(digest_len, capacity), s)
    // bits of security against collision, preimage and second-preimage attacks.
    // The security level is thus of the order of 256 bits for any field of size at least 102 bits.
    // The rate and capacity are chosen to achieve minimal number of rounds 8.
    pub fn digest(&self, mut message: Vec<T>) -> [T; RESCUE_KEY_COUNT] {
        message.push(T::from(F::ONE));
        if message.len() % self.rate != 0 {
            message.extend(vec![
                T::from(F::ZERO);
                self.rate - (message.len() % self.rate)
            ]);
        }
        let mut state = Matrix::new((self.desc.m, 1), T::from(F::ZERO));
        message.chunks(self.rate).for_each(|chunk| {
            chunk.iter().copied().enumerate().for_each(|(i, s)| {
                *state.index_mut((i, 0)) = *state.index_mut((i, 0)) + s;
            });
            state = self.desc.permute(&state);
        });

        state
            .into_iter()
            .take(self.digest_len)
            .collect::<Vec<T>>()
            .try_into()
            .unwrap_or_else(|v: Vec<T>| {
                panic!(
                    "Expected a Vec of length {} (found {})",
                    self.digest_len,
                    v.len()
                )
            })
    }
}

impl<F: ActuallyUsedField, T: RescueArg<F>> Default for RescuePrimeHash<F, T> {
    fn default() -> Self {
        Self::new()
    }
}
