//! Implements Bitransform for Keccak and its (unknown...) inverse
//! See https://keccak.team and https://en.wikipedia.org/wiki/SHA-3

extern crate rand;

use rand::prelude::*;

#[cfg(test)]
use std::fmt::Write;

use super::{
    Bitransform,
    BitransformDesc
};

// ----------------------------------------------------------------
// Based on Keccak_FIPS202.py implementation by Gilles Van Assche
// With some "optimizations"

// (x, y) = (1, 0); then (x, y) = (y, (2 * x + 3 * y) % 5), 24 times
const KECCAK_RHO_PI_XY: [(usize, usize); 24] = [(0, 2), (2, 1), (1, 2), (2, 3), (3, 3), (3, 0), (0, 1), (1, 3), (3, 1), (1, 4), (4, 4), (4, 0), (0, 3), (3, 4), (4, 3), (3, 2), (2, 2), (2, 0), (0, 4), (4, 2), (2, 4), (4, 1), (1, 1), (1, 0)];

fn keccak_f1600_on_lanes(rounds: usize, lanes: &mut [[u64; 5]; 5]) {
    let mut r: u64 = 1;

    for _ in 0..rounds {
        // θ
        let mut c = [0u64; 5];
        for x in 0..5 {
            let lanes_x = & lanes[x];
            c[x] = lanes_x[0] ^ lanes_x[1] ^ lanes_x[2] ^ lanes_x[3] ^ lanes_x[4];
        }
        let d = [
            c[4] ^ c[1].rotate_left(1),
            c[0] ^ c[2].rotate_left(1),
            c[1] ^ c[3].rotate_left(1),
            c[2] ^ c[4].rotate_left(1),
            c[3] ^ c[0].rotate_left(1)
        ];
        for x in 0..5 {
            let (d_x, lanes_x) = (d[x], &mut lanes[x]);
            for y in 0..5 {
                lanes_x[y] ^= d_x;
            }
        }

        // ρ and π
        let mut current = lanes[1][0];
        for t in 0..24u32 {
            let (x, y) = KECCAK_RHO_PI_XY[t as usize];

            let current_next = lanes[x][y];
            lanes[x][y] = current.rotate_left(((t + 1) * (t + 2)) >> 1);
            current = current_next;
        }

        // χ
        for y in 0..5 {
            let t = [lanes[0][y], lanes[1][y], lanes[2][y], lanes[3][y], lanes[4][y]];
            lanes[0][y] = t[0] ^ ((!t[1]) & t[2]);
            lanes[1][y] = t[1] ^ ((!t[2]) & t[3]);
            lanes[2][y] = t[2] ^ ((!t[3]) & t[4]);
            lanes[3][y] = t[3] ^ ((!t[4]) & t[0]);
            lanes[4][y] = t[4] ^ ((!t[0]) & t[1]);
        }

        // ι
        for j in 0..7 {
            r = ((r << 1) ^ ((r >> 7) * 0x71)) & 0xFF;
            if (r & 2) != 0 {
                lanes[0][0] = lanes[0][0] ^ (1u64 << ((1 << j) - 1));
            }
        }
    }
}

fn keccak_f1600(rounds: usize, state: &mut [u8; 200]) {
    let mut lanes = [[0u64; 5]; 5];
    for x in 0..5 {
        for y in 0..5 {
            lanes[x][y] = (0..8).fold(0u64, |lane, i| lane | (state[((x + 5 * y) << 3) + i] as u64) << (i << 3));
        }
    }
    keccak_f1600_on_lanes(rounds, &mut lanes);
    for x in 0..5 {
        for y in 0..5 {
            let offs = (x + 5 * y) << 3;
            state[offs..(offs + 8)].copy_from_slice(& lanes[x][y].to_le_bytes());
        }
    }
}

fn keccak(length: usize, rounds: usize, msg: &[u8]) -> Vec<u8> {
    let capacity = length << 1;

    let rate = 1600 - capacity;
    let rate_in_bytes = rate >> 3;

    let delimited_suffix: u8 = 0x6; // SHA-3

    let mut block_size: usize = 0;
    let mut input_offset: usize = 0;

    let mut state = [0u8; 200];

    // Absorb all the input blocks
    while input_offset < msg.len() {
        block_size = rate_in_bytes.min(msg.len() - input_offset);
        for i in 0..block_size {
            state[i] ^= msg[input_offset + i];
        }
        input_offset += block_size;
        if block_size == rate_in_bytes {
            keccak_f1600(rounds, &mut state);
            block_size = 0;
        }
    }

    // Do the padding and switch to the squeezing phase
    state[block_size] ^= delimited_suffix;
    if ((delimited_suffix & 0x80) != 0) && (block_size == (rate_in_bytes - 1)) {
        keccak_f1600(rounds, &mut state);
    }
    state[rate_in_bytes - 1] ^= 0x80;
    keccak_f1600(rounds, &mut state);

    // Squeeze out all the output blocks
    let mut output = Vec::<u8>::new();
    let mut output_byte_len = capacity >> 4; // in bytes
    while output_byte_len > 0 {
        block_size = rate_in_bytes.min(output_byte_len);
        output.extend(& state[..block_size]);
        output_byte_len -= block_size;
        if output_byte_len > 0 {
            keccak_f1600(rounds, &mut state);
        }
    }

    output
}

// ----------------------------------------------------------------

#[derive(Debug)]
pub struct Keccak {
    length: usize,
    rounds: usize,
    msg_bytes: usize
}

impl Bitransform for Keccak {
    fn id() -> String {
        String::from("Keccak")
    }

    fn info() -> String {
        format!("IN: msg, OUT: SHA3 of IN, PREFS: length (mult. of 8), rounds (>=1), msg_bytes (>=1)")
    }

    fn new(prefs: &[usize]) -> Result<Box<dyn Bitransform>, String> {
        match prefs.len() {
            0 ..= 2 => Err(format!("Not enough prefs (must be 3)")),
            3 => Ok(Box::new(Self{length: prefs[0], rounds: prefs[1], msg_bytes: prefs[2]})),
            _ => Err(format!("Too many prefs (must be 3)"))
        }
    }

    fn desc(&self) -> BitransformDesc {
        BitransformDesc::new(Self::id(), &[self.length, self.rounds, self.msg_bytes])
    }

    fn inp_size(&self) -> usize {
        self.msg_bytes << 3
    }

    fn outp_size(&self) -> usize {
        self.length
    }

    fn calc_func(&self, args: &[u128]) -> Vec<u128> {
        let mut yb = vec![0u128; self.length];
        for i in 0..0x80 {
            let mut msg = vec![0u8; args.len() >> 3];
            (0..args.len()).for_each(|j|
                msg[j >> 3] |= (((args[j] >> i) & 1) as u8) << (j & 7)
            );
            let y = keccak(self.length, self.rounds, &msg);
            (0..self.length).for_each(|j|
                yb[j] |= (((y[j >> 3] >> (j & 7)) & 1) as u128) << i
            );
        }
        yb
    }

    fn rand_inp(&self, rng: &mut ThreadRng) -> Vec<u128> {
        let mut msgb = Vec::<u128>::with_capacity(self.inp_size());
        for _ in 0..self.inp_size() {
            msgb.push(rng.gen::<u128>())
        }
        msgb
    }

    fn score_inp_outp(&self, inp: &[u128], outp: &[u128]) -> f64 {
        let func = self.calc_func(inp);
        let mut score: i32 = 0;

        for i in 0..self.length {
            score -= (func[i] ^ outp[i]).count_ones() as i32; // number of different bits
        }

        (score as f64) / (0x80 as f64)
    }

}

// ----------------------------------------------------------------

#[derive(Debug)]
pub struct InvKeccak {
    direct: Keccak
}

impl Bitransform for InvKeccak {
    fn id() -> String {
        String::from("InvKeccak")
    }

    fn info() -> String {
        format!("IN: SHA3, OUT: msg whose SHA3 is IN, PREFS: length (mult. of 8), rounds (>=1), msg_bytes (>=1)")
    }

    fn new(prefs: &[usize]) -> Result<Box<dyn Bitransform>, String> {
        match prefs.len() {
            0 ..= 2 => Err(format!("Not enough prefs (must be 3)")),
            3 => Ok(Box::new(Self{direct: Keccak{length: prefs[0], rounds: prefs[1], msg_bytes: prefs[2]}})),
            _ => Err(format!("Too many prefs (must be 3)"))
        }
    }

    fn desc(&self) -> BitransformDesc {
        BitransformDesc::new(Self::id(), &[self.direct.length, self.direct.rounds, self.direct.msg_bytes])
    }

    fn inp_size(&self) -> usize {
        self.direct.length
    }

    fn outp_size(&self) -> usize {
        self.direct.msg_bytes << 3
    }

    fn calc_func(&self, args: &[u128]) -> Vec<u128> {
        self.direct.calc_func(args)
    }

    fn rand_inp(&self, rng: &mut ThreadRng) -> Vec<u128> {
        self.direct.calc_func(& self.direct.rand_inp(rng))
    }

    fn score_inp_outp(&self, inp: &[u128], outp: &[u128]) -> f64 {
        let func = self.direct.calc_func(outp);
        let mut score: i32 = 0;

        for i in 0..self.direct.length {
            score -= (func[i] ^ inp[i]).count_ones() as i32; // number of different bits
        }

        (score as f64) / (0x80 as f64)
    }

}

// ----------------------------------------------------------------
// From https://www.di-mgt.com.au/sha_testvectors.html

#[test]
fn test() {
    let tests = [
        (256, String::from(""), "a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a"),
        (384, String::from(""), "0c63a75b845e4f7d01107d852e4c2485c51a50aaaa94fc61995e71bbee983a2ac3713831264adb47fb6bd1e058d5f004"),
        (512, String::from(""), "a69f73cca23a9ac5c8b567dc185a756e97c982164fe25859e0d1dcc1475c80a615b2123af1f5f94c11e3e9402c3ac558f500199d95b6d3e301758586281dcd26"),

        (256, String::from("abc"), "3a985da74fe225b2045c172d6bd390bd855f086e3e9d525b46bfe24511431532"),
        (384, String::from("abc"), "ec01498288516fc926459f58e2c6ad8df9b473cb0fc08c2596da7cf0e49be4b298d88cea927ac7f539f1edf228376d25"),
        (512, String::from("abc"), "b751850b1a57168a5693cd924b6b096e08f621827444f70d884f5d0240d2712e10e116e9192af3c91a7ec57647e3934057340b4cf408d5a56592f8274eec53f0"),

        (256, String::from("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"), "41c0dba2a9d6240849100376a8235e2c82e1b9998a999e21db32dd97496d3376"),
        (384, String::from("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"), "991c665755eb3a4b6bbdfb75c78a492e8c56a22c5c4d7e429bfdbc32b9d4ad5aa04a1f076e62fea19eef51acd0657c22"),
        (512, String::from("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq"), "04a371e84ecfb5b8b77cb48610fca8182dd457ce6f326a0fd3d7ec2f1e91636dee691fbe0c985302ba1b0d8dc78c086346b533b49c030d99a27daf1139d6e75e"),

        (256, String::from("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu"), "916f6061fe879741ca6469b43971dfdb28b1a32dc36cb3254e812be27aad1d18"),
        (384, String::from("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu"), "79407d3b5916b59c3e30b09822974791c313fb9ecc849e406f23592d04f625dc8c709b98b43b3852b337216179aa7fc7"),
        (512, String::from("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu"), "afebb2ef542e6579c50cad06d2e578f9f8dd6881d7dc824d26360feebf18a4fa73e3261122948efcfd492e74e82e2189ed0fb440d187f382270cb455f21dd185")
    ];
    for (length, msg_str, true_hash_str) in tests.iter() {
        let hash = keccak(*length, 24, msg_str.as_bytes());
        let mut hash_str = String::from("");
        hash.into_iter().for_each(|b| {
            let _ = write!(hash_str, "{:02x}", b);
        });
        assert_eq!(hash_str.as_str(), *true_hash_str);
    }
}
