local testimony = require("testimony")
local litls = require("litls.core")

local hex = litls.hex_encode
local unhex = litls.hex_decode

local testify = testimony.new("== LITLS AEAD ==")

-- RFC 8439 Section 2.8.2 - ChaCha20-Poly1305 AEAD test vector
testify:that("ChaCha20-Poly1305 AEAD encryption (RFC 8439 2.8.2)", function()
	local key = unhex("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
	local nonce = unhex("070000004041424344454647")
	local aad = unhex("50515253c0c1c2c3c4c5c6c7")
	local plaintext =
		"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it."

	local ct, tag = litls.chacha20_poly1305_encrypt(key, nonce, aad, plaintext)
	testimony.assert_equal(
		"d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d63dbea45e8ca9671282fafb69da92728b1a71de0a9e060b2905d6a5b67ecd3b3692ddbd7f2d778b8c9803aee328091b58fab324e4fad675945585808b4831d7bc3ff4def08e4b7a9de576d26586cec64b6116",
		hex(ct)
	)
	testimony.assert_equal("1ae10b594f09e26a7e902ecbd0600691", hex(tag))
end)

testify:that("ChaCha20-Poly1305 AEAD decryption (RFC 8439 2.8.2)", function()
	local key = unhex("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
	local nonce = unhex("070000004041424344454647")
	local aad = unhex("50515253c0c1c2c3c4c5c6c7")
	local ct = unhex(
		"d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d63dbea45e8ca9671282fafb69da92728b1a71de0a9e060b2905d6a5b67ecd3b3692ddbd7f2d778b8c9803aee328091b58fab324e4fad675945585808b4831d7bc3ff4def08e4b7a9de576d26586cec64b6116"
	)
	local tag = unhex("1ae10b594f09e26a7e902ecbd0600691")

	local pt = litls.chacha20_poly1305_decrypt(key, nonce, aad, ct, tag)
	testimony.assert_equal(
		"Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.",
		pt
	)
end)

testify:that("ChaCha20-Poly1305 AEAD rejects modified ciphertext", function()
	local key = unhex("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
	local nonce = unhex("070000004041424344454647")
	local aad = unhex("50515253c0c1c2c3c4c5c6c7")
	local ct = unhex("d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d6")
	local bad_tag = unhex("00000000000000000000000000000000")

	local pt, err = litls.chacha20_poly1305_decrypt(key, nonce, aad, ct, bad_tag)
	testimony.assert_nil(pt)
	testimony.assert_equal("authentication failed", err)
end)

testify:that("ChaCha20-Poly1305 AEAD encrypt/decrypt round-trip", function()
	local key = litls.random_bytes(32)
	local nonce = litls.random_bytes(12)
	local aad = "additional data"
	local plaintext = "secret message"

	local ct, tag = litls.chacha20_poly1305_encrypt(key, nonce, aad, plaintext)
	local pt = litls.chacha20_poly1305_decrypt(key, nonce, aad, ct, tag)
	testimony.assert_equal(plaintext, pt)
end)

-- ── AES-128-GCM ──────────────────────────────────────────────

-- NIST SP 800-38D Test Case 3: AES-128-GCM with 64-byte plaintext
testify:that("AES-128-GCM encryption (NIST Test Case 3)", function()
	local key = unhex("feffe9928665731c6d6a8f9467308308")
	local nonce = unhex("cafebabefacedbaddecaf888")
	local pt = unhex(
		"d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a72"
			.. "1c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b391aafd255"
	)

	local ct, tag = litls.aes128_gcm_encrypt(key, nonce, "", pt)
	testimony.assert_equal(
		"42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e"
			.. "21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091473f5985",
		hex(ct)
	)
	testimony.assert_equal("4d5c2af327cd64a62cf35abd2ba6fab4", hex(tag))
end)

-- NIST SP 800-38D Test Case 4: AES-128-GCM with AAD
testify:that("AES-128-GCM encryption with AAD (NIST Test Case 4)", function()
	local key = unhex("feffe9928665731c6d6a8f9467308308")
	local nonce = unhex("cafebabefacedbaddecaf888")
	local aad = unhex("feedfacedeadbeeffeedfacedeadbeefabaddad2")
	local pt = unhex(
		"d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a72"
			.. "1c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39"
	)

	local ct, tag = litls.aes128_gcm_encrypt(key, nonce, aad, pt)
	testimony.assert_equal(
		"42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e"
			.. "21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091",
		hex(ct)
	)
	testimony.assert_equal("5bc94fbc3221a5db94fae95ae7121a47", hex(tag))
end)

-- AES-256-GCM empty plaintext
testify:that("AES-256-GCM empty plaintext", function()
	local key = unhex("0000000000000000000000000000000000000000000000000000000000000000")
	local nonce = unhex("000000000000000000000000")

	local ct, tag = litls.aes256_gcm_encrypt(key, nonce, "", "")
	testimony.assert_equal("", ct)
	testimony.assert_equal("530f8afbc74536b9a963b4f1c4cb738b", hex(tag))
end)

-- AES-128-GCM encrypt/decrypt round-trip
testify:that("AES-128-GCM encrypt/decrypt round-trip", function()
	local key = litls.random_bytes(16)
	local nonce = litls.random_bytes(12)
	local aad = "associated data"
	local plaintext = "secret AES-GCM message"

	local ct, tag = litls.aes128_gcm_encrypt(key, nonce, aad, plaintext)
	local pt = litls.aes128_gcm_decrypt(key, nonce, aad, ct, tag)
	testimony.assert_equal(plaintext, pt)
end)

-- AES-256-GCM encrypt/decrypt round-trip
testify:that("AES-256-GCM encrypt/decrypt round-trip", function()
	local key = litls.random_bytes(32)
	local nonce = litls.random_bytes(12)
	local aad = "aes256 aad"
	local plaintext = "AES-256-GCM test"

	local ct, tag = litls.aes256_gcm_encrypt(key, nonce, aad, plaintext)
	local pt = litls.aes256_gcm_decrypt(key, nonce, aad, ct, tag)
	testimony.assert_equal(plaintext, pt)
end)

-- AES-GCM tag verification failure
testify:that("AES-128-GCM rejects modified ciphertext", function()
	local key = litls.random_bytes(16)
	local nonce = litls.random_bytes(12)
	local ct, tag = litls.aes128_gcm_encrypt(key, nonce, "", "test data")

	-- Modify the tag
	local bad_tag = unhex("00000000000000000000000000000000")
	local pt, err = litls.aes128_gcm_decrypt(key, nonce, "", ct, bad_tag)
	testimony.assert_nil(pt)
	testimony.assert_equal("authentication failed", err)
end)

-- AES-GCM AAD-only (empty plaintext with AAD)
testify:that("AES-128-GCM AAD-only (empty plaintext)", function()
	local key = litls.random_bytes(16)
	local nonce = litls.random_bytes(12)
	local aad = "authenticate this but don't encrypt"

	local ct, tag = litls.aes128_gcm_encrypt(key, nonce, aad, "")
	testimony.assert_equal("", ct)
	testimony.assert_equal(16, #tag)

	-- Verify decrypts successfully
	local pt = litls.aes128_gcm_decrypt(key, nonce, aad, "", tag)
	testimony.assert_equal("", pt)

	-- Different AAD should fail
	local pt2, err = litls.aes128_gcm_decrypt(key, nonce, "wrong aad", "", tag)
	testimony.assert_nil(pt2)
end)

testify:conclude()