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

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

local testify = testimony.new("== LITLS TLS 1.3 Record Layer ==")

-- Test AEAD encrypt/decrypt with known keys from RFC 8448
-- Using AES-128-GCM (suite 0x1301)

testify:that("AES-128-GCM encrypt then decrypt with RFC 8448 server HS key", function()
	local key = unhex("3fce516009c21727d0f2e4e86ee403bc")
	local iv = unhex("5d313eb2671276ee13000b30")

	-- Construct nonce: iv XOR seq=0 => iv itself
	local nonce = iv

	-- Simulate a record: encrypt "hello" with AAD = record header
	local aad = unhex("1703030016") -- type 0x17, version 0x0303, length 22 (5+1+16)
	local plaintext = "hello"

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

	-- Decrypt should recover plaintext
	local pt = litls.aes128_gcm_decrypt(key, nonce, aad, ct, tag)
	testimony.assert_equal("hello", pt)
end)

testify:that("AES-128-GCM rejects tampered ciphertext", function()
	local key = unhex("3fce516009c21727d0f2e4e86ee403bc")
	local iv = unhex("5d313eb2671276ee13000b30")
	local aad = unhex("1703030016")

	local ct, tag = litls.aes128_gcm_encrypt(key, iv, aad, "test")
	-- Flip a bit in ciphertext
	local bad_ct = string.char(bit.bxor(ct:byte(1), 1)) .. ct:sub(2)
	local pt, err = litls.aes128_gcm_decrypt(key, iv, aad, bad_ct, tag)
	testimony.assert_nil(pt)
end)

testify:that("AES-128-GCM rejects tampered tag", function()
	local key = unhex("3fce516009c21727d0f2e4e86ee403bc")
	local iv = unhex("5d313eb2671276ee13000b30")
	local aad = unhex("1703030016")

	local ct, tag = litls.aes128_gcm_encrypt(key, iv, aad, "test")
	-- Flip a bit in tag
	local bad_tag = string.char(bit.bxor(tag:byte(1), 1)) .. tag:sub(2)
	local pt, err = litls.aes128_gcm_decrypt(key, iv, aad, ct, bad_tag)
	testimony.assert_nil(pt)
end)

-- Test with ChaCha20-Poly1305
testify:that("ChaCha20-Poly1305 encrypt then decrypt", function()
	local key = unhex("b67b7d690cc16c4e75e54213cb2d37b4e9c912bcded9105d42befd59d391ad38")
	local nonce = unhex("5d313eb2671276ee13000b30")
	local aad = unhex("17030300ff") -- example record header

	local plaintext = "TLS 1.3 record layer test data"
	local ct, tag = litls.chacha20_poly1305_encrypt(key, nonce, aad, plaintext)
	testimony.assert_not_nil(ct)
	testimony.assert_equal(#plaintext, #ct)

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

-- Test nonce construction: seq XOR iv
testify:that("Nonce construction: seq XOR iv", function()
	-- For seq=0, nonce = iv
	local iv = unhex("5d313eb2671276ee13000b30")
	-- seq=0: nonce should equal iv
	testimony.assert_equal(hex(iv), "5d313eb2671276ee13000b30")

	-- For seq=1, nonce = iv XOR (0...01)
	-- Last byte: 0x30 XOR 0x01 = 0x31
	-- Expected: 5d313eb2671276ee13000b31
	local nonce_bytes = {}
	for i = 1, 12 do
		nonce_bytes[i] = iv:byte(i)
	end
	-- XOR seq=1 into last 8 bytes
	nonce_bytes[12] = bit.bxor(nonce_bytes[12], 1)
	local nonce_hex = ""
	for i = 1, 12 do
		nonce_hex = nonce_hex .. string.format("%02x", nonce_bytes[i])
	end
	testimony.assert_equal("5d313eb2671276ee13000b31", nonce_hex)
end)

-- Test AES-256-GCM (for SHA-384 cipher suites)
testify:that("AES-256-GCM encrypt then decrypt", function()
	local key = unhex("9f02283b6c9c07efc26bb9f2ac92e356a11af9f05531f856ad47116b45a95032")
	local nonce = unhex("cf782b88dd83549aadf1e984")
	local aad = unhex("1703030020")

	local plaintext = "AES-256-GCM record test"
	local ct, tag = litls.aes256_gcm_encrypt(key, nonce, aad, plaintext)
	testimony.assert_not_nil(ct)

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

testify:conclude()