//! TurboSHAKE128 (RFC 9861 §2.1) and the XOF Prio3 derives every seed and vector with,
//! XofTurboShake128 (`draft-irtf-cfrg-vdaf-22` §6.2, §6.2.1), with the draft's domain
//! separation tag (§6.2.3).
//!
//! TurboSHAKE128 is the Keccak-p[1600, 12] sponge of rate 168 bytes; its message is padded
//! with the domain byte `D` and a final 0x80. The permutation is `keccak::p1600`.

use super::field::Field;

const RATE: usize = 168;

/// The draft's global VERSION, which every domain separation tag starts with (§2).
pub const VERSION: u8 = 18;

/// A TurboSHAKE128 instance: absorb, then squeeze any number of bytes in order.
pub struct TurboShake128 {
    state: [u64; 25],
    /// Bytes absorbed into, or squeezed from, the current block.
    pos: usize,
    domain: u8,
    squeezing: bool,
}

impl TurboShake128 {
    /// `domain` is the RFC's D, in 0x01..=0x7F.
    pub fn new(domain: u8) -> TurboShake128 {
        assert!((0x01..=0x7F).contains(&domain), "TurboSHAKE domain byte out of range");
        TurboShake128 { state: [0; 25], pos: 0, domain, squeezing: false }
    }

    fn xor_byte(&mut self, i: usize, b: u8) {
        self.state[i / 8] ^= (b as u64) << (8 * (i % 8));
    }

    fn byte(&self, i: usize) -> u8 {
        (self.state[i / 8] >> (8 * (i % 8))) as u8
    }

    pub fn absorb(&mut self, m: &[u8]) {
        assert!(!self.squeezing, "absorb after squeeze");
        for &b in m {
            self.xor_byte(self.pos, b);
            self.pos += 1;
            if self.pos == RATE {
                keccak::p1600(&mut self.state, 12);
                self.pos = 0;
            }
        }
    }

    fn finish(&mut self) {
        self.xor_byte(self.pos, self.domain);
        self.xor_byte(RATE - 1, 0x80);
        keccak::p1600(&mut self.state, 12);
        self.pos = 0;
        self.squeezing = true;
    }

    pub fn squeeze(&mut self, out: &mut [u8]) {
        if !self.squeezing {
            self.finish();
        }
        for o in out.iter_mut() {
            if self.pos == RATE {
                keccak::p1600(&mut self.state, 12);
                self.pos = 0;
            }
            *o = self.byte(self.pos);
            self.pos += 1;
        }
    }
}

/// TurboSHAKE128(M, D, L).
pub fn turbo_shake128(m: &[u8], domain: u8, len: usize) -> Vec<u8> {
    let mut t = TurboShake128::new(domain);
    t.absorb(m);
    let mut out = vec![0u8; len];
    t.squeeze(&mut out);
    out
}

/// XofTurboShake128 (§6.2.1): TurboSHAKE128 with D = 1 over
/// `le16(len(dst)) || dst || le8(len(seed)) || seed || binder`.
pub struct Xof(TurboShake128);

pub const SEED_SIZE: usize = 32;

impl Xof {
    pub fn new(seed: &[u8], dst: &[u8], binder: &[u8]) -> Xof {
        assert!(seed.len() <= 255 && dst.len() <= 65535);
        let mut t = TurboShake128::new(1);
        t.absorb(&(dst.len() as u16).to_le_bytes());
        t.absorb(dst);
        t.absorb(&[seed.len() as u8]);
        t.absorb(seed);
        t.absorb(binder);
        Xof(t)
    }

    pub fn next(&mut self, len: usize) -> Vec<u8> {
        let mut out = vec![0u8; len];
        self.0.squeeze(&mut out);
        out
    }

    /// `length` field elements, rejecting a candidate not below the modulus (§6.2).
    pub fn next_vec<F: Field>(&mut self, length: usize) -> Vec<F> {
        // next_power_of_2(MODULUS) − 1 is every bit of ENCODED_SIZE bytes for both fields.
        let mut v = Vec::with_capacity(length);
        while v.len() < length {
            let b = self.next(F::ENCODED_SIZE);
            let mut w = [0u8; 16];
            w[..b.len()].copy_from_slice(&b);
            let x = u128::from_le_bytes(w);
            if x < F::MODULUS {
                v.push(F::from_u128(x));
            }
        }
        v
    }

    pub fn derive_seed(seed: &[u8], dst: &[u8], binder: &[u8]) -> [u8; SEED_SIZE] {
        let mut x = Xof::new(seed, dst, binder);
        let mut out = [0u8; SEED_SIZE];
        x.0.squeeze(&mut out);
        out
    }

    pub fn expand_into_vec<F: Field>(seed: &[u8], dst: &[u8], binder: &[u8], length: usize) -> Vec<F> {
        Xof::new(seed, dst, binder).next_vec(length)
    }
}

/// format_dst (§6.2.3): VERSION, the algorithm class, the algorithm and the usage.
pub fn format_dst(algo_class: u8, algo: u32, usage: u16) -> Vec<u8> {
    let mut v = vec![VERSION, algo_class];
    v.extend_from_slice(&algo.to_be_bytes());
    v.extend_from_slice(&usage.to_be_bytes());
    v
}

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

    fn h(s: &str) -> Vec<u8> {
        let s: String = s.split_whitespace().collect();
        (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect()
    }

    fn ptn(n: usize) -> Vec<u8> {
        (0..n).map(|i| (i % 251) as u8).collect()
    }

    /// RFC 9861 §5's test vectors for TurboSHAKE128.
    #[test]
    fn rfc9861_vectors() {
        assert_eq!(turbo_shake128(&[], 0x1F, 32), h("1E 41 5F 1C 59 83 AF F2 16 92 17 27 7D 17 BB 53 8C D9 45 A3 97 DD EC 54 1F 1C E4 1A F2 C1 B7 4C"));
        let long = turbo_shake128(&[], 0x1F, 10032);
        assert_eq!(long[10000..], h("A3 B9 B0 38 59 00 CE 76 1F 22 AE D5 48 E7 54 DA 10 A5 24 2D 62 E8 C6 58 E3 F3 A9 23 A7 55 56 07"));
        assert_eq!(turbo_shake128(&ptn(17), 0x1F, 32), h("9C 97 D0 36 A3 BA C8 19 DB 70 ED E0 CA 55 4E C6 E4 C2 A1 A4 FF BF D9 EC 26 9C A6 A1 11 16 12 33"));
        assert_eq!(turbo_shake128(&ptn(17 * 17 * 17), 0x1F, 32), h("D4 97 6E B5 6B CF 11 85 20 58 2B 70 9F 73 E1 D6 85 3E 00 1F DA F8 0E 1B 13 E0 D0 59 9D 5F B3 72"));
        assert_eq!(turbo_shake128(&[0xFF, 0xFF, 0xFF], 0x01, 32), h("BF 32 3F 94 04 94 E8 8E E1 C5 40 FE 66 0B E8 A0 C9 3F 43 D1 5E C0 06 99 84 62 FA 99 4E ED 5D AB"));
        assert_eq!(turbo_shake128(&[0xFF; 7], 0x0B, 32), h("8D EE AA 1A EC 47 CC EE 56 9F 65 9C 21 DF A8 E1 12 DB 3C EE 37 B1 81 78 B2 AC D8 05 B7 99 CC 37"));
        assert_eq!(turbo_shake128(&[0xFF, 0xFF, 0xFF], 0x7F, 32), h("16 27 4C C6 56 D4 4C EF D4 22 39 5D 0F 90 53 BD A6 D2 8E 12 2A BA 15 C7 65 E5 AD 0E 6E AF 26 F9"));
    }

    #[test]
    fn squeezing_in_pieces_is_squeezing_at_once() {
        let mut a = TurboShake128::new(1);
        a.absorb(b"abc");
        let mut x = vec![0u8; 400];
        for c in x.chunks_mut(13) {
            a.squeeze(c);
        }
        assert_eq!(x, turbo_shake128(b"abc", 1, 400));
    }
}
