// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub use s2n_quic_core::random::Generator;

/// Provides random number generation support for an endpoint
pub trait Provider: 'static {
    type Generator: 'static + Generator;
    type Error: core::fmt::Display + Send + Sync;

    fn start(self) -> Result<Self::Generator, Self::Error>;
}

pub use self::rand::{Provider as Default, Random, ReseedingRng};

impl_provider_utils!();

mod rand {
    use core::convert::Infallible;
    use rand::{
        rand_core::{Rng, SeedableRng, TryRng},
        rngs::ChaCha20Rng,
    };
    use s2n_quic_core::random;

    // Number of generated bytes after which to reseed the public and private random
    // generators.
    //
    // This value is based on THREAD_RNG_RESEED_THRESHOLD from
    // [rand::rngs::thread.rs](https://github.com/rust-random/rand/blob/ef75e56cf5824d33c55622bf84a70ec6e22761ba/src/rngs/thread.rs#L39)
    const RESEED_THRESHOLD: u64 = 1024 * 64;

    /// A ChaCha-based RNG that periodically reseeds from a given entropy source.
    ///
    /// This replaces the removed `rand::rngs::ReseedingRng` with equivalent
    /// functionality: a ChaCha CSPRNG that reseeds from the entropy source after
    /// `RESEED_THRESHOLD` bytes have been generated.
    #[derive(Debug)]
    pub struct ReseedingRng<S: TryRng + Send + 'static> {
        inner: ChaCha20Rng,
        entropy: S,
        bytes_until_reseed: u64,
    }

    impl<S: TryRng + Send + 'static> ReseedingRng<S> {
        pub fn new(mut entropy: S) -> Self {
            let inner = ChaCha20Rng::try_from_rng(&mut entropy)
                .expect("entropy source failed during initial seed");
            Self {
                inner,
                entropy,
                bytes_until_reseed: RESEED_THRESHOLD,
            }
        }

        pub fn fill_bytes(&mut self, dest: &mut [u8]) {
            // We should first check whether the rng needs to be reseeded before generate random bytes.
            // This matches the behavior of rand: https://github.com/rust-random/rand/blob/0.9.1/src/rngs/reseeding.rs#L163.
            if self.bytes_until_reseed == 0 {
                self.inner = ChaCha20Rng::try_from_rng(&mut self.entropy)
                    .expect("entropy source failed during reseed");
                self.bytes_until_reseed = RESEED_THRESHOLD;
            }
            Rng::fill_bytes(&mut self.inner, dest);
            let len = dest.len() as u64;
            self.bytes_until_reseed = self.bytes_until_reseed.saturating_sub(len);
        }
    }

    /// A random number generator backed by a ChaCha CSPRNG that periodically
    /// reseeds from a given entropy source `S`.
    #[derive(Debug)]
    pub struct Random<S: TryRng + Send + 'static> {
        public: ReseedingRng<S>,
        private: ReseedingRng<S>,
    }

    impl<S: TryRng + Send + 'static> Random<S> {
        pub fn new(public_entropy: S, private_entropy: S) -> Self {
            Self {
                public: ReseedingRng::new(public_entropy),
                private: ReseedingRng::new(private_entropy),
            }
        }
    }

    impl<S: TryRng + Send + 'static> random::Generator for Random<S> {
        fn public_random_fill(&mut self, dest: &mut [u8]) {
            self.public.fill_bytes(dest)
        }

        fn private_random_fill(&mut self, dest: &mut [u8]) {
            self.private.fill_bytes(dest)
        }
    }

    // Default provider using SysRng as entropy source

    use ::rand::rngs::SysRng;

    #[derive(Debug, Default)]
    pub struct Provider(Generator);

    impl super::Provider for Provider {
        type Generator = Generator;
        type Error = Infallible;

        fn start(self) -> Result<Self::Generator, Self::Error> {
            Ok(self.0)
        }
    }

    impl super::TryInto for Generator {
        type Provider = Provider;
        type Error = Infallible;

        fn try_into(self) -> Result<Self::Provider, Self::Error> {
            Ok(Provider(self))
        }
    }

    /// The default random generator, using the OS system RNG as the entropy source.
    #[derive(Debug)]
    pub struct Generator(Random<SysRng>);

    impl Default for Generator {
        fn default() -> Self {
            Self(Random::new(SysRng, SysRng))
        }
    }

    impl random::Generator for Generator {
        fn public_random_fill(&mut self, dest: &mut [u8]) {
            self.0.public_random_fill(dest)
        }

        fn private_random_fill(&mut self, dest: &mut [u8]) {
            self.0.private_random_fill(dest)
        }
    }

    #[cfg(test)]
    mod tests {
        use s2n_quic_core::random::Generator;

        #[test]
        fn generator_test() {
            let mut generator = super::Generator::default();

            let mut dest_1 = [0; 20];
            let mut dest_2 = [0; 20];

            generator.public_random_fill(&mut dest_1);
            generator.public_random_fill(&mut dest_2);

            assert_ne!(dest_1, dest_2);

            generator.private_random_fill(&mut dest_1);
            generator.private_random_fill(&mut dest_2);

            assert_ne!(dest_1, dest_2);
        }
    }
}
