//! SHAKE128 / SHAKE256 PRNGs — NIST SHA-3 extendable-output functions used
//! as CSPRNGs.
//!
//! SHAKE128 and SHAKE256 are XOFs defined in FIPS 202. They produce an
//! arbitrary-length output from a variable-length input. To use them as
//! PRNGs we feed the seed as input and read the XOF output as a keystream.
//!
//! Per `docs/MANIFEST.md` §6.1, SHAKE256 is the recommended PRNG for the
//! "Air Gap" profile (offline classified destruction).

use crate::{caps, PrngError, PrngProvider, PrngState};
use sha3::{Shake128, Shake256};
use sha3::digest::{Update, ExtendableOutput, XofReader};

pub struct Shake128Prng;
pub struct Shake256Prng;

pub struct Shake128State {
    reader: <Shake128 as ExtendableOutput>::Reader,
}

pub struct Shake256State {
    reader: <Shake256 as ExtendableOutput>::Reader,
}

impl PrngState for Shake128State {
    fn generate(&mut self, out: &mut [u8]) -> Result<(), PrngError> {
        XofReader::read(&mut self.reader, out);
        Ok(())
    }
}

impl PrngState for Shake256State {
    fn generate(&mut self, out: &mut [u8]) -> Result<(), PrngError> {
        XofReader::read(&mut self.reader, out);
        Ok(())
    }
}

impl PrngProvider for Shake128Prng {
    fn name(&self) -> &'static str { "SHAKE128 (CSPRNG)" }
    fn capabilities(&self) -> &'static [&'static str] {
        &[caps::CSPRNG, caps::XOF, caps::FIPS_ELIGIBLE]
    }
    fn min_seed_bytes(&self) -> usize { 16 }
    fn state_size(&self) -> usize { 64 }

    fn init(&self, seed: &[u8]) -> Result<Box<dyn PrngState>, PrngError> {
        if seed.len() < 16 {
            return Err(PrngError::SeedTooShort { got: seed.len(), need: 16 });
        }
        let mut h = Shake128::default();
        Update::update(&mut h, b"scuttle-shake128-prng-v1");
        Update::update(&mut h, seed);
        Ok(Box::new(Shake128State { reader: h.finalize_xof() }))
    }

    fn self_test(&self) -> Result<(), PrngError> {
        // Determinism: same seed → same output.
        let seed = [0u8; 32];
        let mut a = self.init(&seed)?;
        let mut b = self.init(&seed)?;
        let mut oa = [0u8; 64];
        let mut ob = [0u8; 64];
        a.generate(&mut oa)?;
        b.generate(&mut ob)?;
        if oa != ob {
            return Err(PrngError::SelfTestFailed {
                provider: "SHAKE128 (CSPRNG)", detail: "determinism check failed".into(),
            });
        }
        // Different seed → different output.
        let mut c = self.init(&[0xffu8; 32])?;
        let mut oc = [0u8; 64];
        c.generate(&mut oc)?;
        if oa == oc {
            return Err(PrngError::SelfTestFailed {
                provider: "SHAKE128 (CSPRNG)",
                detail: "different seeds produced identical output".into(),
            });
        }
        Ok(())
    }
}

impl PrngProvider for Shake256Prng {
    fn name(&self) -> &'static str { "SHAKE256 (CSPRNG)" }
    fn capabilities(&self) -> &'static [&'static str] {
        &[caps::CSPRNG, caps::XOF, caps::FIPS_ELIGIBLE]
    }
    fn min_seed_bytes(&self) -> usize { 32 }
    fn state_size(&self) -> usize { 64 }

    fn init(&self, seed: &[u8]) -> Result<Box<dyn PrngState>, PrngError> {
        if seed.len() < 32 {
            return Err(PrngError::SeedTooShort { got: seed.len(), need: 32 });
        }
        let mut h = Shake256::default();
        Update::update(&mut h, b"scuttle-shake256-prng-v1");
        Update::update(&mut h, seed);
        Ok(Box::new(Shake256State { reader: h.finalize_xof() }))
    }

    fn self_test(&self) -> Result<(), PrngError> {
        let seed = [0u8; 32];
        let mut a = self.init(&seed)?;
        let mut b = self.init(&seed)?;
        let mut oa = [0u8; 64];
        let mut ob = [0u8; 64];
        a.generate(&mut oa)?;
        b.generate(&mut ob)?;
        if oa != ob {
            return Err(PrngError::SelfTestFailed {
                provider: "SHAKE256 (CSPRNG)", detail: "determinism check failed".into(),
            });
        }
        let mut c = self.init(&[0xffu8; 32])?;
        let mut oc = [0u8; 64];
        c.generate(&mut oc)?;
        if oa == oc {
            return Err(PrngError::SelfTestFailed {
                provider: "SHAKE256 (CSPRNG)",
                detail: "different seeds produced identical output".into(),
            });
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn shake128_self_test() { Shake128Prng.self_test().unwrap(); }

    #[test]
    fn shake256_self_test() { Shake256Prng.self_test().unwrap(); }
}