use crate::common::*;
use crate::{Code, Error, Hash, Result};
use core::fmt;
use core::fmt::Formatter;

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 H0: [u32; 8] = [
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F,
    0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
];

const MB_SIZE: usize = 64;

pub const SIZE: usize = 32;

fn u32_to_chunk_in_place(x: u32, chunk: &mut [u8]) {
    assert!(chunk.len() >= 4);
    chunk[0] = ((x & 0xff000000) >> 24) as u8;
    chunk[1] = ((x & 0x00ff0000) >> 16) as u8;
    chunk[2] = ((x & 0x0000ff00) >> 8) as u8;
    chunk[3] = (x & 0x000000ff) as u8;
}

#[derive(Debug, Clone, Copy)]
pub struct SHA256 {
    mlen: u64,
    i_hash: [u32; 8],

    pub(crate) hresult: Code,
    complete: bool,

    mbi: usize,
    mb: [u8; MB_SIZE],
}

impl fmt::Display for SHA256 {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "SHA256<{:?}, ", self.hresult)?;
        if !self.complete {
            write!(f, "in")?;
        }

        write!(f, "complete>")
    }
}

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

impl SHA256 {
    pub fn new() -> Self {
        Self {
            mlen: 0,
            i_hash: H0,
            hresult: Code::OK,
            complete: false,
            mbi: 0,
            mb: [0; MB_SIZE],
        }
    }

    fn chunk_to_u32(&self, offset: usize) -> u32 {
        let mut chunk: u32 = 0;
        let mut i: usize = offset;

        while i < offset + 4 {
            chunk <<= 8;
            chunk += self.mb[i] as u32;
            i += 1;
        }

        chunk
    }

    fn update_message_block(&mut self) {
        let mut w: [u32; MB_SIZE] = [0; MB_SIZE];
        let mut chunk = 0;
        let mut i: usize = 0;

        while i < 16 {
            w[i] = self.chunk_to_u32(chunk);
            chunk += 4;
            i += 1;
        }

        self.mbi = 0;
        i = 16;
        while i < 64 {
            w[i] = sha_σ1(w[i - 2])
                .wrapping_add(w[i - 7])
                .wrapping_add(sha_σ0(w[i - 15]))
                .wrapping_add(w[i - 16]);
            i += 1;
        }

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

        i = 0;
        while i < 64 {
            let t1 = h
                .wrapping_add(sha_Σ1(e))
                .wrapping_add(sha_ch(e, f, g))
                .wrapping_add(K[i])
                .wrapping_add(w[i]);
            let t2 = sha_Σ0(a).wrapping_add(sha_maj(a, b, c));
            h = g;
            g = f;
            f = e;
            e = d.wrapping_add(t1);
            d = c;
            c = b;
            b = a;
            a = t1.wrapping_add(t2);

            i += 1
        }

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

    fn add_length(&mut self, δ: usize) -> Result<()> {
        let mδ = self.mlen + δ as u64;
        if mδ < self.mlen {
            return Err(Error::with(Code::InputTooLong));
        }

        self.mlen = mδ;
        Ok(())
    }

    fn pad_message(&mut self, pc: u8) -> Result<()> {
        if self.mbi < (MB_SIZE - 8) {
            self.mb[self.mbi] = pc;
            self.mbi += 1;
        } else {
            let mut pc_add = false;

            if self.mbi < MB_SIZE - 1 {
                self.mb[self.mbi] = pc;
                self.mbi += 1;
                pc_add = true;
            }

            while self.mbi < MB_SIZE {
                self.mb[self.mbi] = 0;
                self.mbi += 1;
            }

            self.update_message_block();
            if !pc_add {
                self.mb[self.mbi] = pc;
                self.mbi += 1;
            }

            // Assumption: updating the message block has not left the
            // context in a corrupted state.
            debug_assert_eq!(self.hresult, Code::OK);
        }

        while self.mbi < (MB_SIZE - 8) {
            self.mb[self.mbi] = 0;
            self.mbi += 1;
        }

        const LSTART: usize = MB_SIZE - 8;
        self.mb[LSTART] = (self.mlen >> 56) as u8;
        self.mb[LSTART + 1] =
            ((self.mlen & 0x00ff000000000000) >> 48) as u8;
        self.mb[LSTART + 2] =
            ((self.mlen & 0x0000ff0000000000) >> 40) as u8;
        self.mb[LSTART + 3] =
            ((self.mlen & 0x000000ff00000000) >> 32) as u8;
        self.mb[LSTART + 4] =
            ((self.mlen & 0x00000000ff000000) >> 24) as u8;
        self.mb[LSTART + 5] =
            ((self.mlen & 0x0000000000ff0000) >> 16) as u8;
        self.mb[LSTART + 6] =
            ((self.mlen & 0x000000000000ff00) >> 8) as u8;
        self.mb[LSTART + 7] = (self.mlen & 0x00000000000000ff) as u8;

        self.update_message_block();

        // Assumption: updating the message block has not left the
        // context in a corrupted state.
        debug_assert_eq!(self.hresult, Code::OK);
        Ok(())
    }

    fn copy_out_digest(&self, digest: &mut [u8]) {
        u32_to_chunk_in_place(self.i_hash[0], digest);
        u32_to_chunk_in_place(self.i_hash[1], &mut digest[4..]);
        u32_to_chunk_in_place(self.i_hash[2], &mut digest[8..]);
        u32_to_chunk_in_place(self.i_hash[3], &mut digest[12..]);
        u32_to_chunk_in_place(self.i_hash[4], &mut digest[16..]);
        u32_to_chunk_in_place(self.i_hash[5], &mut digest[20..]);
        u32_to_chunk_in_place(self.i_hash[6], &mut digest[24..]);
        u32_to_chunk_in_place(self.i_hash[7], &mut digest[28..]);
    }
}

impl Hash for SHA256 {
    fn reset(&mut self) -> Result<()> {
        // The message block is set to the initial hash vector.
        self.i_hash[0] = H0[0];
        self.i_hash[1] = H0[1];
        self.i_hash[2] = H0[2];
        self.i_hash[3] = H0[3];
        self.i_hash[4] = H0[4];
        self.i_hash[5] = H0[5];
        self.i_hash[6] = H0[6];
        self.i_hash[7] = H0[7];

        self.mb.fill(0);
        self.mbi = 0;
        self.hresult = Code::OK;
        self.complete = false;
        self.mlen = 0;
        Ok(())
    }

    fn update(&mut self, msg: &[u8]) -> Result<()> {
        if msg.is_empty() {
            self.hresult = Code::OK;
            return Ok(());
        }

        if self.hresult != Code::OK {
            return Err(Error::with(self.hresult));
        }

        if self.complete {
            return Err(Error::with(Code::InvalidState));
        }

        let mut i: usize = 0;
        while i < msg.len() {
            self.mb[self.mbi] = msg[i];
            self.mbi += 1;
            self.add_length(8)?;
            if self.mbi == MB_SIZE {
                self.update_message_block();
            }

            i += 1;
        }

        if self.hresult != Code::OK {
            Err(Error::with(self.hresult))
        } else {
            Ok(())
        }
    }

    fn finalize(&mut self, digest: &mut [u8]) -> Result<()> {
        if digest.len() < SIZE {
            return Err(Error::with(Code::BufferTooSmall));
        }

        if self.hresult != Code::OK {
            return Err(Error::with(self.hresult));
        }

        if self.complete {
            return Err(Error::with(Code::InvalidState));
        }

        self.pad_message(0x80)?;

        let mut i: usize = 0;
        while i < self.mb.len() {
            self.mb[i] = 0;
            i += 1;
        }

        self.complete = true;
        self.mlen = 0;

        self.copy_out_digest(digest);
        self.hresult = Code::OK;

        Ok(())
    }

    fn result(&self, digest: &mut [u8]) -> Result<()> {
        if digest.len() < SIZE {
            return Err(Error::with(Code::BufferTooSmall));
        }

        if self.hresult != Code::OK {
            return Err(Error::with(self.hresult));
        }

        if !self.complete {
            return Err(Error::with(Code::HashNotFinalized));
        }

        self.copy_out_digest(digest);
        Ok(())
    }

    fn size() -> usize {
        SIZE
    }
}

fn run_test(input: &[u8], expected: &[u8]) -> Result<()> {
    let mut h = SHA256::default();
    let mut d: [u8; SIZE] = [0; SIZE];

    h.update(input)?;
    h.finalize(&mut d)?;
    assert_eq!(d, expected);

    d = [0; SIZE];
    h.result(&mut d)?;
    assert_eq!(d, expected);

    Ok(())
}

fn run_reset_cycle() -> Result<()> {
    let mut h = SHA256::default();
    let mut d: [u8; SIZE] = [0; SIZE];

    h.update(b"hello, world")?;
    h.finalize(&mut d)?;
    assert_eq!(&d, HELLO_WORLD);

    d = [0; SIZE];
    h.reset()?;
    h.finalize(&mut d)?;
    assert_eq!(&d, EMPTY_VECTOR);

    h.reset()?;
    h.update(b"hello, world")?;
    h.finalize(&mut d)?;
    assert_eq!(&d, HELLO_WORLD);

    Ok(())
}

const EMPTY_VECTOR: &[u8; SIZE] = &[
    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,
];

const HELLO_WORLD: &[u8; SIZE] = &[
    0x09, 0xca, 0x7e, 0x4e, 0xaa, 0x6e, 0x8a, 0xe9, 0xc7, 0xd2, 0x61,
    0x16, 0x71, 0x29, 0x18, 0x48, 0x83, 0x64, 0x4d, 0x07, 0xdf, 0xba,
    0x7c, 0xbf, 0xbc, 0x4c, 0x8a, 0x2e, 0x08, 0x36, 0x0d, 0x5b,
];

/// self_test runs a quick test cycle, and can be used to validate
/// correctness in systems on startup.
pub fn self_test() -> Result<()> {
    run_test(b"", EMPTY_VECTOR)?;
    run_test(b"hello, world", HELLO_WORLD)?;
    run_reset_cycle()?;
    Ok(())
}