// Copyright (c) 2026 blackshirt. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// AES-GCM (NIST SP 800-38D) built on top of the pure-V `crypto.aes` block
// cipher. Its build to support for QUIC protocol and another cases.
// QUIC (RFC 9001) mandates AES-128-GCM for Initial packets and allows
// AES-128-GCM / AES-256-GCM for 1-RTT packets. The GHASH multiplication uses
// the straightforward bit-by-bit algorithm, which is constant in structure and
// easy to audit (performance is adequate for handshakes and modest streams).
module aes

import crypto.cipher

// gcm_tag_size is the size of the GCM authentication tag in bytes.
pub const gcm_tag_size = 16

// gcm_nonce_size is the size of the GCM nonce in bytes (only support 96-bit nonce variant).
pub const gcm_nonce_size = 12

// zero_block is a preallocated 16-bytes of zeros block. Its currently sized by Aes block size.
const zero_block = []u8{len: block_size}

// maximum plaintext length per invocation GCM: MUST be <= 2^36 - 32 octets.
const max_plaintext_size = u64(1) << 36 - 32

// maximum ciphertext length per invocation GCM: MUST be <= 2^36 - 16 octets.
const max_ciphertext_size = u64(1) << 36 - 16

// maximum length of additional authenticated data. Even technically its was not limited,
// we give it a limit to reduce the risk. Intended purposes of additional data was
// as an additional authenticated data included into output, not act as a primary input.
const max_aad_size = max_u32

// AesGcm holds the AES block cipher and the derived GHASH subkey for a key.
@[heap; noinit]
pub struct AesGcm implements cipher.AEAD {
mut:
	block cipher.Block
	h     []u8 // GHASH subkey H = E(K, 0^128)
}

// new_aes_gcm builds AES-GCM state for a 16, 24 or 32-byte key.
@[direct_array_access]
pub fn new_aes_gcm(key []u8) !&AesGcm {
	if key.len != 16 && key.len != 24 && key.len != 32 {
		return error('AES-GCM key must be 16, 24 or 32 bytes, got ${key.len}')
	}
	block := new_cipher(key)
	mut h := zero_block.clone()
	block.encrypt(mut h, zero_block)
	return &AesGcm{
		block: block
		h:     h
	}
}

// nonce_size returns the size of nonce (in bytes).
pub fn (g &AesGcm) nonce_size() int {
	return gcm_nonce_size
}

// overhead returns the maximum difference between the lengths of a plaintext and its ciphertext.
pub fn (g &AesGcm) overhead() int {
	return gcm_tag_size
}

// encrypt encrypts `plaintext` with the given 12-byte `nonce` and additional
// authenticated data `ad`, returning ciphertext with the 16-byte tag appended.
@[direct_array_access]
pub fn (g &AesGcm) encrypt(plaintext []u8, nonce []u8, ad []u8) ![]u8 {
	// In V language, unlike C and Go, int is always a 32 bit integer, so technically safe
	// to assume if the arrays length not would exceed 32-bit integer (its maybe changed on the future).
	// For safety purposes, we check agains the limit.
	if u64(plaintext.len) > max_plaintext_size {
		return error('AES-GCM encrypt: plaintext size exceed the limit')
	}
	if u64(ad.len) > max_aad_size {
		return error('AES-GCM encrypt: additional data size exceed the limit')
	}
	if nonce.len != gcm_nonce_size {
		return error('AES-GCM nonce must be ${gcm_nonce_size} bytes')
	}
	j0 := j0_from_nonce(nonce)
	mut ctr := j0.clone()
	inc32(mut ctr)
	ciphertext := g.gctr(ctr, plaintext)

	mut ghash_in := pad_block(ad)
	ghash_in << pad_block(ciphertext)
	ghash_in << len_block(ad.len, ciphertext.len)
	s := g.ghash(zero_block, ghash_in)
	tag := g.gctr(j0, s)

	mut out := []u8{cap: ciphertext.len + gcm_tag_size}
	out << ciphertext
	out << tag
	return out
}

// decrypt verifies and decrypts `ciphertext` (which must include the trailing
// 16-byte tag) using `nonce` and additional authenticated data `ad`.
@[direct_array_access]
pub fn (g &AesGcm) decrypt(ciphertext []u8, nonce []u8, ad []u8) ![]u8 {
	// Check the array size for safety
	if u64(ciphertext.len) > max_ciphertext_size {
		return error('AES-GCM decrypt: ciphertext size exceed the limit')
	}
	if u64(ad.len) > max_aad_size {
		return error('AES-GCM decrypt: additional data size exceed the limit')
	}
	if nonce.len != gcm_nonce_size {
		return error('AES-GCM nonce must be ${gcm_nonce_size} bytes')
	}
	if ciphertext.len < gcm_tag_size {
		return error('AES-GCM ciphertext shorter than tag')
	}
	ct := ciphertext[..ciphertext.len - gcm_tag_size]
	tag := ciphertext[ciphertext.len - gcm_tag_size..]
	j0 := j0_from_nonce(nonce)

	mut ghash_in := pad_block(ad)
	ghash_in << pad_block(ct)
	ghash_in << len_block(ad.len, ct.len)
	s := g.ghash(zero_block, ghash_in)
	expected := g.gctr(j0, s)
	if !bytes_equal(expected, tag) {
		return error('AES-GCM authentication failed')
	}
	mut ctr := j0.clone()
	inc32(mut ctr)
	return g.gctr(ctr, ct)
}

// ghash computes GHASH_H over `data`, which the caller must already have padded
// to a multiple of the block size, starting from the running value `y`.
@[direct_array_access]
fn (g &AesGcm) ghash(y []u8, data []u8) []u8 {
	mut acc := y.clone()
	mut off := 0
	for off < data.len {
		for j in 0 .. block_size {
			acc[j] ^= data[off + j]
		}
		acc = gf_mult(acc, g.h)
		off += block_size
	}
	return acc
}

// gctr applies the GCM counter mode to `input` starting from counter block
// `icb`, returning the keystream-XORed output.
@[direct_array_access]
fn (g &AesGcm) gctr(icb []u8, input []u8) []u8 {
	if input.len == 0 {
		return []u8{}
	}
	mut out := []u8{len: input.len}
	mut ctr := icb.clone()
	mut ks := zero_block.clone()
	mut off := 0
	for off < input.len {
		g.block.encrypt(mut ks, ctr)
		n := if input.len - off < block_size { input.len - off } else { block_size }
		for j in 0 .. n {
			out[off + j] = input[off + j] ^ ks[j]
		}
		inc32(mut ctr)
		off += block_size
	}
	return out
}

// Helpers

// gf_mult multiplies two 16-byte blocks in GF(2^128) using the reduction
// polynomial from NIST SP 800-38D (the bit-reflected x^128 + x^7 + x^2 + x + 1).
@[direct_array_access]
fn gf_mult(x []u8, y []u8) []u8 {
	mut z := zero_block.clone()
	mut v := y.clone()

	// Process all 128 bits
	for i in 0 .. 128 {
		// Extract bit i from x (MSB first, left to right)
		// byte index: i >> 3 or (i / 8)
		// bit position: 7 - (i & 7) (7, 6, 5, ..., 0)
		bit := (x[i >> 3] >> (7 - u8(i & 7))) & 1

		// Create constant-time mask
		// If bit = 1: mask = 0xFF (all 1s)
		// If bit = 0: mask = 0x00 (all 0s)
		// This is constant-time: -(i8(bit)) produces correct result
		mask := u8(-(i8(bit)))

		// Conditional XOR: z ^= v * bit (constant-time)
		// Instead of: if bit == 1 { z[j] ^= v[j] }
		// We do:      z[j] ^= (v[j] & mask)
		for j in 0 .. block_size {
			z[j] ^= (v[j] & mask)
		}

		// Right shift v by 1 bit (always performed)
		// This is constant-time - no branches
		lsb := v[15] & 1 // Save LSB before shift

		for j := 15; j > 0; j-- {
			v[j] = (v[j] >> 1) | ((v[j - 1] & 1) << 7)
		}
		v[0] >>= 1

		// Polynomial reduction (constant-time)
		// If lsb = 1: v[0] ^= 0xe1
		// If lsb = 0: v[0] unchanged
		// Polynomial: x^128 + x^7 + x^2 + x + 1 = 0xe1 in lowest bits
		//
		// Instead of: if lsb == 1 { v[0] ^= 0xe1 }
		// We do:      v[0] ^= (0xe1 & mask_for_lsb)
		red_mask := u8(-(i8(lsb)))
		v[0] ^= (0xe1 & red_mask)
	}
	return z
}

// pad_block returns `data` right-padded with zeros to a multiple of 16 bytes.
@[direct_array_access]
fn pad_block(data []u8) []u8 {
	temp := data.clone()
	rem := data.len % block_size
	if rem == 0 {
		return temp
	}
	mut out := []u8{cap: data.len + block_size - rem}
	out << temp
	out << []u8{len: block_size - rem} // padding
	return out
}

// inc32 increments the last 32 bits of the 16-byte counter block in place.
@[direct_array_access]
fn inc32(mut ctr []u8) {
	mut c := (u32(ctr[12]) << 24) | (u32(ctr[13]) << 16) | (u32(ctr[14]) << 8) | u32(ctr[15])
	// detect for wrapping u32 counter
	if (c + 1) == 0 {
		panic('AES-GCM: inc32 overflow the counter')
	}
	c += 1
	ctr[12] = u8(c >> 24)
	ctr[13] = u8(c >> 16)
	ctr[14] = u8(c >> 8)
	ctr[15] = u8(c)
}

// len_block returns the GHASH length block: the bit lengths of the AAD and the
// ciphertext, each encoded as a big-endian 64-bit integer.
fn len_block(ad_len int, ct_len int) []u8 {
	a := u64(ad_len) * 8
	c := u64(ct_len) * 8
	mut b := zero_block.clone()
	for i in 0 .. 8 {
		b[7 - i] = u8(a >> (8 * u32(i)))
		b[15 - i] = u8(c >> (8 * u32(i)))
	}
	return b
}

// j0 derives the pre-counter block for a 96-bit nonce: nonce || 0^31 || 1.
@[direct_array_access]
fn j0_from_nonce(nonce []u8) []u8 {
	mut j := zero_block.clone()
	for i in 0 .. 12 {
		j[i] = nonce[i]
	}
	j[15] = 1
	return j
}

// bytes_equal compares two equal-length byte slices in constant time.
// similar to `crypto.internal.subtle.constant_time_compare()` ones
@[direct_array_access]
fn bytes_equal(a []u8, b []u8) bool {
	if a.len != b.len {
		return false
	}
	mut diff := u8(0)
	for i in 0 .. a.len {
		diff |= a[i] ^ b[i]
	}
	return diff == 0
}