// Copyright (c) 2023 Yuki Kishimoto
// Distributed under the MIT software license

use alloc::vec::Vec;

const K: [u32; 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 SHA256_INIT: [u32; 8] = [
    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
];

pub(crate) fn hash(mut data: Vec<u8>) -> [u8; 32] {
    let original_len: usize = data.len();
    let bit_len: u64 = (original_len * 8) as u64;

    // Pre-processing: padding the message
    data.push(0x80); // Append a '1' bit
    while data.len() % 64 != 56 {
        data.push(0x00);
    }

    // Append the original bit length as a 64-bit big-endian integer
    let mut bit_len_bytes: [u8; 8] = [0; 8];
    for i in 0..8 {
        bit_len_bytes[7 - i] = ((bit_len >> (i * 8)) & 0xff) as u8;
    }
    data.extend_from_slice(&bit_len_bytes);

    let mut hash: [u32; 8] = SHA256_INIT;

    for chunk in data.chunks(64) {
        let mut w: [u32; 64] = [0u32; 64];
        for (i, chunk_byte) in chunk.iter().enumerate() {
            w[i / 4] |= u32::from(*chunk_byte) << (24 - (i % 4) * 8);
        }

        for i in 16..64 {
            let s0: u32 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3);
            let s1: u32 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10);
            w[i] = w[i - 16]
                .wrapping_add(s0)
                .wrapping_add(w[i - 7])
                .wrapping_add(s1);
        }

        let mut a: u32 = hash[0];
        let mut b: u32 = hash[1];
        let mut c: u32 = hash[2];
        let mut d: u32 = hash[3];
        let mut e: u32 = hash[4];
        let mut f: u32 = hash[5];
        let mut g: u32 = hash[6];
        let mut h: u32 = hash[7];

        for i in 0..64 {
            let s1: u32 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
            let ch: u32 = (e & f) ^ ((!e) & g);
            let temp1: u32 = h
                .wrapping_add(s1)
                .wrapping_add(ch)
                .wrapping_add(K[i])
                .wrapping_add(w[i]);
            let s0: u32 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
            let maj: u32 = (a & b) ^ (a & c) ^ (b & c);
            let temp2: u32 = s0.wrapping_add(maj);

            h = g;
            g = f;
            f = e;
            e = d.wrapping_add(temp1);
            d = c;
            c = b;
            b = a;
            a = temp1.wrapping_add(temp2);
        }

        hash[0] = hash[0].wrapping_add(a);
        hash[1] = hash[1].wrapping_add(b);
        hash[2] = hash[2].wrapping_add(c);
        hash[3] = hash[3].wrapping_add(d);
        hash[4] = hash[4].wrapping_add(e);
        hash[5] = hash[5].wrapping_add(f);
        hash[6] = hash[6].wrapping_add(g);
        hash[7] = hash[7].wrapping_add(h);
    }

    let mut result: [u8; 32] = [0u8; 32];
    for (i, &h) in hash.iter().enumerate() {
        result[i * 4] = (h >> 24) as u8;
        result[i * 4 + 1] = (h >> 16) as u8;
        result[i * 4 + 2] = (h >> 8) as u8;
        result[i * 4 + 3] = h as u8;
    }

    result
}

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

    const HASHES: [(&str, [u8; 32]); 10] = [
        ("Bitcoin: A Peer-to-Peer Electronic Cash System", [0xef, 0xb5, 0xc6, 0x72, 0x9d, 0x8c, 0xe3, 0xe0,
            0x3f, 0xd0, 0x3a, 0xec, 0x34, 0x05, 0x40, 0xb2,
            0x4a, 0x78, 0x84, 0x54, 0xd4, 0x5e, 0x71, 0x70,
            0x89, 0xb1, 0xe5, 0x92, 0x43, 0xe1, 0x6f, 0x43]
        ),
        ("", [
            0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14, 0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
            0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c, 0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55
        ]
        ),
        ("Hello, world!", [
            0x31, 0x5f, 0x5b, 0xdb, 0x76, 0xd0, 0x78, 0xc4, 0x3b, 0x8a, 0xc0, 0x06, 0x4e, 0x4a, 0x01, 0x64,
            0x61, 0x2b, 0x1f, 0xce, 0x77, 0xc8, 0x69, 0x34, 0x5b, 0xfc, 0x94, 0xc7, 0x58, 0x94, 0xed, 0xd3
        ]
        ),
        ("Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.", [
            0x18, 0x78, 0xf6, 0x45, 0xa8, 0x78, 0xa7, 0x49, 0x7d, 0x49, 0x7f, 0xf2, 0x10, 0x95, 0xf3, 0xc4,
            0xdb, 0xd5, 0x82, 0x54, 0x5b, 0x50, 0xf3, 0x0f, 0x94, 0x51, 0x30, 0xd3, 0xa2, 0x20, 0xfe, 0x65
        ]
        ),
        ("💡🔒🌍", [
            0xc6, 0x19, 0x1a, 0x7c, 0x1c, 0x77, 0x4b, 0xcb, 0x9b, 0x43, 0x94, 0x17, 0x23, 0xe5, 0x4a, 0xa6,
            0xe5, 0xd3, 0xc6, 0x93, 0x99, 0x4f, 0xad, 0xcf, 0x40, 0x68, 0xa1, 0x61, 0xf8, 0x31, 0x79, 0x78
        ]
        ),
        ("hello", [
            0x2c, 0xf2, 0x4d, 0xba, 0x5f, 0xb0, 0xa3, 0x0e, 0x26, 0xe8, 0x3b, 0x2a, 0xc5, 0xb9, 0xe2, 0x9e,
            0x1b, 0x16, 0x1e, 0x5c, 0x1f, 0xa7, 0x42, 0x5e, 0x73, 0x04, 0x33, 0x62, 0x93, 0x8b, 0x98, 0x24
        ]
        ),
        ("Rust implementation of SHA-256", [
            0x88, 0x31, 0x18, 0x98, 0x3a, 0xe2, 0xf8, 0x76, 0x44, 0x56, 0x71, 0x4f, 0x3b, 0x33, 0xf6, 0x2b,
            0x8b, 0x9a, 0xd0, 0x92, 0x02, 0x1a, 0x2a, 0xdf, 0xf9, 0x7a, 0xdc, 0x20, 0x8e, 0xee, 0x09, 0x48
        ]
        ),
        ("abc", [
            0xba, 0x78, 0x16, 0xbf, 0x8f, 0x01, 0xcf, 0xea, 0x41, 0x41, 0x40, 0xde, 0x5d, 0xae, 0x22, 0x23,
            0xb0, 0x03, 0x61, 0xa3, 0x96, 0x17, 0x7a, 0x9c, 0xb4, 0x10, 0xff, 0x61, 0xf2, 0x00, 0x15, 0xad
        ]
        ),
        ("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq", [
            0x24, 0x8d, 0x6a, 0x61, 0xd2, 0x06, 0x38, 0xb8, 0xe5, 0xc0, 0x26, 0x93, 0x0c, 0x3e, 0x60, 0x39,
            0xa3, 0x3c, 0xe4, 0x59, 0x64, 0xff, 0x21, 0x67, 0xf6, 0xec, 0xed, 0xd4, 0x19, 0xdb, 0x06, 0xc1
        ]
        ),
        ("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu", [
            0xcf, 0x5b, 0x16, 0xa7, 0x78, 0xaf, 0x83, 0x80, 0x03, 0x6c, 0xe5, 0x9e, 0x7b, 0x04, 0x92, 0x37,
            0x0b, 0x24, 0x9b, 0x11, 0xe8, 0xf0, 0x7a, 0x51, 0xaf, 0xac, 0x45, 0x03, 0x7a, 0xfe, 0xe9, 0xd1
        ]
        ),
    ];

    #[test]
    fn test_sha256() {
        for (data, expected) in HASHES.iter() {
            let hash = hash(data.as_bytes().to_vec());
            assert_eq!(&hash, expected);
        }
    }
}

#[cfg(bench)]
mod benches {
    use super::*;
    use crate::test::{black_box, Bencher};

    #[bench]
    pub fn sha256_hash(bh: &mut Bencher) {
        bh.iter(|| {
            black_box(hash("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"));
        });
    }
}
