#![warn(clippy::all, clippy::pedantic, clippy::nursery)]

use isa_l::{
    ec_encode_data, ec_init_tables_owned, gf_gen_cauchy1_matrix, gf_gen_decode_matrix_simple,
};

#[derive(Debug)]
pub enum ErrorKind {
    InsufficientFragments,
    NoDecodeMatrix,
}

#[derive(Debug)]
pub struct ErasureCoder {
    k: usize,
    p: usize,
    encode_matrix: Vec<u8>,
    encode_gftbls: Vec<u8>,
}

impl ErasureCoder {
    #[must_use]
    pub const fn k(&self) -> usize {
        self.k
    }

    #[must_use]
    pub const fn p(&self) -> usize {
        self.p
    }

    #[must_use]
    pub const fn m(&self) -> usize {
        self.k + self.p
    }

    #[must_use]
    pub fn encode_matrix(&self) -> &[u8] {
        &self.encode_matrix[..]
    }

    #[must_use]
    pub fn new(k: usize, p: usize) -> Self {
        let encode_matrix = gf_gen_cauchy1_matrix(k, k + p);
        Self::new_with_encode_matrix(k, p, encode_matrix)
    }

    #[must_use]
    pub fn new_with_encode_matrix(k: usize, p: usize, encode_matrix: Vec<u8>) -> Self {
        let encode_gftbls = ec_init_tables_owned(k, p, &encode_matrix[k * k..]);

        Self {
            k,
            p,
            encode_matrix,
            encode_gftbls,
        }
    }

    /// Calculate the parities over `k` data slices of data and place it into `p` buffer slices.
    ///
    /// The length of all data slices and all buffer slices must be equal.
    ///
    /// # Panics
    ///
    /// Panics if the length of `data` is not equal to `k` or if the length of `bufs` is not equal
    /// to `p`.
    pub fn encode<T, M>(&self, data: impl AsRef<[T]>, mut bufs: impl AsMut<[M]>)
    where
        T: AsRef<[u8]>,
        M: AsMut<[u8]>,
    {
        assert_eq!(data.as_ref().len(), self.k);
        assert_eq!(bufs.as_mut().len(), self.p);

        let block_size = data.as_ref().iter().next().unwrap().as_ref().len();

        ec_encode_data(block_size, self.k, self.p, &self.encode_gftbls, data, bufs)
    }

    /// Decode erased blocks given their indexes and the available blocks
    ///
    /// Uses `gf_gen_decode_matrix_simple` from the [`isa-l`] crate.
    ///
    /// # Errors
    ///
    /// This function will return an error if the number of provided blocks is smaller than `k` or
    /// if no decode matrix can be found.
    ///
    /// [`isa-l`]: https://crates.io/crates/isa-l
    pub fn decode<T: AsRef<[u8]>, M: AsMut<[u8]>>(
        &self,
        blocks: &[T],
        erased_idxs: &[usize],
        bufs: impl AsMut<[M]>,
    ) -> Result<(), ErrorKind> {
        self.decode_with_fn(blocks, erased_idxs, gf_gen_decode_matrix_simple, bufs)
    }

    /// Decode erased blocks given their indexes and the available blocks
    ///
    /// Uses a custom function to generate a decode matrix with the signature:
    /// ```ignore
    /// fn gf_gen_decode_matrix_simple(
    ///     encode_matrix: &[u8],
    ///     erased_idxs: &[usize],
    ///     k: usize,
    ///     m: usize,
    /// ) -> Option<Vec<u8>>;
    /// ```
    ///
    /// # Errors
    ///
    /// This function will return an error if the number of provided blocks is smaller than `k` or
    /// if no decode matrix can be found.
    pub fn decode_with_fn<'a, T: AsRef<[u8]>, M: AsMut<[u8]>>(
        &'a self,
        blocks: &[T],
        erased_idxs: &'a [usize],
        gen_decode_matrix: fn(&'a [u8], &'a [usize], usize, usize) -> Option<Vec<u8>>,
        bufs: impl AsMut<[M]>,
    ) -> Result<(), ErrorKind> {
        if blocks.len() < self.k {
            return Err(ErrorKind::InsufficientFragments);
        }

        let nerrs = erased_idxs.len();
        let block_size = blocks.first().unwrap().as_ref().len();

        let decode_matrix = gen_decode_matrix(&self.encode_matrix, erased_idxs, self.k, self.m())
            .ok_or_else(|| ErrorKind::NoDecodeMatrix)?;

        // Recover data
        let decode_g_tables =
            ec_init_tables_owned(self.k, self.p, &decode_matrix[..self.k * self.p]);

        ec_encode_data(
            block_size,
            self.k,
            nerrs,
            &decode_g_tables,
            &blocks[..self.k],
            bufs,
        );
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::ErasureCoder;
    use rand::{RngCore, SeedableRng};
    use rand_pcg::Pcg64Mcg;

    #[test]
    fn test_vandermonde_not_decodable() {
        let k = 10;
        let p = 6;

        #[rustfmt::skip]
        let vandermonde_encode_matrix = vec![
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
            0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
            0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1D, 0x3A,
            0x01, 0x04, 0x10, 0x40, 0x1D, 0x74, 0xCD, 0x13, 0x4C, 0x2D,
            0x01, 0x08, 0x40, 0x3A, 0xCD, 0x26, 0x2D, 0x75, 0x8F, 0x0C,
            0x01, 0x10, 0x1D, 0xCD, 0x4C, 0xB4, 0x8F, 0x18, 0x9D, 0x25,
            0x01, 0x20, 0x74, 0x26, 0xB4, 0x03, 0x60, 0x9C, 0x6A, 0xC1,
        ];

        let c = ErasureCoder::new(k, p);

        let vandermonde_not_decodable_erasures = [[2, 4, 7, 11, 12, 15], [4, 6, 9, 11, 12, 15]];

        for erased_idxs in &vandermonde_not_decodable_erasures {
            let mut decoded_data = vec![vec![0]; p];
            let bufs = decoded_data.iter_mut().collect::<Vec<_>>();
            let is_ok = c.decode(&vec![&vec![1]; 10], erased_idxs, bufs).is_ok();
            assert!(
                is_ok || c.encode_matrix == vandermonde_encode_matrix,
                "Should not be decodable with Vandermonde encode matrix only."
            );
        }
    }

    #[test]
    fn test_encode_decode() {
        let mut rng: Pcg64Mcg = SeedableRng::seed_from_u64(0);

        let k = 12;
        let p = 2;
        let bs = 1024 * 1024;

        let data = {
            let mut dest = vec![0; k * bs];
            rng.fill_bytes(&mut dest);
            dest
        };

        let c = ErasureCoder::new(k, p);

        let mut parity = vec![vec![0; bs]; p];
        let bufs = parity.iter_mut().collect::<Vec<_>>();
        c.encode(data.chunks_exact(bs).collect::<Vec<_>>(), bufs);

        let mut data = data;
        data.extend(parity.into_iter().flatten());
        let encoded_data = data;

        let erased_idxs = vec![0, 12];
        let erased_encoded_data = encoded_data
            .chunks_exact(bs)
            .enumerate()
            .filter_map(|(i, b)| {
                if erased_idxs.contains(&i) {
                    None
                } else {
                    Some(b)
                }
            })
            .collect::<Vec<_>>();

        let mut decoded_data = vec![vec![0; bs]; p];
        let bufs = decoded_data.iter_mut().collect::<Vec<_>>();
        c.decode(&erased_encoded_data, &erased_idxs, bufs).unwrap();

        for (idx, block) in erased_idxs.iter().zip(decoded_data.iter()) {
            assert_eq!(block.as_slice(), &encoded_data[idx * bs..][..bs]);
        }
    }
}