use crate::generators::aes_ctr::states::State;
use crate::generators::aes_ctr::{AesBlockCipher, AesCtrGenerator, TableIndex};
use crate::generators::{widening_mul, BytesPerChild, ChildrenCount, ForkError};

/// A type alias for the children iterator closure type.
pub type ParallelChildrenClosure<BlockCipher> =
    fn((usize, (Box<BlockCipher>, TableIndex, BytesPerChild))) -> AesCtrGenerator<BlockCipher>;

/// A type alias for the parallel children iterator type.
pub type ParallelChildrenIterator<BlockCipher> = rayon::iter::Map<
    rayon::iter::Zip<
        rayon::range::Iter<usize>,
        rayon::iter::RepeatN<(Box<BlockCipher>, TableIndex, BytesPerChild)>,
    >,
    fn((usize, (Box<BlockCipher>, TableIndex, BytesPerChild))) -> AesCtrGenerator<BlockCipher>,
>;

impl<BlockCipher: AesBlockCipher> AesCtrGenerator<BlockCipher> {
    /// Tries to fork the current generator into `n_child` generators each able to output
    /// `child_bytes` random bytes as a parallel iterator.
    ///
    /// # Notes
    ///
    /// This method requires the "parallel" feature.
    pub fn par_try_fork(
        &mut self,
        n_children: ChildrenCount,
        n_bytes: BytesPerChild,
    ) -> Result<ParallelChildrenIterator<BlockCipher>, ForkError>
    where
        BlockCipher: Send + Sync,
    {
        use rayon::prelude::*;

        if n_children.0 == 0 {
            return Err(ForkError::ZeroChildrenCount);
        }
        if n_bytes.0 == 0 {
            return Err(ForkError::ZeroBytesPerChild);
        }
        if !self.is_fork_in_bound(n_children, n_bytes) {
            return Err(ForkError::ForkTooLarge);
        }

        // rayon zip is only implemented for IndexedParallelIterator which only exists for ranges of
        // u8, u16, u32, usize, i8, i16, i32, isize.
        // This check can fail on 32b platforms
        if n_children.0 > (usize::MAX as u64) {
            return Err(ForkError::ForkTooLarge);
        }

        // The state currently stored in the parent generator points to the table index of the last
        // generated byte. The first index to be generated is the next one :
        let first_index = self.state.table_index().incremented();
        let output = (0..n_children.0 as usize)
            .into_par_iter()
            .zip(rayon::iter::repeat_n(
                (self.block_cipher.clone(), first_index, n_bytes),
                n_children.0 as usize,
            ))
            .map(
                // This map is a little weird because we need to cast the closure to a fn pointer
                // that matches the signature of `ChildrenIterator<BlockCipher>`. Unfortunately,
                // the compiler does not manage to coerce this one automatically.
                (|(i, (block_cipher, first_index, n_bytes))| {
                    // The first index to be outputted by the child is the `first_index` shifted by
                    // the proper amount of `child_bytes`.
                    let child_first_index =
                        first_index.increased(widening_mul(n_bytes.0, i as u64));
                    // The bound of the child is the first index of its next sibling.
                    let child_bound_index =
                        first_index.increased(widening_mul(n_bytes.0, (i + 1) as u64));
                    AesCtrGenerator::from_block_cipher(
                        block_cipher,
                        child_first_index,
                        child_bound_index,
                    )
                }) as ParallelChildrenClosure<BlockCipher>,
            );
        // The parent next index is the bound of the last child.
        let next_index = first_index.increased(widening_mul(n_bytes.0, n_children.0));
        self.state = State::new(next_index);

        Ok(output)
    }
}

#[cfg(test)]
pub mod aes_ctr_parallel_generic_tests {

    use super::*;
    use crate::generators::aes_ctr::aes_ctr_generic_test::{any_key, any_valid_fork};
    use rayon::prelude::*;

    const REPEATS: usize = 1_000_000;

    /// Check the property:
    ///     On a valid fork, the table index of the first child is the same as the table index of
    ///     the parent before the fork.
    pub fn prop_fork_first_state_table_index<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let original_generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let mut forked_generator = original_generator.clone();
            let first_child = forked_generator
                .par_try_fork(nc, nb)
                .unwrap()
                .find_first(|_| true)
                .unwrap();
            assert_eq!(original_generator.table_index(), first_child.table_index());
        }
    }

    /// Check the property:
    ///     On a valid fork, the table index of the first byte outputted by the parent after the
    ///     fork, is the bound of the last child of the fork.
    pub fn prop_fork_last_bound_table_index<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let mut parent_generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let last_child = parent_generator
                .par_try_fork(nc, nb)
                .unwrap()
                .find_last(|_| true)
                .unwrap();
            assert_eq!(
                parent_generator.table_index().incremented(),
                last_child.get_bound()
            );
        }
    }

    /// Check the property:
    ///     On a valid fork, the bound of the parent does not change.
    pub fn prop_fork_parent_bound_table_index<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let original_generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let mut forked_generator = original_generator.clone();
            forked_generator
                .par_try_fork(nc, nb)
                .unwrap()
                .find_last(|_| true)
                .unwrap();
            assert_eq!(original_generator.get_bound(), forked_generator.get_bound());
        }
    }

    /// Check the property:
    ///     On a valid fork, the parent table index is increased of the number of children
    ///     multiplied by the number of bytes per child.
    pub fn prop_fork_parent_state_table_index<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let original_generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let mut forked_generator = original_generator.clone();
            forked_generator
                .par_try_fork(nc, nb)
                .unwrap()
                .find_last(|_| true)
                .unwrap();
            assert_eq!(
                forked_generator.table_index(),
                // Decrement accounts for the fact that the table index stored is the previous one
                t.increased(widening_mul(nc.0, nb.0)).decremented()
            );
        }
    }

    /// Check the property:
    ///     On a valid fork, the bytes outputted by the children in the fork order form the same
    ///     sequence the parent would have had outputted no fork had happened.
    pub fn prop_fork<G: AesBlockCipher>() {
        for _ in 0..1000 {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let bytes_to_go = nc.0 * nb.0;
            let original_generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let mut forked_generator = original_generator.clone();
            let initial_output: Vec<u8> = original_generator
                .take(usize::try_from(bytes_to_go).unwrap())
                .collect();
            let forked_output: Vec<u8> = forked_generator
                .par_try_fork(nc, nb)
                .unwrap()
                .flat_map(|child| child.collect::<Vec<_>>())
                .collect();
            assert_eq!(initial_output, forked_output);
        }
    }

    /// Check the property:
    ///     On a valid fork, all children got a number of remaining bytes equals to the number of
    ///     bytes per child given as fork input.
    pub fn prop_fork_children_remaining_bytes<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let mut generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            assert!(generator
                .par_try_fork(nc, nb)
                .unwrap()
                .all(|c| c.remaining_bytes().0 == nb.0 as u128));
        }
    }

    /// Check the property:
    ///     On a valid fork, the number of remaining bytes of the parent is reduced by the
    ///     number of children multiplied by the number of bytes per child.
    pub fn prop_fork_parent_remaining_bytes<G: AesBlockCipher>() {
        for _ in 0..REPEATS {
            let (t, nc, nb, i) = any_valid_fork().next().unwrap();
            let k = any_key().next().unwrap();
            let bytes_to_go = nc.0 * nb.0;
            let mut generator = AesCtrGenerator::<G>::new(
                k,
                Some(t),
                Some(t.increased(widening_mul(nc.0, nb.0) + i)),
            );
            let before_remaining_bytes = generator.remaining_bytes();
            let _ = generator.par_try_fork(nc, nb).unwrap();
            let after_remaining_bytes = generator.remaining_bytes();
            assert_eq!(
                before_remaining_bytes.0 - after_remaining_bytes.0,
                bytes_to_go as u128
            );
        }
    }
}
