//! The "joky" cipher: a deliberately non-cryptographic stream cipher.
//!
//! Two stages, picked for very different reasons:
//!
//! 1. A keyed 256-byte substitution table (an RC4-style key schedule) turns the
//!    password and a per-connection nonce into 256 bits of seed material. This
//!    runs exactly once per connection, so it can afford to be byte-at-a-time.
//!    It is the "precomputed table" stage.
//!
//! 2. A `xoshiro256++` generator expands that seed into the actual keystream.
//!    This is the hot path: one 64-bit word per four cheap integer ops, and the
//!    XOR loop vectorizes. It is what makes encryption fast.
//!
//! None of this is secure. There is no authentication of the payload, the
//! keystream is a plain PRNG, and reusing a nonce reuses the keystream. It is an
//! educational toy, not a tunnel you should trust with anything real.

/// RC4-style keyed permutation. Used both as a key-derivation function and to
/// produce short authentication tags during the handshake.
pub struct KeyedTable {
    s: [u8; 256],
    i: u8,
    j: u8,
}

impl KeyedTable {
    pub fn new(key: &[u8]) -> Self {
        let mut s = [0u8; 256];
        for (idx, slot) in s.iter_mut().enumerate() {
            *slot = idx as u8;
        }
        let mut j: u8 = 0;
        if !key.is_empty() {
            for i in 0..256 {
                j = j.wrapping_add(s[i]).wrapping_add(key[i % key.len()]);
                s.swap(i, j as usize);
            }
        }
        Self { s, i: 0, j: 0 }
    }

    #[inline]
    fn next_byte(&mut self) -> u8 {
        self.i = self.i.wrapping_add(1);
        self.j = self.j.wrapping_add(self.s[self.i as usize]);
        self.s.swap(self.i as usize, self.j as usize);
        let idx = self.s[self.i as usize].wrapping_add(self.s[self.j as usize]);
        self.s[idx as usize]
    }

    pub fn fill(&mut self, out: &mut [u8]) {
        for b in out.iter_mut() {
            *b = self.next_byte();
        }
    }
}

/// xoshiro256++ — David Blackman and Sebastiano Vigna's PRNG. Fast, long period
/// (2^256 - 1), good statistical quality. Not cryptographically secure, which is
/// exactly the point here.
struct Xoshiro256pp {
    s: [u64; 4],
}

impl Xoshiro256pp {
    fn from_seed(seed: [u8; 32]) -> Self {
        let mut s = [0u64; 4];
        for (word, chunk) in s.iter_mut().zip(seed.chunks_exact(8)) {
            let mut bytes = [0u8; 8];
            bytes.copy_from_slice(chunk);
            *word = u64::from_le_bytes(bytes);
        }
        if s == [0, 0, 0, 0] {
            s = [0x9e3779b97f4a7c15, 1, 2, 3];
        }
        Self { s }
    }

    #[inline(always)]
    fn next(&mut self) -> u64 {
        let result = self.s[0]
            .wrapping_add(self.s[3])
            .rotate_left(23)
            .wrapping_add(self.s[0]);
        let t = self.s[1] << 17;
        self.s[2] ^= self.s[0];
        self.s[3] ^= self.s[1];
        self.s[1] ^= self.s[2];
        self.s[0] ^= self.s[3];
        self.s[2] ^= t;
        self.s[3] = self.s[3].rotate_left(45);
        result
    }
}

/// A one-directional keystream. Encryption and decryption are the same XOR
/// operation, so a relay can apply this to bytes regardless of which way they
/// happen to be travelling — as long as both endpoints keep their per-direction
/// streams in lockstep with the byte sequence.
pub struct StreamCipher {
    rng: Xoshiro256pp,
    leftover: [u8; 8],
    have: usize,
}

impl StreamCipher {
    pub fn from_seed(seed: [u8; 32]) -> Self {
        Self {
            rng: Xoshiro256pp::from_seed(seed),
            leftover: [0u8; 8],
            have: 0,
        }
    }

    /// XOR `data` in place with the next stretch of keystream. Keystream byte k
    /// always lands on data byte k regardless of how the caller chunks its
    /// reads, which is what lets the two endpoints stay aligned.
    pub fn apply(&mut self, data: &mut [u8]) {
        let mut idx = 0;

        while self.have > 0 && idx < data.len() {
            data[idx] ^= self.leftover[8 - self.have];
            self.have -= 1;
            idx += 1;
        }

        while idx + 8 <= data.len() {
            let k = self.rng.next().to_le_bytes();
            for t in 0..8 {
                data[idx + t] ^= k[t];
            }
            idx += 8;
        }

        let rem = data.len() - idx;
        if rem > 0 {
            let k = self.rng.next().to_le_bytes();
            for t in 0..rem {
                data[idx + t] ^= k[t];
            }
            self.leftover = k;
            self.have = 8 - rem;
        }
    }
}

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

    fn seed_for(key: &[u8]) -> [u8; 32] {
        let mut seed = [0u8; 32];
        KeyedTable::new(key).fill(&mut seed);
        seed
    }

    #[test]
    fn roundtrips_in_one_pass() {
        let seed = seed_for(b"hunter2");
        let mut enc = StreamCipher::from_seed(seed);
        let mut dec = StreamCipher::from_seed(seed);

        let plain = b"the quick brown fox jumps over the lazy dog".to_vec();
        let mut buf = plain.clone();
        enc.apply(&mut buf);
        assert_ne!(buf, plain, "ciphertext should differ from plaintext");
        dec.apply(&mut buf);
        assert_eq!(buf, plain);
    }

    #[test]
    fn chunking_does_not_change_keystream() {
        let seed = seed_for(b"correct horse battery staple");
        let data: Vec<u8> = (0..5000u32).map(|x| (x * 7) as u8).collect();

        let mut whole = data.clone();
        StreamCipher::from_seed(seed).apply(&mut whole);

        let mut piecewise = data.clone();
        let mut c = StreamCipher::from_seed(seed);
        let mut off = 0;
        for step in [1usize, 7, 8, 9, 100, 333, 1, 2, 3] {
            let end = (off + step).min(piecewise.len());
            c.apply(&mut piecewise[off..end]);
            off = end;
        }
        c.apply(&mut piecewise[off..]);

        assert_eq!(
            whole, piecewise,
            "stream position must not depend on chunking"
        );
    }

    #[test]
    fn different_passwords_diverge() {
        let mut a = [0u8; 64];
        let mut b = [0u8; 64];
        StreamCipher::from_seed(seed_for(b"password-a")).apply(&mut a);
        StreamCipher::from_seed(seed_for(b"password-b")).apply(&mut b);
        assert_ne!(a, b);
    }

    #[test]
    fn no_long_zero_runs_in_keystream() {
        let mut buf = [0u8; 4096];
        StreamCipher::from_seed(seed_for(b"pw")).apply(&mut buf);
        let mut run = 0;
        let mut worst = 0;
        for &byte in buf.iter() {
            if byte == 0 {
                run += 1;
                worst = worst.max(run);
            } else {
                run = 0;
            }
        }
        assert!(
            worst < 5,
            "keystream has a suspicious run of zero bytes: {worst}"
        );
    }
}
