use crate::hash::ffi::SHA2Outcome;

pub const SHA2_256_HASH_BITS: usize = 256;
pub const SHA2_256_HASH_BYTES: usize = SHA2_256_HASH_BITS / 8;
pub const SHA2_256_HASH_PARTS: usize = SHA2_256_HASH_BYTES / 4 /* storing with u32 */;
pub const SHA2_256_BLOCK_SIZE: usize = 2 * SHA2_256_HASH_BYTES;

pub type Digest256 = [u32; SHA2_256_HASH_PARTS];
pub(crate) const DIGEST256_EMPTY: Digest256 = [0; SHA2_256_HASH_PARTS];

macro_rules! show_internals {
    ($block:block) => {
        #[cfg(feature = "show_internals")]
        $block
    };
}

#[derive(PartialEq, Eq, Clone, Hash)]
pub struct Sha2_256Context {
    intermediate_hash: Digest256,
    block: [u8; SHA2_256_BLOCK_SIZE],
    block_idx: usize,
    length_low: usize,
    length_high: usize,
}

impl Sha2_256Context {
    const K: [u32; SHA2_256_BLOCK_SIZE] = [
        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 H: [u32; SHA2_256_HASH_PARTS] = [
        0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab,
        0x5be0cd19,
    ];
    const PAD_BYTE: u8 = 0x80;

    pub fn new() -> Self {
        Self {
            intermediate_hash: Self::H,
            block: [0; SHA2_256_BLOCK_SIZE],
            block_idx: Default::default(),
            length_low: Default::default(),
            length_high: Default::default(),
        }
    }

    #[inline]
    pub fn clear(&mut self) {
        *self = Self::new();
    }

    pub fn input(&mut self, data: &[u8]) -> Result<(), SHA2Outcome> {
        if data.is_empty() {
            return Ok(());
        }
        for byte in data {
            self.block[self.block_idx] = *byte;
            self.block_idx += 1;
            if self.add_length(8).is_ok() && self.block_idx == SHA2_256_BLOCK_SIZE {
                self.process_block();
            }
        }
        Ok(())
    }

    pub fn finish(mut self) -> Result<Digest256, SHA2Outcome> {
        self.finalize();
        Ok(self.intermediate_hash)
    }

    fn process_block(&mut self) {
        show_internals!({
            self.show_internal_state("start processing");
        });

        let mut w = [0; SHA2_256_BLOCK_SIZE];

        #[allow(clippy::needless_range_loop)]
        for t in 0..16 {
            let t4 = t * 4;
            w[t] = ((self.block[t4] as u32) << 24)
                | ((self.block[t4 + 1] as u32) << 16)
                | ((self.block[t4 + 2] as u32) << 8)
                | (self.block[t4 + 3] as u32);
        }

        // NOTE: The sigma functions used for initializing W and for calculating the alphabet
        // variables differ.
        for t in 16..SHA2_256_BLOCK_SIZE {
            w[t] = Self::sha256_sigma1_w(w[t - 2])
                .wrapping_add(w[t - 7])
                .wrapping_add(Self::sha256_sigma0_w(w[t - 15]))
                .wrapping_add(w[t - 16]);
        }

        show_internals!({
            for (i, w) in w.iter().enumerate() {
                println!("w[{i}]:\t{w:032b}")
            }
        });

        let mut a = self.intermediate_hash[0];
        let mut b = self.intermediate_hash[1];
        let mut c = self.intermediate_hash[2];
        let mut d = self.intermediate_hash[3];
        let mut e = self.intermediate_hash[4];
        let mut f = self.intermediate_hash[5];
        let mut g = self.intermediate_hash[6];
        let mut h = self.intermediate_hash[7];

        show_internals!({
            println!("A:\t{a:032b}");
            println!("B:\t{b:032b}");
            println!("C:\t{c:032b}");
            println!("D:\t{d:032b}");
            println!("E:\t{e:032b}");
            println!("F:\t{f:032b}");
            println!("G:\t{g:032b}");
            println!("H:\t{h:032b}");
        });

        let mut temp1;
        let mut temp2;
        #[allow(clippy::needless_range_loop)]
        for t in 0..SHA2_256_BLOCK_SIZE {
            show_internals!({
                println!("Iter {t}");
            });

            temp1 = h
                .wrapping_add(Self::sha256_sigma1(e))
                .wrapping_add(Self::sha_ch(e, f, g))
                .wrapping_add(Self::K[t])
                .wrapping_add(w[t]);
            temp2 = Self::sha256_sigma0(a).wrapping_add(Self::sha_maj(a, b, c));
            h = g;
            g = f;
            f = e;
            e = d.wrapping_add(temp1);
            d = c;
            c = b;
            b = a;
            a = temp1.wrapping_add(temp2);

            show_internals!({
                println!("temp1:\t{temp1:032b} {temp1:08x}");
                println!("temp2:\t{temp2:032b} {temp2:08x}");
                println!("A:\t{a:032b} {a:08x}");
                println!("B:\t{b:032b} {b:08x}");
                println!("C:\t{c:032b} {c:08x}");
                println!("D:\t{d:032b} {d:08x}");
                println!("E:\t{e:032b} {e:08x}");
                println!("F:\t{f:032b} {f:08x}");
                println!("G:\t{g:032b} {g:08x}");
                println!("H:\t{h:032b} {h:08x}");
                for (i, h) in self.intermediate_hash.iter().enumerate() {
                    println!("h[{i}]:\t{h:032b} {h:08x}");
                }
            });
        }

        show_internals!({
            self.show_internal_state("assigning intermediate_hash");
        });

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

        self.block_idx = 0;

        show_internals!({
            self.show_internal_state("processing end");
            for (i, h) in self.intermediate_hash.iter().enumerate() {
                println!("h[{i}]:\t{h:032b}");
            }
        });
    }

    fn add_length(&mut self, len: usize) -> Result<(), SHA2Outcome> {
        let add_temp = self.length_low;

        // Update length_low and check for overflow
        self.length_low = self.length_low.wrapping_add(len);
        if self.length_low < add_temp {
            // Overflow occurred, increment length_high
            self.length_high = self.length_high.wrapping_add(1);
            if self.length_high == 0 {
                return Err(SHA2Outcome::shaInputTooLong);
            }
        }

        Ok(())
    }

    #[inline]
    fn finalize(&mut self) {
        self.pad_block();
    }

    fn pad_block(&mut self) {
        if self.block_idx >= (SHA2_256_BLOCK_SIZE - 8) {
            self.block[self.block_idx] = Self::PAD_BYTE;
            self.block_idx += 1;
            while self.block_idx < SHA2_256_BLOCK_SIZE {
                self.block[self.block_idx] = 0;
                self.block_idx += 1;
            }
            self.process_block();
        } else {
            self.block[self.block_idx] = Self::PAD_BYTE;
            self.block_idx += 1;
        }

        while self.block_idx < (SHA2_256_BLOCK_SIZE - 8) {
            self.block[self.block_idx] = 0;
            self.block_idx += 1;
        }

        self.block[56] = (self.length_high >> 24) as u8;
        self.block[57] = (self.length_high >> 16) as u8;
        self.block[58] = (self.length_high >> 8) as u8;
        self.block[59] = (self.length_high) as u8;
        self.block[60] = (self.length_low >> 24) as u8;
        self.block[61] = (self.length_low >> 16) as u8;
        self.block[62] = (self.length_low >> 8) as u8;
        self.block[63] = (self.length_low) as u8;

        self.process_block();
    }

    #[cfg(feature = "show_internals")]
    fn show_internal_state(&self, msg: impl core::fmt::Display) {
        println!("{msg}\n{self:#x?}")
    }
}

/// internal functions for SHA2 process_block
impl Sha2_256Context {
    #[inline]
    fn sha256_shr(bits: u32, word: u32) -> u32 {
        word >> bits
    }

    #[inline]
    fn sha256_rotr(bits: u32, word: u32) -> u32 {
        word.rotate_right(bits)
    }

    /// This is a different function for W
    #[inline]
    fn sha256_sigma0_w(word: u32) -> u32 {
        Self::sha256_rotr(7, word) ^ Self::sha256_rotr(18, word) ^ Self::sha256_shr(3, word)
    }

    /// This is a different function for W
    #[inline]
    fn sha256_sigma1_w(word: u32) -> u32 {
        Self::sha256_rotr(17, word) ^ Self::sha256_rotr(19, word) ^ Self::sha256_shr(10, word)
    }

    // NOTE: these do not use shift right for the last internal variable!
    #[inline]
    fn sha256_sigma0(word: u32) -> u32 {
        Self::sha256_rotr(2, word) ^ Self::sha256_rotr(13, word) ^ Self::sha256_rotr(22, word)
    }

    #[inline]
    fn sha256_sigma1(word: u32) -> u32 {
        Self::sha256_rotr(6, word) ^ Self::sha256_rotr(11, word) ^ Self::sha256_rotr(25, word)
    }

    #[inline]
    fn sha_ch(x: u32, y: u32, z: u32) -> u32 {
        (x & (y ^ z)) ^ z
    }

    #[inline]
    fn sha_maj(x: u32, y: u32, z: u32) -> u32 {
        (x & y) ^ (x & z) ^ (y & z)
    }
}

impl Default for Sha2_256Context {
    fn default() -> Self {
        Self::new()
    }
}

pub fn sha2_256_oneshot(data: &[u8]) -> Result<Digest256, SHA2Outcome> {
    let mut ctx = Sha2_256Context::default();
    ctx.input(data)?;
    ctx.finish()
}

#[cfg(feature = "show_internals")]
#[allow(unused)]
mod debug {
    use super::*;

    struct InlineArrHex2<'a, T: core::fmt::Debug + core::fmt::LowerHex>(&'a [T]);

    impl<'a, T> core::fmt::Debug for InlineArrHex2<'a, T>
    where
        T: core::fmt::Debug + core::fmt::LowerHex,
    {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            writeln!(f, "[")?;
            for (i, t) in self.0.iter().enumerate() {
                write!(f, "0x{t:02x}")?;
                if i < self.0.len() - 1 {
                    write!(f, ", ")?;
                }
                if i % 8 == 7 && i < self.0.len() - 1 {
                    writeln!(f)?;
                }
            }
            write!(f, "\n]")?;
            Ok(())
        }
    }

    impl core::fmt::Debug for Sha2_256Context {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            f.debug_struct("Sha2_256Context")
                .field("intermediate_hash", &InlineArrHex2(&self.intermediate_hash))
                .field("block", &InlineArrHex2(&self.block))
                .field("block_idx", &self.block_idx)
                .field("length_low", &self.length_low)
                .field("length_high", &self.length_high)
                .finish()
        }
    }
}
#[cfg(feature = "show_internals")]
#[allow(unused)]
use debug::*;

#[cfg(test)]
mod test {
    use super::*;
    use crate::assert_eq_bin;
    use crate::hash::test::TEST_VALUES;
    use assert_hex::*;
    use pretty_assertions::{assert_eq, assert_ne};

    #[test]
    fn test_sha256_check() {
        let mut dig;
        for (input, expected_output) in TEST_VALUES.iter().copied() {
            dig = sha2_256_oneshot(input.as_bytes()).unwrap();
            println!(
                "Binary of first seg matches: {}\n{:032b}\n{:032b}",
                dig[0] == expected_output[0],
                dig[0],
                expected_output[0]
            );
            assert_eq_hex!(dig, expected_output)
        }
    }

    #[test]
    fn test_sha256_sigma0() {
        let a = 0b01101010000010011110011001100111;
        assert_eq_bin!(
            Sha2_256Context::sha256_sigma0(a),
            0b11001110001000001011010001111110
        );
    }

    #[test]
    fn test_sha256_sigma1() {
        let a = 0b01010001000011100101001001111111;
        assert_eq_bin!(
            Sha2_256Context::sha256_sigma1(a),
            0b00110101100001110010011100101011
        );
    }
}
