// FIPS 202 / Keccak-f[1600]
// Lane indexing:
//     A[x, y] = state[x + 5*y]
// Lanes are little-endian when converted to/from bytes.

const RHO: [[u32; 5]; 5] = [
    [0, 36, 3, 41, 18],
    [1, 44, 10, 45, 2],
    [62, 6, 43, 15, 61],
    [28, 55, 25, 21, 56],
    [27, 20, 39, 8, 14],
];

const ROUND_CONSTANTS: [u64; 24] = [
    0x0000_0000_0000_0001,
    0x0000_0000_0000_8082,
    0x8000_0000_0000_808a,
    0x8000_0000_8000_8000,
    0x0000_0000_0000_808b,
    0x0000_0000_8000_0001,
    0x8000_0000_8000_8081,
    0x8000_0000_0000_8009,
    0x0000_0000_0000_008a,
    0x0000_0000_0000_0088,
    0x0000_0000_8000_8009,
    0x0000_0000_8000_000a,
    0x0000_0000_8000_808b,
    0x8000_0000_0000_008b,
    0x8000_0000_0000_8089,
    0x8000_0000_0000_8003,
    0x8000_0000_0000_8002,
    0x8000_0000_0000_0080,
    0x0000_0000_0000_800a,
    0x8000_0000_8000_000a,
    0x8000_0000_8000_8081,
    0x8000_0000_0000_8080,
    0x0000_0000_8000_0001,
    0x8000_0000_8000_8008,
];

/// Keccak-f[1600] permutation.
///
/// State layout:
///   state[x + 5*y] == A[x,y]
#[inline]
pub fn keccak_f1600(state: &mut [u64; 25]) {
    let mut c = [0u64; 5];
    let mut d = [0u64; 5];
    let mut b = [0u64; 25];

    for &rc in &ROUND_CONSTANTS {
        // 0 step
        for x in 0..5 {
            c[x] = state[x] ^ state[x + 5] ^ state[x + 10] ^ state[x + 15] ^ state[x + 20];
        }

        for x in 0..5 {
            d[x] = c[(x + 4) % 5] ^ c[(x + 1) % 5].rotate_left(1);
        }

        for x in 0..5 {
            for y in 0..5 {
                state[x + 5 * y] ^= d[x];
            }
        }

        // rho + pi step
        // B[y, 2x + 3y] = ROT(A[x,y], r[x,y])
        for x in 0..5 {
            for y in 0..5 {
                let new_x = y;
                let new_y = (2 * x + 3 * y) % 5;

                b[new_x + 5 * new_y] = state[x + 5 * y].rotate_left(RHO[x][y]);
            }
        }

        // Khi
        for y in 0..5 {
            for x in 0..5 {
                let current = b[x + 5 * y];
                let next = b[(x + 1) % 5 + 5 * y];
                let next2 = b[(x + 2) % 5 + 5 * y];

                state[x + 5 * y] = current ^ ((!next) & next2);
            }
        }

        // iota
        state[0] ^= rc;
    }
}

/// XOR one byte into the sponge state.
///
/// Keccak lanes are little-endian:
/// byte 0 -> bits 0..7 of lane 0
/// byte 1 -> bits 8..15 of lane 0
/// ...
#[inline]
fn xor_byte(state: &mut [u64; 25], offset: usize, value: u8) {
    let lane = offset / 8;
    let shift = (offset % 8) * 8;

    state[lane] ^= (value as u64) << shift;
}

/// Read one byte from the sponge state.
#[inline]
fn get_byte(state: &[u64; 25], offset: usize) -> u8 {
    let lane = offset / 8;
    let shift = (offset % 8) * 8;

    ((state[lane] >> shift) & 0xff) as u8
}

/// Generic Keccak sponge.
///
/// `rate` and `capacity` are specified in bits.
/// `suffix` is the domain-separation / delimited-suffix byte.
///
/// Examples:
///   SHA-3: suffix = 0x06
///   SHAKE: suffix = 0x1f
pub fn keccak(
    rate: usize,
    capacity: usize,
    input: &[u8],
    suffix: u8,
    output_len: usize,
) -> Vec<u8> {
    assert_eq!(rate + capacity, 1600, "rate + capacity must equal 1600");

    assert_eq!(rate % 8, 0, "rate must be a multiple of 8 bits");

    let rate_bytes = rate / 8;

    let mut state = [0u64; 25];
    let mut input_offset = 0;

    // Absorb
    // Full blocks
    while input_offset + rate_bytes <= input.len() {
        for i in 0..rate_bytes {
            xor_byte(&mut state, i, input[input_offset + i]);
        }

        keccak_f1600(&mut state);

        input_offset += rate_bytes;
    }

    // Remaining partial block
    let block_size = input.len() - input_offset;

    for i in 0..block_size {
        xor_byte(&mut state, i, input[input_offset + i]);
    }

    // Padding
    // delimitedSuffix followed by the final 0x80 bit.

    xor_byte(&mut state, block_size, suffix);

    if (suffix & 0x80) != 0 && block_size == rate_bytes - 1 {
        keccak_f1600(&mut state);
    }

    xor_byte(&mut state, rate_bytes - 1, 0x80);

    keccak_f1600(&mut state);

    // Squeeze
    let mut output = Vec::with_capacity(output_len);

    while output.len() < output_len {
        let block_size = (output_len - output.len()).min(rate_bytes);

        for i in 0..block_size {
            output.push(get_byte(&state, i));
        }

        if output.len() < output_len {
            keccak_f1600(&mut state);
        }
    }

    output
}

// SHA-3

pub fn sha3_224(input: &[u8]) -> [u8; 28] {
    keccak(1152, 448, input, 0x06, 28).try_into().unwrap()
}

pub fn sha3_256(input: &[u8]) -> [u8; 32] {
    keccak(1088, 512, input, 0x06, 32).try_into().unwrap()
}

pub fn sha3_384(input: &[u8]) -> [u8; 48] {
    keccak(832, 768, input, 0x06, 48).try_into().unwrap()
}

pub fn sha3_512(input: &[u8]) -> [u8; 64] {
    keccak(576, 1024, input, 0x06, 64).try_into().unwrap()
}

// SHAKE

pub fn shake128(input: &[u8], output_len: usize) -> Vec<u8> {
    keccak(1344, 256, input, 0x1f, output_len)
}

pub fn shake256(input: &[u8], output_len: usize) -> Vec<u8> {
    keccak(1088, 512, input, 0x1f, output_len)
}

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

    fn hex_decode(s: &str) -> Vec<u8> {
        assert!(
            s.len() % 2 == 0,
            "hex string must contain an even number of characters"
        );

        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    fn assert_hex(actual: &[u8], expected: &str) {
        let expected = hex_decode(expected);

        assert_eq!(
            actual,
            expected.as_slice(),
            "\nexpected: {}\nactual:   {}",
            expected
                .iter()
                .map(|b| format!("{b:02x}"))
                .collect::<String>(),
            actual
                .iter()
                .map(|b| format!("{b:02x}"))
                .collect::<String>(),
        );
    }

    // SHA3-224
    #[test]
    fn sha3_224_nist() {
        // NIST CAVP SHA3_224ShortMsg.rsp
        let vectors = [
            (
                "",
                "6b4e03423667dbb73b6e15454f0eb1abd4597f9a1b078e3f5b5a6bc7",
            ),
            (
                "6b2b92584146a433bee8b947cc1f35b617b73f5b1e0376ac8bdadfe5bfdf2263b205f74dfa53db7a29e5078f5c34a268119736ba390961f6",
                "132cfa7e71fe0991abbd88ef588ac95ac9289b1d775b42033567dd33",
            ),
            (
                "d4f757d1c33b9c0b38b4e93e8e2483ec51b4861299f1d650961457496d86614d42a36e3696bf168fd4663efc26e88cd58d151e1531467b73f69dc9ce4f8d41ce579ce1c91e6760e340e7677abdf4fec1040745aa5144640a39b8c4f884df80753a691653003d634fa5bfce81f94ec3f6",
                "be11259377f09821d9dc358592b6565d8ef2b414dfaa7db5609fb751",
            ),
        ];

        for (message, expected) in vectors {
            let input = hex_decode(message);
            assert_hex(&sha3_224(&input), expected);
        }
    }

    // SHA3-256
    #[test]
    fn sha3_256_nist() {
        // NIST CAVP SHA3_256ShortMsg.rsp
        let vectors = [
            (
                "",
                "a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a",
            ),
            (
                "00ff6c96b7aa3cf27d036cf20af7031434113252574bda9cf9244d85aef2593d3a7a83bff6be904b75164a1766828042bc3f4f090d98a03d",
                "d000eafca34815783bed9b050c6901c97f2e77d4771a0ed724dd8f6ff1448791",
            ),
            (
                "8d93627c0b7cbf61a7fe70e78c2c8ed23b1344b4cfed31bd85980dd37b4690e5b8758f7d6d2269957a39a1ac3451cc196696ae9e9606a04089e13456095a1ce1e593481b3ac84f53f1cb10f789b099f316c948398ad52fa13474bdf486de9b431bd5d57ef9d83a42139a05f112b2bd08",
                "344ec86642eabb206b2fd930e4c5dde78aa878577d6c271cb0069d4999495652",
            ),
        ];

        for (message, expected) in vectors {
            let input = hex_decode(message);
            assert_hex(&sha3_256(&input), expected);
        }
    }

    // SHA3-384
    #[test]
    fn sha3_384_nist() {
        // NIST CAVP SHA3_384ShortMsg.rsp
        let vectors = [
            (
                "",
                "0c63a75b845e4f7d01107d852e4c2485c51a50aaaa94fc61995e71bbee983a2ac3713831264adb47fb6bd1e058d5f004",
            ),
            (
                "5415c2596aa7d21e855be98491bd702357c19f21f46294f98a8aa37b3532ee1541ca35509adbef9d83eb99528ba14ef0bd2998a718da861c3f16fe6971",
                "8f9fd7d879d6b51ee843e1fbcd40bb67449ae744db9f673e3452f028cb0189d9cb0fef7bdb5c760d63fea0e3ba3dd8d1",
            ),
            (
                "00ce225eaea24843406fa42cc8450e66f76ac9f549b8591f7d40942f4833fc734a034c8741c551d57ddafb5d94ceb4b25680f045038306e6bcc53e88386e2b45b80b3ba23dec8c13f8ca01c202ae968c4d0df04cdb38395d2df42a5aff646928",
                "81d6e0d96575a9b8ca083ee9ec2ead57ddf72b97d7709086a2f4a749d3f61d16423463487562c7f09aba1b26e8cae47b",
            ),
        ];

        for (message, expected) in vectors {
            let input = hex_decode(message);
            assert_hex(&sha3_384(&input), expected);
        }
    }

    // SHA3-512
    #[test]
    fn sha3_512_nist() {
        // NIST CAVP SHA3_512ShortMsg.rsp
        let vectors = [
            (
                "",
                "a69f73cca23a9ac5c8b567dc185a756e97c982164fe25859e0d1dcc1475c80a615b2123af1f5f94c11e3e9402c3ac558f500199d95b6d3e301758586281dcd26",
            ),
            (
                "302fa84fdaa82081b1192b847b81ddea10a9f05a0f04138fd1da84a39ba5e18e18bc3cea062e6df92ff1ace89b3c5f55043130108abf631e",
                "8c8eaae9a445643a37df34cfa6a7f09deccab2a222c421d2fc574bbc5641e504354391e81eb5130280b1226812556d474e951bb78dbdd9b77d19f647e2e7d7be",
            ),
            (
                "0ce9f8c3a990c268f34efd9befdb0f7c4ef8466cfdb01171f8de70dc5fefa92acbe93d29e2ac1a5c2979129f1ab08c0e77de7924ddf68a209cdfa0adc62f85c18637d9c6b33f4ff8",
                "b018a20fcf831dde290e4fb18c56342efe138472cbe142da6b77eea4fce52588c04c808eb32912faa345245a850346faec46c3a16d39bd2e1ddb1816bc57d2da",
            ),
        ];

        for (message, expected) in vectors {
            let input = hex_decode(message);
            assert_hex(&sha3_512(&input), expected);
        }
    }

    // SHA-3 standard "abc" vectors
    #[test]
    fn sha3_abc() {
        assert_hex(
            &sha3_224(b"abc"),
            "e642824c3f8cf24ad09234ee7d3c766fc9a3a5168d0c94ad73b46fdf",
        );

        assert_hex(
            &sha3_256(b"abc"),
            "3a985da74fe225b2045c172d6bd390bd855f086e3e9d525b46bfe24511431532",
        );

        assert_hex(
            &sha3_384(b"abc"),
            "ec01498288516fc926459f58e2c6ad8df9b473cb0fc08c2596da7cf0e49be4b298d88cea927ac7f539f1edf228376d25",
        );

        assert_hex(
            &sha3_512(b"abc"),
            "b751850b1a57168a5693cd924b6b096e08f621827444f70d884f5d0240d2712e10e116e9192af3c91a7ec57647e3934057340b4cf408d5a56592f8274eec53f0",
        );
    }

    // SHAKE128
    #[test]
    fn shake128_nist() {
        // FIPS 202 / standard known-answer vectors.
        assert_hex(
            &shake128(b"", 32),
            "7f9c2ba4e88f827d616045507605853ed73b8093f6efbc88eb1a6eacfa66ef26",
        );

        assert_hex(
            &shake128(b"abc", 32),
            "5881092dd818bf5cf8a3ddb793fbcba74097d5c526a6d35f97b83351940f2cc8",
        );
    }

    // SHAKE256
    #[test]
    fn shake256_nist() {
        assert_hex(
            &shake256(b"", 64),
            "46b9dd2b0ba88d13233b3feb743eeb243fcd52ea62b81b82b50c27646ed5762f\
d75dc4ddd8c0f200cb05019d67b592f6fc821c49479ab48640292eacb3b7c4be",
        );

        assert_hex(
            &shake256(b"abc", 64),
            "483366601360a8771c6863080cc4114d\
8db44530f8f1e1ee4f94ea37e78b5739\
d5a15bef186a5386c75744c0527e1faa\
9f8726e462a12a4feb06bd8801e751e4",
        );
    }

    #[test]
    fn shake_prefix_property() {
        let short = shake128(b"abc", 16);
        let long = shake128(b"abc", 64);

        assert_eq!(&long[..16], &short[..]);

        let short = shake256(b"abc", 32);
        let long = shake256(b"abc", 64);

        assert_eq!(&long[..32], &short[..]);
    }

    #[test]
    fn shake_multiple_squeeze_blocks() {
        let output128 = shake128(b"abc", 200);
        let output256 = shake256(b"abc", 200);

        assert_eq!(output128.len(), 200);
        assert_eq!(output256.len(), 200);

        // Verify the first bytes against the standard vectors.
        assert_hex(
            &output128[..32],
            "5881092dd818bf5cf8a3ddb793fbcba74097d5c526a6d35f97b83351940f2cc8",
        );

        assert_hex(
            &output256[..64],
            "483366601360a8771c6863080cc4114d\
8db44530f8f1e1ee4f94ea37e78b5739\
d5a15bef186a5386c75744c0527e1faa\
9f8726e462a12a4feb06bd8801e751e4",
        );
    }
}