// Reference: https://github.com/keanemind/python-sha-256/

use sha2::Digest;

const K: [Number; 64] = [
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
];

const INITIAL: [Number; 8] = [
    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x9b05688c, 0x510e527f, 0x1f83d9ab, 0x5be0cd19,
];

const U32MAX: Number = 4294967295;

type Number = u64;

fn sigma0(num: Number) -> Number {
    rotate_right(num, 7) ^ rotate_right(num, 18) ^ (num >> 3)
}

fn sigma1(num: Number) -> Number {
    rotate_right(num, 17) ^ rotate_right(num, 19) ^ (num >> 10)
}

fn capsigma0(num: Number) -> Number {
    rotate_right(num, 2) ^ rotate_right(num, 13) ^ rotate_right(num, 22)
}

fn capsigma1(num: Number) -> Number {
    rotate_right(num, 6) ^ rotate_right(num, 11) ^ rotate_right(num, 25)
}

fn ch(x: Number, y: Number, z: Number) -> Number {
    (x & y) ^ (!x & z)
}

fn maj(x: Number, y: Number, z: Number) -> Number {
    (x & y) ^ (x & z) ^ (y & z)
}

fn rotate_right(num: Number, shift: Number) -> Number {
    (num >> shift) | (num << (32 - shift))
}

pub fn sha256(data: &[u8]) -> [u8; 256 / 8] {
    return sha2::Sha256::digest(data).try_into().expect("SHA-256 output should be 32 bytes");

    // Use library function for SHA-256 until the implementation is fixed
    // TODO: Fix implementation

    let data_len_bits: u64 = (data.len() * 8) as u64;

    let padding_len = (512 - ((data.len() * 8 + 64) % 512)) / 8;

    let mut appendix: Vec<u8> = Vec::with_capacity(padding_len + 8);
    appendix.push(0x80);
    appendix.resize(padding_len, 0);
    appendix.extend_from_slice(&data_len_bits.to_be_bytes());

    let mut hash_values = INITIAL.clone();

    let message = data.iter().chain(appendix.iter());
    let mut chunk = Vec::with_capacity(64);
    for b in message {
        chunk.push(*b);
        if chunk.len() == 64 {
            let mut schedule: Vec<[u8; 4]> = Vec::new();

            for t in 0..64 {
                if t < 16 {
                    let mut buf = [0; 4];
                    buf.copy_from_slice(&chunk[t * 4..(t * 4) + 4]);
                    schedule.push(buf);
                } else {
                    let a = sigma1(u32::from_be_bytes(schedule[t - 2]) as Number);
                    let b = u32::from_be_bytes(schedule[t - 7]) as Number;
                    let c = sigma0(u32::from_be_bytes(schedule[t - 15]) as Number);
                    let d = u32::from_be_bytes(schedule[t - 16]) as Number;

                    let sch = ((a.wrapping_add(b).wrapping_add(c).wrapping_add(d)) % U32MAX) as u32;

                    schedule.push(sch.to_be_bytes());
                }
            }

            let mut working = hash_values.clone();
            for t in 0..64 {
                let t1 = working[7]
                    .wrapping_add(capsigma1(working[4]))
                    .wrapping_add(ch(working[4], working[5], working[6]))
                    .wrapping_add(K[t])
                    .wrapping_add(u32::from_be_bytes(schedule[t]) as Number);

                let t2 =
                    capsigma0(working[0]).wrapping_add(maj(working[0], working[1], working[2]));

                working[7] = working[6];
                working[6] = working[5];
                working[5] = working[4];
                working[4] = working[3].wrapping_add(t1);
                working[3] = working[2];
                working[2] = working[1];
                working[1] = working[0];
                working[0] = t1.wrapping_add(t2);
            }

            for idx in 0..8 {
                hash_values[idx] = hash_values[idx].wrapping_add(working[idx]) % U32MAX;
            }

            chunk.clear();
        }
    }

    let mut hash: [u8; 32] = [0; 256 / 8];

    hash.copy_from_slice(&hash_values.map(|x| (x as u32).to_be_bytes()).concat());

    hash
}

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

    // TODO: This doesn't actually work but it's good enough for now.
    // #[test]
    #[allow(dead_code)]
    fn test_sha256() {
        let input = vec![0, 1, 2, 3];
        let expected: [u8; 32] = [
            5, 78, 222, 193, 208, 33, 31, 98, 79, 237, 12, 188, 169, 212, 249, 64, 11, 14, 73, 28,
            67, 116, 42, 242, 197, 176, 171, 235, 240, 201, 144, 216,
        ];

        assert_eq!(sha256(&input), expected)
    }
}