// Copyright (c) 2025 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// This file contains an experimental port of a Rust reference implementation of nonce-misuse
// resistant and key-committing authenticated encryption scheme called ChaCha20-Poly1305-PSIV,
// It backed by `chacha20` stream cipher and `poly1305` message authentication code module.
// Its originally described by Michiel Verbauwhede and the teams on his papers.
// See the detail on the [A Robust Variant of ChaCha20-Poly1305](https://eprint.iacr.org/2025/222).
module chacha20poly1305

import crypto.cipher
import encoding.binary
import crypto.internal.subtle
import x.crypto.chacha20
import x.crypto.poly1305

// new_psiv creates a new Chacha20Poly1305RE with PSIV construct to operate on.
@[direct_array_access]
pub fn new_psiv(key []u8) !&Chacha20Poly1305RE {
	if key.len != key_size {
		return error('new_psiv: bad key size')
	}
	// derives and initializes the new key for later purposes
	pol_key := fk_k(key)
	mac_key := fm_k(key)
	enc_key := fe_k(key)

	mut s := chacha20.State{}
	mut x64 := [64]u8{}
	unsafe { vmemcpy(x64, pol_key[0], 36) }
	unpack_into_state(mut s, x64)
	ws := chacha20_core(s)

	// For poly1305 mac, we only take a first 32-bytes of the state as a key
	mut poly1305_key := []u8{len: 32}
	pack32_from_state(mut poly1305_key, ws)
	po := poly1305.new(poly1305_key)!

	// set the values
	c := &Chacha20Poly1305RE{
		mac_key: mac_key
		enc_key: enc_key
		po:      po
	}
	return c
}

// psiv_encrypt encrypts plaintext with provided key, nonce and additional data ad.
// It returns a ciphertext plus message authentication code (mac) contained
// within the end of ciphertext
@[direct_array_access]
pub fn psiv_encrypt(plaintext []u8, key []u8, nonce []u8, ad []u8) ![]u8 {
	mut c := new_psiv(key)!
	out := c.encrypt(plaintext, nonce, ad)!
	unsafe { c.free() }
	return out
}

// psiv_decrypt decrypts the ciphertext with provided key, nonce and additional data in ad.
// It also tries to validate message authentication code within ciphertext compared with
// calculated tag. It returns successfully decrypted message or error on fails.
@[direct_array_access]
pub fn psiv_decrypt(ciphertext []u8, key []u8, nonce []u8, ad []u8) ![]u8 {
	mut c := new_psiv(key)!
	out := c.decrypt(ciphertext, nonce, ad)!
	unsafe { c.free() }
	return out
}

// Chacha20Poly1305RE is a Chacha20Poly1305 opaque with nonce-misuse resistent
// and key-commiting AEAD scheme with PSIV construct.
@[noinit]
pub struct Chacha20Poly1305RE implements cipher.AEAD {
mut:
	// flag that marked this instance should not be used again, set on .free call
	done bool
	// underlying derived keys, set on instance creation with new_psiv.
	mac_key [36]u8
	enc_key [36]u8
	po      &poly1305.Poly1305 = unsafe { nil }
}

// free releases resources taken by c. Dont use c after `.free` call.
@[unsafe]
pub fn (mut c Chacha20Poly1305RE) free() {
	// if it already marked as done, just return
	if c.done {
		return
	}
	unsafe {
		// we reset derived keys
		vmemset(c.mac_key, 0, c.mac_key.len)
		vmemset(c.enc_key, 0, c.enc_key.len)
		c.po = nil
	}
	// mark this instance as done, no longer usable
	c.done = true
}

// nonce_size return the size of the nonce of underlying c.
// Currently, it only support for standard 12-bytes nonce.
pub fn (c &Chacha20Poly1305RE) nonce_size() int {
	return nonce_size
}

// overhead returns difference between the lengths of a plaintext and its ciphertext.
// Its normally returns a tag size produced by this scheme.
pub fn (c &Chacha20Poly1305RE) overhead() int {
	return tag_size
}

// encrypt encrypts and authenticates the provided plaintext along with a nonce, and
// to be authenticated additional data in `ad`. It returns a ciphertext with message authenticated
// code stored within the end of ciphertext.
@[direct_array_access]
pub fn (c Chacha20Poly1305RE) encrypt(plaintext []u8, nonce []u8, ad []u8) ![]u8 {
	if c.done {
		return error('Chacha20Poly1305RE.encrypt: instance marked as done, no longer usable')
	}
	if nonce.len != nonce_size {
		return error('Chacha20Poly1305RE.encrypt: bad nonce length, only support 12-bytes nonce')
	}

	// clone the initial poly1305 state and updates it with additional data ad
	mut po_ad := c.po.clone()
	update_with_padding(mut po_ad, ad)
	// make a clone of updated poly1305
	mut po_ad_clone := po_ad.clone()

	// setup output buffer
	mut out := []u8{len: plaintext.len + tag_size}
	// write out an authentication tag into the last tag_size bytes of output
	psiv_gen_tag(mut out[plaintext.len..], mut po_ad_clone, plaintext, ad.len, c.mac_key, nonce)
	// write out authenticated encrypted plaintext into the first plaintext.len bytes of output
	psiv_encrypt_internal(mut out[0..plaintext.len], plaintext, c.enc_key, out[plaintext.len..],
		nonce)!

	// return the result
	return out
}

// decrypt decrypts the ciphertext with provided key, nonce and additional data in ad.
// It also tries to validate message authenticated code within ciphertext compared with
// calculated tag. It returns successfully decrypted message or error on fails.
@[direct_array_access]
pub fn (c Chacha20Poly1305RE) decrypt(ciphertext []u8, nonce []u8, ad []u8) ![]u8 {
	if c.done {
		return error('Chacha20Poly1305RE.decrypt: instance marked as done, no longer usable')
	}
	if ciphertext.len < tag_size {
		return error('Chacha20Poly1305RE.decrypt: insufficient ciphertext length')
	}
	if nonce.len != nonce_size {
		return error('Chacha20Poly1305RE.decrypt: invalid nonce length provided')
	}
	enc := ciphertext[0..ciphertext.len - c.overhead()]
	tag := ciphertext[ciphertext.len - c.overhead()..]

	// updates a clone of poly1305 with additional data
	mut po_with_ad := c.po.clone()
	update_with_padding(mut po_with_ad, ad)
	mut poad_clone := po_with_ad.clone()

	// generates authenticated encrypted plaintext with associated mac
	mut out := []u8{len: enc.len}
	psiv_encrypt_internal(mut out, enc, c.enc_key, tag, nonce)!

	mut mac := []u8{len: tag_size}
	psiv_gen_tag(mut mac, mut poad_clone, out, ad.len, c.mac_key, nonce)

	// check if authentication tag was matching or error on fails.
	if subtle.constant_time_compare(mac, tag) != 1 {
		unsafe {
			out.free()
			mac.free()
		}
		return error('unmatching tag')
	}
	// return the decrypted ciphertext
	return out
}

// The AEAD_CHACHA20_POLY1305 PSIV construct helpers
//

// psiv_encrypt_internal is an internal encryption routine used by the core of psiv construct
// for encrypting (or decrypting) message.
@[direct_array_access]
fn psiv_encrypt_internal(mut dst []u8, plaintext []u8, dkey [36]u8, tag []u8, nonce []u8) ! {
	// loads the counter from the first 8-bytes of the tag input
	mut ctr := binary.little_endian_u64(tag[0..8])

	// setup some temporary vars
	mut tc := []u8{len: 8} // counter buffer
	mut s := chacha20.State{}
	mut b64 := [64]u8{} // state buffer
	mut tt := merge_drv_key(dkey, nonce, tag[0..8], tag[8..16])

	mut j := 0
	mut n := 0

	// process for every bytes on plaintext input
	for plaintext[n..].len > 0 {
		// how many block of bytes available to process on
		want_len := if plaintext[n..].len < 64 { plaintext[n..].len } else { 64 }
		// loads current counter
		binary.little_endian_put_u64(mut tc, ctr)

		// updates derived keys with current counter, scrambled with chacha20_core and
		// puts state into b64 buffer
		unsafe { vmemcpy(tt[48], tc.data, tc.len) }
		unpack_into_state(mut s, tt)
		ws := chacha20_core(s)
		pack64_from_state(mut b64, ws)

		// xor every bytes of plaintext with bytes on b64, stores result in dst
		for i in 0 .. want_len {
			dst[j] = plaintext[j] ^ b64[i]
			j++
		}
		// updates current counter and returns error on overflow.
		ctr += 1
		if ctr == 0 {
			return error('counter overflowing')
		}
		n += want_len
	}
	// explicitly reset (release) temporary allocated resources and return the result.
	unsafe {
		tc.free()
		s.reset()
		vmemset(b64, 0, b64.len)
		vmemset(tt, 0, tt.len)
	}
}

// psiv_gen_tag computes a tag from the key, nonce, and Poly1305 tag of the associated data
// and plaintext using the ChaCha20 permutation with the feed-forward, truncating the output.
@[direct_array_access]
fn psiv_gen_tag(mut out []u8, mut po poly1305.Poly1305, input []u8, ad_len int, mac_key [36]u8, nonce []u8) {
	// updates poly1305 mac by input message, associated data length and input length.
	update_with_padding(mut po, input)
	po.update(length_to_block(ad_len, input.len))

	// produces 16-bytes of mac from current poly1305 state.
	po.finish(mut out)

	// The tag was produced from derived key scrambled with chacha20 quarter round routine,
	// and then truncating the output into 16-bytes tag.
	drv_key := merge_drv_key(mac_key, nonce, out[0..8], out[8..16])
	mut x := chacha20.State{}
	unpack_into_state(mut x, drv_key)
	ws := chacha20_core(x)

	// truncating state output into tag_sized bytes. As a note, we reuse buffer previously allocated
	// to store the result.
	pack16_from_state(mut out, ws)

	// explicitly releases (reset) temporary allocated resources
	unsafe {
		vmemset(drv_key, 0, drv_key.len)
		ws.reset()
		x.reset()
	}
}

// update_with_padding updates poly1305 mac with data, padding the tail block if necessary.
@[direct_array_access; inline]
fn update_with_padding(mut po poly1305.Poly1305, data []u8) {
	po.update(data)
	rem := data.len % tag_size
	if rem != 0 {
		block := []u8{len: tag_size}
		po.update(block[..tag_size - rem])
	}
}

// merge_drv_key merges provided bytes into 64-bytes key
@[direct_array_access; inline]
fn merge_drv_key(dkey [36]u8, nonce []u8, tag_ctr []u8, tag_rest []u8) [64]u8 {
	mut x64 := [64]u8{}

	// 0..36
	for i := 0; i < dkey.len; i++ {
		x64[i] = dkey[i]
	}
	// 36..48
	for i := 0; i < nonce.len; i++ {
		x64[36 + i] = nonce[i]
	}
	// 48..56
	for i := 0; i < tag_ctr.len; i++ {
		x64[i + 48] = tag_ctr[i]
	}
	// 56..64
	for i := 0; i < tag_rest.len; i++ {
		x64[i + 56] = tag_rest[i]
	}

	return x64
}

// fk_k maps and transforms 32-bytes of key into 36-bytes of new key used to
// derive a poly1305 construction.
// See the papers doc on the 3.3 Additional Details part, on page 12-13
@[direct_array_access; inline]
fn fk_k(key []u8) [36]u8 {
	// fk(K) = 	K1 ∥ K2 ∥ K3 ∥ 03 ∥ K5 ∥ K6 ∥ K7 ∥ 0c ∥ K9 ∥ K10 ∥ K11 ∥ 30
	//			∥ K4 ∥ K8 ∥ K12 ∥ c0 ∥ K13 ∥ K14 ∥  · · ·  ∥ K32
	// with 0-based index
	// 			K0 ∥ K1 ∥ K2 ∥ 03 ∥ K4 ∥ K5 ∥ K6 ∥ 0c ∥ K8 ∥ K9 ∥ K10 ∥ 30
	//			∥ K3 ∥ K7 ∥ K11 ∥ c0 ∥ K12 ∥ K13 ∥  · · ·  ∥ K31
	mut x := [36]u8{}
	// 0 .. 4
	for i := 0; i < 3; i++ {
		x[i] = key[i]
	}
	x[3] = u8(0x03)

	// 4 .. 8
	for i := 4; i < 7; i++ {
		x[i] = key[i]
	}
	x[7] = 0x0c

	// 8 .. 12
	for i := 8; i < 11; i++ {
		x[i] = key[i]
	}
	x[11] = 0x30

	// 12 .. 16
	x[12] = key[3]
	x[13] = key[7]
	x[14] = key[11]
	x[15] = 0xc0

	// 16 .. 36
	for i := 16; i < 36; i++ {
		x[i] = key[i - 4]
	}

	return x
}

// fm_k maps and transforms 32-bytes of key into 36-bytes of message authentication key.
// It later used for psiv tag generation.
@[direct_array_access; inline]
fn fm_k(key []u8) [36]u8 {
	// fm(K) = 	K1 ∥ K2 ∥ K3 ∥ 05 ∥ K5 ∥ K6 ∥ K7 ∥ 0a ∥ K9 ∥ K10 ∥ K11 ∥ 50 ∥
	//			K4 ∥ K8 ∥ K12 ∥ a0 ∥ K13 ∥ K14 ∥  · · ·  ∥ K32 ,
	// Or, with 0-based index
	// fm(K) = 	K0 ∥ K1 ∥ K2 ∥ 05 ∥ K4 ∥ K5 ∥ K6 ∥ 0a ∥ K8 ∥ K9 ∥ K10 ∥ 50 ∥
	//			K3 ∥ K7 ∥ K11 ∥ a0 ∥ K12 ∥ K13 ∥  · · ·  ∥ K31
	mut x := [36]u8{}
	// 0 .. 4
	for i := 0; i < 3; i++ {
		x[i] = key[i]
	}
	x[3] = u8(0x05)

	// 4 .. 8
	for i := 4; i < 7; i++ {
		x[i] = key[i]
	}
	x[7] = 0x0a

	// 8 .. 12
	for i := 8; i < 11; i++ {
		x[i] = key[i]
	}
	x[11] = 0x50

	// 12 .. 16
	x[12] = key[3]
	x[13] = key[7]
	x[14] = key[11]
	x[15] = 0xa0

	// 16 .. 36
	for i := 16; i < 36; i++ {
		x[i] = key[i - 4]
	}

	return x
}

// fe_k maps and transforms 32-bytes of key into 36-bytes of new encryption key
@[direct_array_access; inline]
fn fe_k(key []u8) [36]u8 {
	// fe(K) = 	K1 ∥ K2 ∥ K3 ∥ 06 ∥ K5 ∥ K6 ∥ K7 ∥ 09 ∥ K9 ∥ K10 ∥ K11 ∥ 60 ∥
	// 			K4 ∥ K8 ∥ K12 ∥ 90 ∥ K13 ∥ K14 ∥  · · ·  ∥ K32
	// Or, with 0-based index
	// fe(K) = 	K0 ∥ K1 ∥ K2 ∥ 06 ∥ K4 ∥ K5 ∥ K6 ∥ 09 ∥ K8 ∥ K9 ∥ K10 ∥ 60 ∥
	// 			K3 ∥ K7 ∥ K11 ∥ 90 ∥ K12 ∥ K13 ∥  · · ·  ∥ K31
	mut x := [36]u8{}
	// 0 .. 4
	for i := 0; i < 3; i++ {
		x[i] = key[i]
	}
	x[3] = u8(0x06)

	// 4 .. 8
	for i := 4; i < 7; i++ {
		x[i] = key[i]
	}
	x[7] = 0x09

	// 8 .. 12
	for i := 8; i < 11; i++ {
		x[i] = key[i]
	}
	x[11] = 0x60

	// 12 .. 16
	x[12] = key[3]
	x[13] = key[7]
	x[14] = key[11]
	x[15] = 0x90

	// 16 .. 36
	for i := 16; i < 36; i++ {
		x[i] = key[i - 4]
	}

	return x
}

// unpack_into_state deserializes (in little-endian form) 64-bytes of data in x into state s.
@[direct_array_access; inline]
fn unpack_into_state(mut s chacha20.State, x [64]u8) {
	for i := 0; i < 16; i++ {
		s[i] = u32(x[i * 4]) | (u32(x[i * 4 + 1]) << u32(8)) | (u32(x[i * 4 + 2]) << u32(16)) | (u32(x[
			i * 4 + 3]) << u32(24))
	}
}

// pack64_from_state serializes state s into 64-bytes output in little-endian form.
@[direct_array_access; inline]
fn pack64_from_state(mut out [64]u8, s chacha20.State) {
	mut j := 0
	for v in s {
		out[j] = u8(v)
		out[j + 1] = u8(v >> u32(8))
		out[j + 2] = u8(v >> u32(16))
		out[j + 3] = u8(v >> u32(24))
		j += 4
	}
}

// pack32_from_state serializes only a half of state s into 32-bytes output in little-endian form.
@[direct_array_access; inline]
fn pack32_from_state(mut out []u8, s chacha20.State) {
	mut j := 0
	for i in 0 .. 8 {
		out[j] = u8(s[i])
		out[j + 1] = u8(s[i] >> u32(8))
		out[j + 2] = u8(s[i] >> u32(16))
		out[j + 3] = u8(s[i] >> u32(24))
		j += 4
	}
}

// pack16_from_state serializes the first quartet of state s into 16-bytes output in little-endian form.
@[direct_array_access; inline]
fn pack16_from_state(mut out []u8, s chacha20.State) {
	mut j := 0
	for i in 0 .. 4 {
		out[j] = u8(s[i])
		out[j + 1] = u8(s[i] >> u32(8))
		out[j + 2] = u8(s[i] >> u32(16))
		out[j + 3] = u8(s[i] >> u32(24))
		j += 4
	}
}

// chacha20_core performs chacha20 quarter round on the state s.
// It returns a copy of updated state after quarter round.
@[direct_array_access; inline]
fn chacha20_core(s chacha20.State) chacha20.State {
	mut ws := s.clone()
	ws.qround(10)
	for i := 0; i < 16; i++ {
		ws[i] += s[i]
	}
	return ws
}

// length_to_block transforms two's length in len1 and len2 into 16-bytes block
@[inline]
fn length_to_block(len1 int, len2 int) []u8 {
	mut block := []u8{len: 16}
	binary.little_endian_put_u64(mut block[0..8], u64(len1))
	binary.little_endian_put_u64(mut block[8..16], u64(len2))

	return block
}