use std::cmp::min;

use super::constants::{SM3_BLOCK_SIZE as BLOCK_SIZE, SM3_IV as IV};


pub struct SM3 {
    buffer: [u8; BLOCK_SIZE as usize],
    total_length: usize,
    v: [u32; 8],
}

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

impl SM3 {

    pub fn new() -> Self {
        Self {
            buffer: [0u8; BLOCK_SIZE as usize],
            total_length: 0,
            v: IV,
        }
    }

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

            self.total_length += current_length as usize;
            offset += current_length as usize;

            if self.total_length % BLOCK_SIZE as usize == 0 {
                self.f533(self.f532());
            }
        }
    }

    pub fn r#final(&mut self) -> [u8; 32] {
        let k = (120 - self.total_length % BLOCK_SIZE as usize) as u8 % BLOCK_SIZE;

        let mut message = [0u8; BLOCK_SIZE as usize];
        message[0] = 128;
        for i in 0..8 {
            message[1 + (k as usize - 1) + i] = (self.total_length * 8 / 256_usize.pow((7 - i) as u32)) as u8;
        }
        self.update(&message[..1 + (k as usize - 1) + 8]);

        let mut bytes = [0u8; 32];
        for i in 0..self.v.len() {
            bytes[i * 4..i * 4 + 4].copy_from_slice(self.v[i].to_be_bytes().as_slice());
        }

        bytes
    }

    const fn f42_t(j: u8) -> u32 {
        match j {
            0..=15 => 0x79CC4519,
            16..=63 => 0x7A879D8A,
            _ => unreachable!(),
        }
    }

    const fn f43_ff(j: u8, x: u32, y: u32, z: u32) -> u32 {
        match j {
            0..=15 => x ^ y ^ z,
            16..=63 => (x & y) | (x & z) | (y & z),
            _ => unreachable!(),
        }
    }

    const fn f43_gg(j: u8, x: u32, y: u32, z: u32) -> u32 {
        match j {
            0..=15 => x ^ y ^ z,
            16..=63 => (x & y) | (!x & z),
            _ => unreachable!(),
        }
    }

    const fn f44_p0(x: u32) -> u32 {
        x ^ x.rotate_left(9) ^ x.rotate_left(17)
    }

    const fn f44_p1(x: u32) -> u32 {
        x ^ x.rotate_left(15) ^ x.rotate_left(23)
    }

    fn f532(&self) -> ([u32; 68], [u32; 64]) {
        let mut words: [u32; 16] = [0u32; 16];
        for (i, item) in words.iter_mut().enumerate() {
            *item = (self.buffer[i * 4] as u32) << 24 | (self.buffer[i * 4 + 1] as u32) << 16 | (self.buffer[i * 4 + 2] as u32) << 8 | (self.buffer[i * 4 + 3] as u32);
        }

        let mut w: [u32; 68] = [0u32; 68];
        w[..16].copy_from_slice(&words[..16]);
        for j in 16..=67 {
            w[j] = Self::f44_p1(w[j - 16] ^ w[j - 9] ^ w[j - 3].rotate_left(15)) ^ w[j - 13].rotate_left(7) ^ w[j - 6];
        }
        let mut w_1: [u32; 64] = [0u32; 64];
        for j in 0..=63 {
            w_1[j] = w[j] ^ w[j + 4];
        }

        (w, w_1)
    }

    fn f533(&mut self, (w, w_1): ([u32; 68], [u32; 64])) {
        let mut v = self.v;
        for j in 0..=63 {
            let ss1: u32 = ((v[0].rotate_left(12)).wrapping_add(v[4]).wrapping_add(Self::f42_t(j).rotate_left((j & 31) as u32))).rotate_left(7);
            let ss2: u32 = ss1 ^ v[0].rotate_left(12);
            let tt1: u32 = Self::f43_ff(j, v[0], v[1], v[2]).wrapping_add(v[3]).wrapping_add(ss2).wrapping_add(w_1[j as usize]);
            let tt2: u32 = Self::f43_gg(j, v[4], v[5], v[6]).wrapping_add(v[7]).wrapping_add(ss1).wrapping_add(w[j as usize]);
            v[3] = v[2];
            v[2] = v[1].rotate_left(9);
            v[1] = v[0];
            v[0] = tt1;
            v[7] = v[6];
            v[6] = v[5].rotate_left(19);
            v[5] = v[4];
            v[4] = Self::f44_p0(tt2);
        }

        for (i, item) in self.v.iter_mut().enumerate() {
            *item ^= v[i]
        }
    }
}
