//! BLAKE2b and BLAKE2s hash functions (RFC 7693).
//!
//! BLAKE2b (64-bit, 1..64 byte output) and BLAKE2s (32-bit, 1..32
//! byte output) are cryptographic hash functions designed as faster
//! replacements for MD5 / SHA-1. Not FIPS-approved, but widely used
//! in Argon2, libsodium, WireGuard, and many file-integrity tools.

use crate::Hasher;

// ====================================================================
// BLAKE2b (64-bit)
// ====================================================================

const BLAKE2B_IV: [u64; 8] = [
    0x6A09E667F3BCC908,
    0xBB67AE8584CAA73B,
    0x3C6EF372FE94F82B,
    0xA54FF53A5F1D36F1,
    0x510E527FADE682D1,
    0x9B05688C2B3E6C1F,
    0x1F83D9ABFB41BD6B,
    0x5BE0CD19137E2179,
];

const BLAKE2B_SIGMA: [[usize; 16]; 12] = [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
    [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
    [11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
    [7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
    [9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
    [2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
    [12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
    [13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
    [6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
    [10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0],
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
    [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
];

#[inline(always)]
fn blake2b_g(v: &mut [u64; 16], a: usize, b: usize, c: usize, d: usize, x: u64, y: u64) {
    v[a] = v[a].wrapping_add(v[b]).wrapping_add(x);
    v[d] = (v[d] ^ v[a]).rotate_right(32);
    v[c] = v[c].wrapping_add(v[d]);
    v[b] = (v[b] ^ v[c]).rotate_right(24);
    v[a] = v[a].wrapping_add(v[b]).wrapping_add(y);
    v[d] = (v[d] ^ v[a]).rotate_right(16);
    v[c] = v[c].wrapping_add(v[d]);
    v[b] = (v[b] ^ v[c]).rotate_right(63);
}

fn blake2b_compress(h: &mut [u64; 8], block: &[u8; 128], t: u128, last: bool) {
    let mut v = [0u64; 16];
    v[..8].copy_from_slice(h);
    v[8..16].copy_from_slice(&BLAKE2B_IV);
    v[12] ^= t as u64;
    v[13] ^= (t >> 64) as u64;
    if last {
        v[14] = !v[14];
    }

    let mut m = [0u64; 16];
    for i in 0..16 {
        m[i] = u64::from_le_bytes([
            block[8 * i],
            block[8 * i + 1],
            block[8 * i + 2],
            block[8 * i + 3],
            block[8 * i + 4],
            block[8 * i + 5],
            block[8 * i + 6],
            block[8 * i + 7],
        ]);
    }

    for i in 0..12 {
        let s = &BLAKE2B_SIGMA[i];
        blake2b_g(&mut v, 0, 4, 8, 12, m[s[0]], m[s[1]]);
        blake2b_g(&mut v, 1, 5, 9, 13, m[s[2]], m[s[3]]);
        blake2b_g(&mut v, 2, 6, 10, 14, m[s[4]], m[s[5]]);
        blake2b_g(&mut v, 3, 7, 11, 15, m[s[6]], m[s[7]]);
        blake2b_g(&mut v, 0, 5, 10, 15, m[s[8]], m[s[9]]);
        blake2b_g(&mut v, 1, 6, 11, 12, m[s[10]], m[s[11]]);
        blake2b_g(&mut v, 2, 7, 8, 13, m[s[12]], m[s[13]]);
        blake2b_g(&mut v, 3, 4, 9, 14, m[s[14]], m[s[15]]);
    }

    for i in 0..8 {
        h[i] ^= v[i] ^ v[i + 8];
    }
}

/// BLAKE2b hasher (RFC 7693). Default: 64-byte output, 128-byte blocks.
///
/// Not FIPS-approved. Used by Argon2, libsodium, WireGuard.
#[derive(Clone)]
pub struct Blake2b {
    h: [u64; 8],
    buf: [u8; 128],
    buf_len: usize,
    total: u128,
    out_len: usize,
}

impl Blake2b {
    /// Create a BLAKE2b hasher with a custom output length (1..=64 bytes).
    pub fn with_output_len(out_len: usize) -> Self {
        assert!(out_len >= 1 && out_len <= 64, "BLAKE2b output must be 1..64");
        let mut h = BLAKE2B_IV;
        // Parameter block: digest_length || key_length=0 || fanout=1 || depth=1
        h[0] ^= 0x01010000 ^ (out_len as u64);
        Self {
            h,
            buf: [0u8; 128],
            buf_len: 0,
            total: 0,
            out_len,
        }
    }
}

impl Hasher for Blake2b {
    const OUTPUT_LEN: usize = 64;
    const BLOCK_LEN: usize = 128;

    fn new() -> Self {
        Self::with_output_len(64)
    }

    fn update(&mut self, data: &[u8]) {
        let mut pos = 0;
        while pos < data.len() {
            // If buffer is full and there's more data, compress.
            if self.buf_len == 128 {
                self.total += 128;
                let block = self.buf;
                blake2b_compress(&mut self.h, &block, self.total, false);
                self.buf_len = 0;
            }
            let take = (128 - self.buf_len).min(data.len() - pos);
            self.buf[self.buf_len..self.buf_len + take].copy_from_slice(&data[pos..pos + take]);
            self.buf_len += take;
            pos += take;
        }
    }

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; self.out_len];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        self.total += self.buf_len as u128;
        // Zero-pad the remaining buffer.
        for b in &mut self.buf[self.buf_len..] {
            *b = 0;
        }
        blake2b_compress(&mut self.h, &self.buf, self.total, true);
        let len = out.len().min(self.out_len);
        let mut full = [0u8; 64];
        for i in 0..8 {
            full[8 * i..8 * i + 8].copy_from_slice(&self.h[i].to_le_bytes());
        }
        out[..len].copy_from_slice(&full[..len]);
    }
}

// ====================================================================
// BLAKE2s (32-bit)
// ====================================================================

const BLAKE2S_IV: [u32; 8] = [
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
];

const BLAKE2S_SIGMA: [[usize; 16]; 10] = [
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
    [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
    [11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
    [7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
    [9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
    [2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
    [12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
    [13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
    [6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
    [10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0],
];

#[inline(always)]
fn blake2s_g(v: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize, x: u32, y: u32) {
    v[a] = v[a].wrapping_add(v[b]).wrapping_add(x);
    v[d] = (v[d] ^ v[a]).rotate_right(16);
    v[c] = v[c].wrapping_add(v[d]);
    v[b] = (v[b] ^ v[c]).rotate_right(12);
    v[a] = v[a].wrapping_add(v[b]).wrapping_add(y);
    v[d] = (v[d] ^ v[a]).rotate_right(8);
    v[c] = v[c].wrapping_add(v[d]);
    v[b] = (v[b] ^ v[c]).rotate_right(7);
}

fn blake2s_compress(h: &mut [u32; 8], block: &[u8; 64], t: u64, last: bool) {
    let mut v = [0u32; 16];
    v[..8].copy_from_slice(h);
    v[8..16].copy_from_slice(&BLAKE2S_IV);
    v[12] ^= t as u32;
    v[13] ^= (t >> 32) as u32;
    if last {
        v[14] = !v[14];
    }

    let mut m = [0u32; 16];
    for i in 0..16 {
        m[i] = u32::from_le_bytes([block[4 * i], block[4 * i + 1], block[4 * i + 2], block[4 * i + 3]]);
    }

    for i in 0..10 {
        let s = &BLAKE2S_SIGMA[i];
        blake2s_g(&mut v, 0, 4, 8, 12, m[s[0]], m[s[1]]);
        blake2s_g(&mut v, 1, 5, 9, 13, m[s[2]], m[s[3]]);
        blake2s_g(&mut v, 2, 6, 10, 14, m[s[4]], m[s[5]]);
        blake2s_g(&mut v, 3, 7, 11, 15, m[s[6]], m[s[7]]);
        blake2s_g(&mut v, 0, 5, 10, 15, m[s[8]], m[s[9]]);
        blake2s_g(&mut v, 1, 6, 11, 12, m[s[10]], m[s[11]]);
        blake2s_g(&mut v, 2, 7, 8, 13, m[s[12]], m[s[13]]);
        blake2s_g(&mut v, 3, 4, 9, 14, m[s[14]], m[s[15]]);
    }

    for i in 0..8 {
        h[i] ^= v[i] ^ v[i + 8];
    }
}

/// BLAKE2s hasher (RFC 7693). Default: 32-byte output, 64-byte blocks.
///
/// 32-bit oriented sibling of BLAKE2b. Faster on embedded (ARM
/// Cortex-M) and 32-bit RISC-V targets. Used in WireGuard, Zcash
/// (Equihash), and various file-integrity tools.
#[derive(Clone)]
pub struct Blake2s {
    h: [u32; 8],
    buf: [u8; 64],
    buf_len: usize,
    total: u64,
    out_len: usize,
}

impl Blake2s {
    /// Create a BLAKE2s hasher with a custom output length (1..=32 bytes).
    pub fn with_output_len(out_len: usize) -> Self {
        assert!(out_len >= 1 && out_len <= 32, "BLAKE2s output must be 1..32");
        let mut h = BLAKE2S_IV;
        h[0] ^= 0x01010000 ^ (out_len as u32);
        Self {
            h,
            buf: [0u8; 64],
            buf_len: 0,
            total: 0,
            out_len,
        }
    }
}

impl Hasher for Blake2s {
    const OUTPUT_LEN: usize = 32;
    const BLOCK_LEN: usize = 64;

    fn new() -> Self {
        Self::with_output_len(32)
    }

    fn update(&mut self, data: &[u8]) {
        let mut pos = 0;
        while pos < data.len() {
            if self.buf_len == 64 {
                self.total += 64;
                let block = self.buf;
                blake2s_compress(&mut self.h, &block, self.total, false);
                self.buf_len = 0;
            }
            let take = (64 - self.buf_len).min(data.len() - pos);
            self.buf[self.buf_len..self.buf_len + take].copy_from_slice(&data[pos..pos + take]);
            self.buf_len += take;
            pos += take;
        }
    }

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; self.out_len];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        self.total += self.buf_len as u64;
        for b in &mut self.buf[self.buf_len..] {
            *b = 0;
        }
        blake2s_compress(&mut self.h, &self.buf, self.total, true);
        let len = out.len().min(self.out_len);
        let mut full = [0u8; 32];
        for i in 0..8 {
            full[4 * i..4 * i + 4].copy_from_slice(&self.h[i].to_le_bytes());
        }
        out[..len].copy_from_slice(&full[..len]);
    }
}

// ====================================================================
// Tests — RFC 7693 Appendix A
// ====================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Hasher;

    fn hex(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{:02x}", b)).collect()
    }

    // RFC 7693 §Appendix A: BLAKE2b-512("abc")
    #[test]
    fn blake2b_abc() {
        let digest = Blake2b::hash(b"abc");
        assert_eq!(
            hex(&digest),
            "ba80a53f981c4d0d6a2797b69f12f6e94c212f14685ac4b74b12bb6fdbffa2d1\
             7d87c5392aab792dc252d5de4533cc9518d38aa8dbf1925ab92386edd4009923"
        );
    }

    // BLAKE2b-512("")
    #[test]
    fn blake2b_empty() {
        let digest = Blake2b::hash(b"");
        assert_eq!(
            hex(&digest),
            "786a02f742015903c6c6fd852552d272912f4740e15847618a86e217f71f5419\
             d25e1031afee585313896444934eb04b903a685b1448b755d56f701afe9be2ce"
        );
    }

    // BLAKE2b with custom output length (32 bytes).
    #[test]
    fn blake2b_32_empty() {
        let mut h = Blake2b::with_output_len(32);
        h.update(b"");
        let digest = h.finalize();
        assert_eq!(digest.len(), 32);
        assert_eq!(
            hex(&digest),
            "0e5751c026e543b2e8ab2eb06099daa1\
             d1e5df47778f7787faab45cdf12fe3a8"
        );
    }

    // RFC 7693 §Appendix A: BLAKE2s-256("abc")
    #[test]
    fn blake2s_abc() {
        let digest = Blake2s::hash(b"abc");
        assert_eq!(
            hex(&digest),
            "508c5e8c327c14e2e1a72ba34eeb452f37458b209ed63a294d999b4c86675982"
        );
    }

    // BLAKE2s-256("")
    #[test]
    fn blake2s_empty() {
        let digest = Blake2s::hash(b"");
        assert_eq!(
            hex(&digest),
            "69217a3079908094e11121d042354a7c1f55b6482ca1a51e1b250dfd1ed0eef9"
        );
    }

    // Streaming: multi-update produces same result.
    #[test]
    fn blake2b_streaming() {
        let mut h = Blake2b::new();
        h.update(b"a");
        h.update(b"b");
        h.update(b"c");
        let stream = h.finalize();
        let oneshot = Blake2b::hash(b"abc");
        assert_eq!(stream, oneshot);
    }

    #[test]
    fn blake2s_streaming() {
        let mut h = Blake2s::new();
        h.update(b"a");
        h.update(b"b");
        h.update(b"c");
        let stream = h.finalize();
        let oneshot = Blake2s::hash(b"abc");
        assert_eq!(stream, oneshot);
    }
}
