// 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 blake2s implements the Blake2s 256, 224, 160, and
// 128 bit hash algorithms
// as defined in IETF RFC 7693.
// Based off:   https://datatracker.ietf.org/doc/html/rfc7693
// Last updated: November 2015
module blake2s

import encoding.binary

// size128 is the size, in bytes, of a Blake2s 128 checksum.
pub const size128 = 16
// size160 is the size, in bytes, of a Blake2s 160 checksum.
pub const size160 = 20
// size224 is the size, in bytes, of a Blake2s 224 checksum.
pub const size224 = 28
// size256 is the size, in bytes, of a Blake2s 256 checksum.
pub const size256 = 32

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

// G rotation constants
const r1 = 16
const r2 = 12
const r3 = 8
const r4 = 7

// 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 = [
	u32(0x6a09e667),
	u32(0xbb67ae85),
	u32(0x3c6ef372),
	u32(0xa54ff53a),
	u32(0x510e527f),
	u32(0x9b05688c),
	u32(0x1f83d9ab),
	u32(0x5be0cd19),
]

// 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            []u32
	m            [16]u32
	t            u64
}

// string makes a formatted string representation of a Digest structure
pub fn (d Digest) str() string {
	return 'blake2s.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}'
}

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

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

// new224 initializes the digest structure for a Blake2s 224 bit hash
pub fn new224() !&Digest {
	return new_digest(size224, []u8{})!
}

// new_pmac224 initializes the digest structure for a Blake2s 224 bit prefix MAC
pub fn new_pmac224(key []u8) !&Digest {
	return new_digest(size224, key)!
}

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

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

// new126 initializes the digest structure for a Blake2s 128 bit hash
pub fn new128() !&Digest {
	return new_digest(size128, []u8{})!
}

// new_pmac128 initializes the digest structure for a Blake2s 128 bit prefix MAC
pub fn new_pmac128(key []u8) !&Digest {
	return new_digest(size128, key)!
}

struct HashSizeError {
	Error
	size u8
}

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

struct KeySizeError {
	Error
	size i32
}

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

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 32.
//
// key - key used for generating a prefix MAC.  A zero length
//     key is used for just generating a hash.  A key of 1 to
//     32 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 > size256 {
		return HashSizeError{
			size: hash_size
		}
	}

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

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

	if key.len > 0 {
		p0 := 0x01010000 ^ u32(key.len) << 8 ^ u32(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 := 0x01010000 ^ u32(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 += 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 u32 values
	for i in 0 .. 16 {
		d.m[i] = binary.little_endian_u32_at(d.input_buffer, i * 4)
	}

	// 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: 4, cap: 4}
		binary.little_endian_put_u32(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) }
}

// sum256 returns the Blake2s 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) }
}

// sum224 returns the Blake2s 224 bit checksum of the data.
pub fn sum224(data []u8) []u8 {
	mut d := new224() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// sum160 returns the Blake2s 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) }
}

// sum128 returns the Blake2s 128 bit checksum of the data.
pub fn sum128(data []u8) []u8 {
	mut d := new128() or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac256 returns the Blake2s 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) }
}

// pmac224 returns the Blake2s 224 bit prefix MAC of the data.
pub fn pmac224(data []u8, key []u8) []u8 {
	mut d := new_pmac224(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}

// pmac160 returns the Blake2s 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) }
}

// pmac128 returns the Blake2s 128 bit prefix MAC of the data.
pub fn pmac128(data []u8, key []u8) []u8 {
	mut d := new_pmac128(key) or { panic(err) }
	d.write(data) or { panic(err) }
	return d.checksum_internal() or { panic(err) }
}