// 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.

// streaming shake-128/256 xof per FIPS 202
// https://nvlpubs.nist.gov/nistpubs/fips/nist.fips.202.pdf

module sha3

@[noinit]
pub struct Shake {
	rate int // bytes per permutation (168 for shake-128, 136 for shake-256)
mut:
	s            State
	input_buffer []u8
	finalized    bool
	squeeze_buf  []u8
}

// new_shake128 returns a new Shake instance for SHAKE-128 extended output function.
pub fn new_shake128() &Shake {
	return &Shake{
		rate: xof_rate_128
	}
}

// new_shake256 returns a new Shake instance for SHAKE-256 extended output function.
pub fn new_shake256() &Shake {
	return &Shake{
		rate: xof_rate_256
	}
}

// write absorbs more data into the sponge state.
// Panics if called after `read`.
@[direct_array_access]
pub fn (mut s Shake) write(data []u8) {
	if s.finalized {
		panic('sha3: write after read on Shake')
	}
	if data.len == 0 {
		return
	}

	// avoid cloning on each iteration
	mut remaining := unsafe { data[..] }

	if s.input_buffer.len != 0 {
		empty_space := s.rate - s.input_buffer.len

		if remaining.len < empty_space {
			s.input_buffer << remaining
			return
		} else {
			s.input_buffer << remaining[..empty_space]
			remaining = unsafe { remaining[empty_space..] }

			s.s.xor_bytes(s.input_buffer[..s.rate], s.rate)
			s.s.kaccak_p_1600_24()

			s.input_buffer = []u8{}
		}
	}

	for remaining.len >= s.rate {
		s.s.xor_bytes(remaining[..s.rate], s.rate)
		s.s.kaccak_p_1600_24()
		remaining = unsafe { remaining[s.rate..] }
	}

	if remaining.len > 0 {
		s.input_buffer = remaining.clone()
	}
}

fn (mut s Shake) finalize() {
	if s.finalized {
		return
	}
	s.finalized = true

	// pad10*1 with xof domain separator 0x1f (FIPS 202 sec B.2)
	mut padded := s.input_buffer.clone()
	if padded.len == s.rate - 1 {
		padded << u8(0x80 | 0x1f)
	} else {
		padded << u8(0x1f)
		for padded.len < s.rate - 1 {
			padded << u8(0x00)
		}
		padded << u8(0x80)
	}

	s.s.xor_bytes(padded[..s.rate], s.rate)
	s.s.kaccak_p_1600_24()

	state_bytes := s.s.to_bytes()
	s.squeeze_buf = state_bytes[..s.rate].clone()
	s.input_buffer = []u8{}
}

// read squeezes `out_len` bytes from the sponge state.
// Finalizes the sponge on first call; further calls to `write` will panic.
@[direct_array_access]
pub fn (mut s Shake) read(out_len int) []u8 {
	if !s.finalized {
		s.finalize()
	}

	mut result := []u8{cap: out_len}
	mut remaining := out_len

	for remaining > 0 {
		if s.squeeze_buf.len == 0 {
			s.s.kaccak_p_1600_24()
			state_bytes := s.s.to_bytes()
			s.squeeze_buf = state_bytes[..s.rate].clone()
		}

		take := if remaining < s.squeeze_buf.len { remaining } else { s.squeeze_buf.len }
		result << s.squeeze_buf[..take]
		s.squeeze_buf = s.squeeze_buf[take..].clone()
		remaining -= take
	}

	return result
}

// reset clears the sponge state, allowing the Shake instance to be reused.
pub fn (mut s Shake) reset() {
	s.s = State{}
	s.input_buffer = []u8{}
	s.finalized = false
	s.squeeze_buf = []u8{}
}