/*
 * SPDX-License-Identifier: MIT
 * Copyright (c) 2023 - 2026. The DeepCausality Authors and Contributors. All Rights Reserved.
 */

#[cfg(feature = "aead-random")]
use crate::{Rng, RngCore};
#[cfg(feature = "aead-random")]
use chacha20poly1305::ChaCha20Poly1305;
#[cfg(feature = "aead-random")]
use chacha20poly1305::aead::{AeadInPlace, KeyInit};
#[cfg(feature = "aead-random")]
use zeroize::{Zeroize, ZeroizeOnDrop};

#[cfg(feature = "aead-random")]
#[derive(Debug, Zeroize, ZeroizeOnDrop)]
pub struct ChaCha20Rng {
    // Stores the key for the ChaCha20Poly1305 cipher.
    // The key is zeroized when the RNG is dropped.
    key: [u8; 32],
    // Stores the nonce for the ChaCha20Poly1305 cipher.
    // The nonce is zeroized when the RNG is dropped.
    // The nonce serves as a counter.
    nonce: [u8; 12],
    // Stores the random bytes generated by the cipher.
    // The buffer is zeroized when the RNG is dropped.
    buffer: Vec<u8>,
    // Tracks the current index in the buffer.
    #[zeroize(skip)] // Do not zeroize the index as it comprises only public metadata
    index: usize,
}

#[cfg(feature = "aead-random")]
impl Default for ChaCha20Rng {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "aead-random")]
impl ChaCha20Rng {
    // ... (new() implementation remains unchanged) ...
    /// Creates a new ChaCha20Rng seeded from the system's secure RNG mixed with
    /// high-resolution software entropy (Time, Thread ID, Memory Layout).
    /// This provides key material even if the OS RNG is compromised.
    pub fn new() -> Self {
        use getrandom::u64 as getrandom_u64;
        use std::collections::hash_map::RandomState;
        use std::hash::{BuildHasher, Hasher};
        use std::thread;
        use std::time::{Instant, SystemTime};

        // 1. Get Hardware Entropy (OS RNG)
        let mut hardware_seed = [0u8; 32];
        for chunk in hardware_seed.chunks_mut(8) {
            let val = getrandom_u64().expect("Failed to get secure seed for ChaCha20Rng");
            chunk.copy_from_slice(&val.to_ne_bytes()[..chunk.len()]);
        }
        // ... (software entropy gathering omitted for brevity, logic is same) ...
        // 2. Gather Software Entropy
        let mut hasher = RandomState::new().build_hasher();
        hasher.write_u64(
            SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos() as u64,
        );
        hasher.write_u64(Instant::now().elapsed().as_nanos() as u64);
        use std::hash::Hash;
        thread::current().id().hash(&mut hasher);
        let stack_var = 0;
        hasher.write_usize(&stack_var as *const i32 as usize);
        #[cfg(target_arch = "x86_64")]
        {
            unsafe {
                hasher.write_u64(core::arch::x86_64::_rdtsc());
            }
        }
        #[cfg(target_arch = "aarch64")]
        {
            unsafe {
                let mut c: u64;
                std::arch::asm!("mrs {}, cntvct_el0", out(reg) c);
                hasher.write_u64(c);
            }
        }
        let heap_var = Box::new(0u8);
        hasher.write_usize(heap_var.as_ref() as *const u8 as usize);
        let software_entropy = hasher.finish();

        for (i, chunk) in hardware_seed.chunks_mut(8).enumerate() {
            let mut val = u64::from_ne_bytes(chunk.try_into().unwrap());
            val ^= software_entropy.rotate_left(i as u32 * 13);
            chunk.copy_from_slice(&val.to_ne_bytes());
        }

        Self::from_seed(hardware_seed)
    }

    pub fn from_seed(seed: [u8; 32]) -> Self {
        ChaCha20Rng {
            key: seed,
            nonce: [0u8; 12], // Start counter at 0
            // Pre-allocate buffer with enough space for 1024 bytes of data + 16 bytes of tag
            buffer: Vec::with_capacity(1024 + 16),
            index: 0,
        }
    }

    /// Reseeds the RNG with a new key.
    /// The old key is overwritten and zeroized.
    pub fn reseed(&mut self, seed: [u8; 32]) {
        self.key = seed;
        self.nonce = [0u8; 12]; // Reset counter

        // Explicitly zeroize the old buffer to prevent state leakage
        self.buffer.zeroize();
        self.buffer.clear();

        self.index = 0;
    }

    /// Fills the internal buffer with fresh random bytes.
    fn refill_buffer(&mut self) {
        // Create the cipher instance
        let cipher =
            ChaCha20Poly1305::new_from_slice(&self.key).expect("Failed to create ChaCha20Poly1305");

        // Prepare buffer:
        // 1. Clear existing content (keeping capacity)
        self.buffer.clear();
        // 2. Fill with 1024 zeros (encrypting zeros generates keystream)
        //    This uses the pre-allocated capacity, so no new allocation if < 1040 bytes.
        self.buffer.resize(1024, 0u8);

        // Encrypt in-place.
        // The tag (16 bytes) will be appended to the buffer.
        // We reuse the buffer to avoid allocating a new Vec for ciphertext.
        match cipher.encrypt_in_place(&self.nonce.into(), &[], &mut self.buffer) {
            Ok(_) => {
                // The buffer now contains 1024 bytes of keystream + 16 bytes of tag.
                // We truncate the tag as we only need the keystream.
                self.buffer.truncate(1024);

                self.index = 0;

                // Increment nonce (counter)
                self.increment_nonce();
            }
            Err(_) => {
                // Should not happen with valid key/nonce/payload
                panic!("ChaCha20Poly1305 encryption failed during RNG generation");
            }
        }
    }

    fn increment_nonce(&mut self) {
        // Treat the 12-byte nonce as a 96-bit little-endian integer.
        // We split it into:
        // - lower 64 bits (bytes 0..8)
        // - upper 32 bits (bytes 8..12)
        // This ensures a correct full 96-bit increment propagation.

        // Safe conversion using standard library functions (handles alignment and endianness)
        let mut lower = u64::from_le_bytes(self.nonce[0..8].try_into().unwrap());
        let mut upper = u32::from_le_bytes(self.nonce[8..12].try_into().unwrap());

        // Increment lower 64 bits
        let (new_lower, overflow) = lower.overflowing_add(1);
        lower = new_lower;

        // If lower 64 bits overflowed, carry into upper 32 bits
        if overflow {
            let (new_upper, _full_overflow) = upper.overflowing_add(1);
            upper = new_upper;
            // Note: _full_overflow would imply the entire 96-bit space has been exhausted.
            // We allow wrapping behavior (standard for PRNGs).
        }

        // Write back
        self.nonce[0..8].copy_from_slice(&lower.to_le_bytes());
        self.nonce[8..12].copy_from_slice(&upper.to_le_bytes());
    }
}

#[cfg(feature = "aead-random")]
impl RngCore for ChaCha20Rng {
    fn next_u32(&mut self) -> u32 {
        let mut buf = [0u8; 4];
        self.fill_bytes(&mut buf);
        u32::from_le_bytes(buf)
    }

    fn next_u64(&mut self) -> u64 {
        let mut buf = [0u8; 8];
        self.fill_bytes(&mut buf);
        u64::from_le_bytes(buf)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        let mut dest_written = 0;
        while dest_written < dest.len() {
            if self.index >= self.buffer.len() {
                self.refill_buffer();
            }

            let available = self.buffer.len() - self.index;
            let needed = dest.len() - dest_written;
            let to_copy = std::cmp::min(available, needed);

            dest[dest_written..dest_written + to_copy]
                .copy_from_slice(&self.buffer[self.index..self.index + to_copy]);

            self.index += to_copy;
            dest_written += to_copy;
        }
    }
}

#[cfg(feature = "aead-random")]
impl Rng for ChaCha20Rng {}

#[cfg(all(test, feature = "aead-random"))]
mod tests {
    use super::*;

    #[test]
    fn test_initial_increment() {
        let mut rng = ChaCha20Rng::from_seed([0u8; 32]);
        rng.increment_nonce();
        // Lower 64 bits should be 1
        assert_eq!(rng.nonce[0], 1);
        assert_eq!(rng.nonce[1..8], [0u8; 7]);
        // Upper 32 bits should be 0
        assert_eq!(rng.nonce[8..], [0u8; 4]);
    }

    #[test]
    fn test_overflow_64bit() {
        let mut rng = ChaCha20Rng::from_seed([0u8; 32]);

        // Set nonce to max u64
        let max_u64_bytes = u64::MAX.to_le_bytes();
        rng.nonce[0..8].copy_from_slice(&max_u64_bytes);

        // Increment should overflow lower 64 and carry to upper 32
        rng.increment_nonce();

        // Lower 64 should wrap to 0
        assert_eq!(rng.nonce[0..8], [0u8; 8]);
        // Upper 32 should increment to 1
        assert_eq!(rng.nonce[8], 1);
        assert_eq!(rng.nonce[9..], [0u8; 3]);
    }

    #[test]
    fn test_full_96bit_overflow() {
        let mut rng = ChaCha20Rng::from_seed([0u8; 32]);

        // Set nonce to max 96-bit value
        rng.nonce = [0xFF; 12];

        // Increment should wrap everything to 0
        rng.increment_nonce();

        assert_eq!(rng.nonce, [0u8; 12]);
    }
}
