// Copyright (c) 2023 Kim Shrier. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
// Package blake2b implements the Blake2b 512, 384, 256, and
// 160 bit hash algorithms
// as defined in IETF RFC 7693.
// Based off:   https://datatracker.ietf.org/doc/html/rfc7693
// Last updated: November 2015
module blake2b

import encoding.binary
import math.unsigned

// size160 is the size, in bytes, of a Blake2b 160 checksum.
pub const size160 = 20
// size256 is the size, in bytes, of a Blake2b 256 checksum.
pub const size256 = 32
// size384 is the size, in bytes, of a Blake2b 384 checksum.
pub const size384 = 48
// size512 is the size, in bytes, of a Blake2b 512 checksum.
pub const size512 = 64

// block_size is the block size, in bytes, of the Blake2b hash functions.
pub const block_size = 128

// G rotation constants
const r1 = 32
const r2 = 24
const r3 = 16
const r4 = 63

// negative G rotation constants so we can rotate right.
const nr1 = -1 * r1
const nr2 = -1 * r2
const nr3 = -1 * r3
const nr4 = -1 * r4

// initialization vector
const iv = [
	u64(0x6a09e667f3bcc908),
	u64(0xbb67ae8584caa73b),
	u64(0x3c6ef372fe94f82b),
	u64(0xa54ff53a5f1d36f1),
	u64(0x510e527fade682d1),
	u64(0x9b05688c2b3e6c1f),
	u64(0x1f83d9abfb41bd6b),
	u64(0x5be0cd19137e2179),
]

// message word schedule permutations
const sigma = [
	[u8(0), 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
	[u8(14), 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
	[u8(11), 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
	[u8(7), 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
	[u8(9), 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
	[u8(2), 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
	[u8(12), 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
	[u8(13), 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
	[u8(6), 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
	[u8(10), 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0],
]

struct Digest {
	hash_size u8
mut:
	input_buffer []u8
	h            []u64
	m            [16]u64
	t            unsigned.Uint128
}

// string makes a formatted string representation of a Digest structure
pub fn (d Digest) str() string {
	return 'blake2b.Digest{\n    hash_size: ${d.hash_size}\n    input_buffer: ${d.input_buffer}\n    input_buffer.len: ${d.input_buffer.len}\n    h: [0x${d.h[0]:016x}, 0x${d.h[1]:016x}, 0x${d.h[2]:016x}, 0x${d.h[3]:016x},\n        0x${d.h[4]:016x}, 0x${d.h[5]:016x}, 0x${d.h[6]:016x}, 0x${d.h[7]:016x}]\n    m: [0x${d.m[0]:016x}, 0x${d.m[1]:016x}, 0x${d.m[2]:016x}, 0x${d.m[3]:016x},\n        0x${d.m[4]:016x}, 0x${d.m[5]:016x}, 0x${d.m[6]:016x}, 0x${d.m[7]:016x},\n        0x${d.m[8]:016x}, 0x${d.m[9]:016x}, 0x${d.m[10]:016x}, 0x${d.m[11]:016x},\n        0x${d.m[12]:016x}, 0x${d.m[13]:016x}, 0x${d.m[14]:016x}, 0x${d.m[15]:016x}]\n    t: ${d.t}\n}'
}

// new512 initializes the digest structure for a Blake2b 512 bit hash
pub fn new512() !&Digest {
	return new_digest(size512, []u8{})!
}

// new_pmac512 initializes the digest structure for a Blake2b 512 bit prefix MAC
pub fn new_pmac512(key []u8) !&Digest {
	return new_digest(size512, key)!
}

// new384 initializes the digest structure for a Blake2b 384 bit hash
pub fn new384() !&Digest {
	return new_digest(size384, []u8{})!
}

// new_pmac384 initializes the digest structure for a Blake2b 384 bit prefix MAC
pub fn new_pmac384(key []u8) !&Digest {
	return new_digest(size384, key)!
}

// new256 initializes the digest structure for a Blake2b 256 bit hash
pub fn new256() !&Digest {
	return new_digest(size256, []u8{})!
}

// new_pmac256 initializes the digest structure for a Blake2b 256 bit prefix MAC
pub fn new_pmac256(key []u8) !&Digest {
	return new_digest(size256, key)!
}

// new160 initializes the digest structure for a Blake2b 160 bit hash
pub fn new160() !&Digest {
	return new_digest(size160, []u8{})!
}

// new_pmac160 initializes the digest structure for a Blake2b 160 bit prefix MAC
pub fn new_pmac160(key []u8) !&Digest {
	return new_digest(size160, key)!
}

struct HashSizeError {
	Error
	size u8
}

fn (err HashSizeError) msg() string {
	return 'Hash size ${err.size} must be between 1 and ${size512}'
}

struct KeySizeError {
	Error
	size i32
}

fn (err KeySizeError) msg() string {
	return 'Key size ${err.size} must be between 0 and ${size512}'
}

struct InputBufferSizeError {
	Error
	size i32
}

fn (err InputBufferSizeError) msg() string {
	return 'The input buffer size ${err.size} .must be between 0 and ${block_size}'
}

// new_digest creates an initialized digest structure based on
// the hash size and whether or not you specify a MAC key.
//
// hash_size - the number of bytes in the generated hash.
//     Legal values are between 1 and 64.
//
// key - key used for generating a prefix MAC.  A zero length
//     key is used for just generating a hash.  A key of 1 to
//     64 bytes can be used for generating a prefix MAC.
pub fn new_digest(hash_size u8, key []u8) !&Digest {
	if hash_size < 1 || hash_size > size512 {
		return HashSizeError{
			size: hash_size
		}
	}

	if key.len < 0 || key.len > size512 {
		return KeySizeError{
			size: i32(key.len)
		}
	}

	mut d := Digest{
		h:         iv.clone()
		t:         unsigned.uint128_zero
		hash_size: hash_size
	}

	if key.len > 0 {
		p0 := 0x0000000001010000 ^ u64(key.len) << 8 ^ u64(hash_size)

		d.h[0] ^= p0

		d.input_buffer.clear()
		d.input_buffer << key

		pad_length := block_size - key.len
		for _ in 0 .. pad_length {
			d.input_buffer << 0
		}
	} else {
		p0 := 0x0000000001010000 ^ u64(hash_size)

		d.h[0] ^= p0
	}

	return &d
}

// The intent is to only use this method
// when the input buffer is full or when
// the last data to be hashed is in the
// input buffer.
//
// If the input buffer does not contain enough
// data to fill all the message blocks, the
// input buffer is padded with 0 bytes until
// it is full prior to moving the data into
// the message blocks.
//
// The input buffer is emptied.
//
// The total byte count is incremented by the
// number of bytes in the input buffer.
fn (mut d Digest) move_input_to_message_blocks() ! {
	// the number of bytes in the input buffer
	// should never exceed the block size.
	if d.input_buffer.len < 0 || d.input_buffer.len > block_size {
		return InputBufferSizeError{
			size: i32(d.input_buffer.len)
		}
	}

	// keep the hashed data length up to date
	d.t = d.t.add_64(u64(d.input_buffer.len))

	// pad the input buffer if necessary
	if d.input_buffer.len < block_size {
		pad_length := block_size - d.input_buffer.len

		for _ in 0 .. pad_length {
			d.input_buffer << 0
		}
	}

	// treat the input bytes as little endian u64 values
	for i in 0 .. 16 {
		d.m[i] = binary.little_endian_u64_at(d.input_buffer, i * 8)
	}

	// empty the input buffer
	d.input_buffer.clear()

	return
}

// write adds bytes to the hash
pub fn (mut d Digest) write(data []u8) ! {
	// if no data is being added to the hash,
	// just return
	if data.len == 0 {
		return
	}

	// if the input buffer is already full,
	// process the existing input bytes first.
	// this is not the final input.
	if d.input_buffer.len >= block_size {
		d.move_input_to_message_blocks()!
		d.f(false)
	}

	// add data to the input buffer until you
	// run out of space in the input buffer or
	// run out of data, whichever comes first.
	empty_space := block_size - d.input_buffer.len
	mut remaining_data := unsafe { data[..] }

	if empty_space >= data.len {
		// ran out of data first
		// just add it to the input buffer and return
		d.input_buffer << data
		return
	} else {
		// ran out of input buffer space first
		// fill it up and process it.
		d.input_buffer << data[..empty_space]
		remaining_data = unsafe { remaining_data[empty_space..] }

		d.move_input_to_message_blocks()!
		d.f(false)
	}

	// process the data in block size amounts until
	// all the data has been processed
	for remaining_data.len > 0 {
		if block_size >= remaining_data.len {
			// running out of data
			// just add it to the input buffer and return
			d.input_buffer << remaining_data
			return
		}

		// add block size bytes to the input buffer
		// and process it
		d.input_buffer << remaining_data[..block_size]
		remaining_data = unsafe { remaining_data[block_size..] }

		d.move_input_to_message_blocks()!
		d.f(false)
	}

	return
}

fn (mut d Digest) checksum_internal() ![]u8 {
	// process the last input bytes
	d.move_input_to_message_blocks()!
	d.f(true)

	// get the hash into the proper byte order
	mut hash_bytes := []u8{}

	for hash in d.h {
		mut h_bytes := []u8{len: 8, cap: 8}
		binary.little_endian_put_u64(mut h_bytes, hash)
		hash_bytes << h_bytes
	}

	// return the appropriate number of hash bytes
	return hash_bytes[..d.hash_size]
}

// checksum finalizes the hash and returns the generated bytes.
pub fn (mut d Digest) checksum() []u8 {
	return d.checksum_internal() or { panic(err) }
}

// sum512 returns the Blake2b 512 bit checksum of the data.
pub fn sum512(data []u8) []u8 {
	mut d := new512() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// sum384 returns the Blake2b 384 bit checksum of the data.
pub fn sum384(data []u8) []u8 {
	mut d := new384() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// sum256 returns the Blake2b 256 bit checksum of the data.
pub fn sum256(data []u8) []u8 {
	mut d := new256() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// sum160 returns the Blake2b 160 bit checksum of the data.
pub fn sum160(data []u8) []u8 {
	mut d := new160() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac512 returns the Blake2b 512 bit prefix MAC of the data.
pub fn pmac512(data []u8, key []u8) []u8 {
	mut d := new_pmac512(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac384 returns the Blake2b 384 bit prefix MAC of the data.
pub fn pmac384(data []u8, key []u8) []u8 {
	mut d := new_pmac384(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac256 returns the Blake2b 256 bit prefix MAC of the data.
pub fn pmac256(data []u8, key []u8) []u8 {
	mut d := new_pmac256(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac160 returns the Blake2b 160 bit prefix MAC of the data.
pub fn pmac160(data []u8, key []u8) []u8 {
	mut d := new_pmac160(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}