// MD4 implementation based on https://docs.rs/md4/latest/md4/
use crate::bytes::Bytes;
const BLOCK_BYTES: usize = 64;

pub struct Md4Core {
    byte_len: u64,
    state: [u32; 4],
}

impl Default for Md4Core {
    #[inline]
    fn default() -> Self {
        let state = [0x6745_2301, 0xEFCD_AB89, 0x98BA_DCFE, 0x1032_5476];
        Self { state, byte_len: 0 }
    }
}
impl Md4Core {
    fn update(&mut self, bytes: &[u8; BLOCK_BYTES]) {
        compress(&mut self.state, bytes);
        self.byte_len += BLOCK_BYTES as u64;
    }

    pub fn get_padding(&self, buffer: &[u8]) -> Vec<u8> {
        let bit_len = self
            .byte_len
            .wrapping_add(buffer.len() as u64)
            .wrapping_mul(8);

        let mut padding = vec![];
        padding.push(0x80);
        while (buffer.len() + padding.len()) % BLOCK_BYTES != (BLOCK_BYTES - 8) {
            padding.push(0x0);
        }
        padding.append(&mut (bit_len as u32).to_le_bytes().to_vec());
        padding.append(&mut ((bit_len >> 32) as u32).to_le_bytes().to_vec());
        padding
    }

    fn finalize(&mut self, mut buffer: Vec<u8>) -> Bytes {
        buffer.append(&mut self.get_padding(&buffer));
        for bytes in buffer.chunks(BLOCK_BYTES) {
            compress(&mut self.state, &bytes.try_into().unwrap());
        }

        Bytes(
            self.state
                .iter()
                .map(|i| i.to_le_bytes())
                .flatten()
                .collect(),
        )
    }

    pub fn hash(&mut self, bytes: &Bytes) -> Bytes {
        let mut final_bytes = vec![];
        for bytes in bytes.0.chunks(BLOCK_BYTES) {
            if bytes.len() == BLOCK_BYTES {
                self.update(&bytes.try_into().unwrap());
            } else {
                final_bytes = bytes.to_vec();
            }
        }
        self.finalize(final_bytes)
    }

    pub fn fix(&mut self, fixate: [u32; 4], byte_len: u64) {
        for i in 0..4 {
            self.state[i] = fixate[i];
        }
        self.byte_len = byte_len;
    }
}

fn compress(state: &mut [u32; 4], input: &[u8; BLOCK_BYTES]) {
    fn f(x: u32, y: u32, z: u32) -> u32 {
        (x & y) | (!x & z)
    }

    fn g(x: u32, y: u32, z: u32) -> u32 {
        (x & y) | (x & z) | (y & z)
    }

    fn h(x: u32, y: u32, z: u32) -> u32 {
        x ^ y ^ z
    }

    fn op1(a: u32, b: u32, c: u32, d: u32, k: u32, s: u32) -> u32 {
        a.wrapping_add(f(b, c, d)).wrapping_add(k).rotate_left(s)
    }

    fn op2(a: u32, b: u32, c: u32, d: u32, k: u32, s: u32) -> u32 {
        a.wrapping_add(g(b, c, d))
            .wrapping_add(k)
            .wrapping_add(0x5A82_7999)
            .rotate_left(s)
    }

    fn op3(a: u32, b: u32, c: u32, d: u32, k: u32, s: u32) -> u32 {
        a.wrapping_add(h(b, c, d))
            .wrapping_add(k)
            .wrapping_add(0x6ED9_EBA1)
            .rotate_left(s)
    }

    let mut a = state[0];
    let mut b = state[1];
    let mut c = state[2];
    let mut d = state[3];

    // load block to data
    let mut data = [0u32; 16];
    for (o, chunk) in data.iter_mut().zip(input.chunks_exact(4)) {
        *o = u32::from_le_bytes(chunk.try_into().unwrap());
    }

    // round 1
    for &i in &[0, 4, 8, 12] {
        a = op1(a, b, c, d, data[i], 3);
        d = op1(d, a, b, c, data[i + 1], 7);
        c = op1(c, d, a, b, data[i + 2], 11);
        b = op1(b, c, d, a, data[i + 3], 19);
    }

    // round 2
    for i in 0..4 {
        a = op2(a, b, c, d, data[i], 3);
        d = op2(d, a, b, c, data[i + 4], 5);
        c = op2(c, d, a, b, data[i + 8], 9);
        b = op2(b, c, d, a, data[i + 12], 13);
    }

    // round 3
    for &i in &[0, 2, 1, 3] {
        a = op3(a, b, c, d, data[i], 3);
        d = op3(d, a, b, c, data[i + 8], 9);
        c = op3(c, d, a, b, data[i + 4], 11);
        b = op3(b, c, d, a, data[i + 12], 15);
    }

    state[0] = state[0].wrapping_add(a);
    state[1] = state[1].wrapping_add(b);
    state[2] = state[2].wrapping_add(c);
    state[3] = state[3].wrapping_add(d);
}

pub fn hash(bytes: &Bytes) -> Bytes {
    let mut md4 = Md4Core::default();
    md4.hash(bytes)
}

pub fn authenticate(message: &Bytes, key: &Bytes) -> Bytes {
    let mut c = vec![];
    c.append(&mut key.0.to_vec());
    c.append(&mut message.0.to_vec());
    hash(&Bytes(c))
}

pub fn verify(message: &Bytes, key: &Bytes, mac: &Bytes) -> bool {
    return authenticate(&message, &key) == *mac;
}