use std::cmp::min;

use crate::common::paddings::{Method2, PaddingModeImpl};
use crate::libs::utils::{
    bytes_to_dwords, dwords_add, dwords_to_bytes, dwords_as_qwords,
    bytes_to_qwords, qwords_add, qwords_to_bytes, qwords_to_dwords,
};

use super::constants::{
    SHA224_BLOCK_SIZE, SHA224_OUTPUT_SIZE, SHA224_MESSAGE_SIZE, SHA224_IV,
    SHA256_BLOCK_SIZE, SHA256_OUTPUT_SIZE, SHA256_MESSAGE_SIZE, SHA256_IV,
    SHA256_CK,
    SHA384_BLOCK_SIZE, SHA384_OUTPUT_SIZE, SHA384_MESSAGE_SIZE, SHA384_IV,
    SHA512_BLOCK_SIZE, SHA512_OUTPUT_SIZE, SHA512_MESSAGE_SIZE, SHA512_IV,
    SHA512_CK,
};
use super::{CryptographicHash, CryptographicHashImpl};


/// 安全散列算法2流式计算器（Secure Hash Algorithm 2 Stream Mode Calculator）
///
/// 安全散列算法2是一种密码散列函数算法标准，由美国国家安全局研发，由美国国家标准与技术研究院在2001年发布。属于SHA算法之一，是SHA-1的后继者。\
/// SHA-2 (Secure Hash Algorithm 2) is a set of cryptographic hash functions designed by the United States National Security Agency
/// and first published by the National Institute of Standards and Technology as a U.S. federal standard.
///
/// # 示例（Examples）
///
/// ```rust
/// # fn main() -> Result<(), ()> {
/// use neuedu_cryptos::cryptographic_hashs::{CryptographicHash, CryptographicHashImpl, SHA2};
///
/// const MESSAGE: [u8; 3] = [
///     0x61, 0x62, 0x63,
/// ];
///
/// let mut instance = SHA2::new(CryptographicHash::SHA256);
/// instance.update(&MESSAGE)?;
/// let digest = instance.r#final()?;
///
/// assert_eq!(digest, [
///     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,
/// ]);
/// #
/// #     Ok(())
/// # }
/// ```
pub struct SHA2 {
    variant: CryptographicHash,
    block_size: u8,
    output_size: u8,
    message_size: u8,
    h: [u64; 8],
    buffer: [u8; 128],
    total_length: u128,
}

impl SHA2 {

    pub fn new(variant: CryptographicHash) -> Self {
        match variant {
            CryptographicHash::SHA224 => Self {
                variant,
                block_size: SHA224_BLOCK_SIZE,
                output_size: SHA224_OUTPUT_SIZE,
                message_size: SHA224_MESSAGE_SIZE,
                h: SHA224_IV,
                buffer: [0u8; 128],
                total_length: 0,
            },
            CryptographicHash::SHA256 => Self {
                variant,
                block_size: SHA256_BLOCK_SIZE,
                output_size: SHA256_OUTPUT_SIZE,
                message_size: SHA256_MESSAGE_SIZE,
                h: SHA256_IV,
                buffer: [0u8; 128],
                total_length: 0,
            },
            CryptographicHash::SHA384 => Self {
                variant,
                block_size: SHA384_BLOCK_SIZE,
                output_size: SHA384_OUTPUT_SIZE,
                message_size: SHA384_MESSAGE_SIZE,
                h: SHA384_IV,
                buffer: [0u8; 128],
                total_length: 0,
            },
            CryptographicHash::SHA512 => Self {
                variant,
                block_size: SHA512_BLOCK_SIZE,
                output_size: SHA512_OUTPUT_SIZE,
                message_size: SHA512_MESSAGE_SIZE,
                h: SHA512_IV,
                buffer: [0u8; 128],
                total_length: 0,
            },
            _ => panic!(),
        }
    }

    fn padding(&mut self) -> Box<[u8]> {
        let message_length = self.total_length * 8;

        let pad_message = Method2::new(self.block_size.into()).pad(&self.buffer[..(self.total_length % self.block_size as u128) as usize]);

        let mut result = Vec::<u8>::with_capacity(self.block_size as usize * 2);
        let message_size = pad_message.len();
        if 1 <= message_size && message_size <= self.message_size.into() {
            result.append(&mut pad_message.into());
            result.resize((self.block_size - self.message_size).into(), 0);
        } else {
            result.append(&mut pad_message[0..pad_message.len() - self.message_size as usize].into());
        }
        result.append(&mut message_length.to_be_bytes()[16 - self.message_size as usize..].into());

        result.into()
    }

    fn extend_dwords(&self) -> [u32; 64] {
        let words = bytes_to_dwords(&self.buffer);

        let mut w: [u32; 64] = [0; 64];
        w[..16].copy_from_slice(&words[..16]);
        for i in 16..=63 {
            let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15]>> 3);
            let s1 = 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);
        }

        w
    }

    fn extend_qwords(&self) -> [u64; 80] {
        let words = bytes_to_qwords(&self.buffer);

        let mut w: [u64; 80] = [0; 80];
        w[..16].copy_from_slice(&words[..16]);
        for i in 16..=79 {
            let s0 = w[i - 15].rotate_right(1) ^ w[i - 15].rotate_right(8) ^ (w[i - 15]>> 7);
            let s1 = w[i - 2].rotate_right(19) ^ w[i - 2].rotate_right(61) ^ (w[i - 2]>> 6);
            w[i] = w[i - 16].wrapping_add(s0).wrapping_add(w[i - 7]).wrapping_add(s1);
        }

        w
    }

    fn compress_dwords(&mut self, w: [u32; 64]) {
        let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = *qwords_to_dwords(&self.h) else { unreachable!() };
        for (i, item) in w.iter().enumerate() {
            let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
            let ch = (e & f) ^ (!e & g);
            let temp1 = h.wrapping_add(s1).wrapping_add(ch).wrapping_add(SHA256_CK[i] as u32).wrapping_add(*item);
            let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
            let maj = (a & b) ^ (a & c) ^ (b & c);
            let temp2 = 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);
        }

        let result = dwords_add(&qwords_to_dwords(&self.h), &[a, b, c, d, e, f, g, h]);
        self.h.copy_from_slice(&dwords_as_qwords(&result));
    }

    fn compress_qwords(&mut self, w: [u64; 80]) {
        let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = self.h;
        for (i, item) in w.iter().enumerate() {
            let s1 = e.rotate_right(14) ^ e.rotate_right(18) ^ e.rotate_right(41);
            let ch = (e & f) ^ (!e & g);
            let temp1 = h.wrapping_add(s1).wrapping_add(ch).wrapping_add(SHA512_CK[i]).wrapping_add(*item);
            let s0 = a.rotate_right(28) ^ a.rotate_right(34) ^ a.rotate_right(39);
            let maj = (a & b) ^ (a & c) ^ (b & c);
            let temp2 = 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);
        }

        let result = qwords_add(&self.h, &[a, b, c, d, e, f, g, h]);
        self.h.copy_from_slice(&result);
    }
}

impl CryptographicHashImpl for SHA2 {

    fn block_size(&self) -> usize {
        self.block_size.into()
    }

    fn output_size(&self) -> usize {
        self.output_size.into()
    }

    fn update(&mut self, message: &[u8]) -> Result<(), ()> {
        let mut overflow;
        let mut offset: usize = 0;
        while offset < message.len() {
            let current_length = min(
                message.len() - offset,
                (self.block_size as u128 - self.total_length % self.block_size as u128) as usize
            ) as u8;
            self.buffer[(self.total_length % self.block_size as u128) as usize..(self.total_length % self.block_size as u128 + current_length as u128) as usize]
                .copy_from_slice(&message[offset..(offset + current_length as usize)]);

            (self.total_length, overflow) = self.total_length.overflowing_add(current_length as u128);
            if self.total_length.leading_zeros() < (128 - self.message_size).into() || overflow {
                return Err(());
            }
            offset += current_length as usize;

            if self.total_length % self.block_size as u128 == 0 {
                match self.variant {
                    CryptographicHash::SHA224 | CryptographicHash::SHA256 => self.compress_dwords(self.extend_dwords()),
                    CryptographicHash::SHA384 | CryptographicHash::SHA512 => self.compress_qwords(self.extend_qwords()),
                    _ => panic!(),
                };
            }
        }

        Ok(())
    }

    fn r#final(&mut self) -> Result<Vec<u8>, ()> {
        let pad_message = self.padding();
        self.update(&pad_message).unwrap();

        match self.variant {
            CryptographicHash::SHA224 => Ok(dwords_to_bytes(&qwords_to_dwords(&self.h[..7])).into()),
            CryptographicHash::SHA256 => Ok(dwords_to_bytes(&qwords_to_dwords(&self.h)).into()),
            CryptographicHash::SHA384 => Ok(qwords_to_bytes(&self.h[..6]).into()),
            CryptographicHash::SHA512 => Ok(qwords_to_bytes(&self.h).into()),
            _ => unreachable!(),
        }
    }
}
