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 Key Schedule ==")

-- RFC 8448 Section 3: Simple 1-RTT Handshake test vectors
-- These are the known-good intermediate values from the RFC.

-- Test HKDF-Expand-Label directly
testify:that("HKDF-Expand-Label SHA-256 basic", function()
	-- Verify the function exists and produces deterministic output
	local secret = string.rep("\x00", 32)
	local result = litls.tls13_expand_label_256(secret, "key", "", 16)
	testimony.assert_not_nil(result)
	testimony.assert_equal(16, #result)

	-- Same inputs should produce same output
	local result2 = litls.tls13_expand_label_256(secret, "key", "", 16)
	testimony.assert_equal(hex(result), hex(result2))
end)

testify:that("HKDF-Expand-Label SHA-256 with context", function()
	local secret = string.rep("\x0b", 32)
	local context = unhex("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
	local result = litls.tls13_expand_label_256(secret, "derived", context, 32)
	testimony.assert_not_nil(result)
	testimony.assert_equal(32, #result)
end)

-- Test Derive-Secret
testify:that("Derive-Secret SHA-256", function()
	local secret = string.rep("\x00", 32)
	local transcript_hash = litls.sha256("")
	local result = litls.tls13_derive_secret_256(secret, "derived", transcript_hash)
	testimony.assert_not_nil(result)
	testimony.assert_equal(32, #result)
end)

-- Test full key schedule against RFC 8448 values
-- From RFC 8448 Section 3 "Simple 1-RTT Handshake":
testify:that("RFC 8448 early secret derivation", function()
	-- Early Secret = HKDF-Extract(0, 0)
	-- salt = 0 (empty), ikm = 0 (32 zero bytes)
	local zeros = string.rep("\x00", 32)
	local early_secret = litls.hkdf_extract_sha256(zeros, zeros)

	-- RFC 8448: early_secret = 33ad0a1c607ec03b09e6cd9893680ce210adf300aa1f2660e1b22e10f170f92a
	local expected = "33ad0a1c607ec03b09e6cd9893680ce210adf300aa1f2660e1b22e10f170f92a"
	testimony.assert_equal(expected, hex(early_secret))
end)

testify:that("RFC 8448 derived secret from early", function()
	local early_secret = unhex("33ad0a1c607ec03b09e6cd9893680ce210adf300aa1f2660e1b22e10f170f92a")
	local empty_hash = litls.sha256("")

	local derived = litls.tls13_derive_secret_256(early_secret, "derived", empty_hash)
	-- RFC 8448: 6f2615a108c702c5678f54fc9dbab69716c076189c48250cebeac3576c3611ba
	local expected = "6f2615a108c702c5678f54fc9dbab69716c076189c48250cebeac3576c3611ba"
	testimony.assert_equal(expected, hex(derived))
end)

testify:that("RFC 8448 handshake secret", function()
	local derived = unhex("6f2615a108c702c5678f54fc9dbab69716c076189c48250cebeac3576c3611ba")
	-- Shared secret from RFC 8448 (ECDHE)
	local shared_secret = unhex("8bd4054fb55b9d63fdfbacf9f04b9f0d35e6d63f537563efd46272900f89492d")

	local hs_secret = litls.hkdf_extract_sha256(derived, shared_secret)
	-- RFC 8448: 1dc826e93606aa6fdc0aadc12f741b01046aa6b99f691ed221a9f0ca043fbeac
	local expected = "1dc826e93606aa6fdc0aadc12f741b01046aa6b99f691ed221a9f0ca043fbeac"
	testimony.assert_equal(expected, hex(hs_secret))
end)

testify:that("RFC 8448 client handshake traffic secret", function()
	local hs_secret = unhex("1dc826e93606aa6fdc0aadc12f741b01046aa6b99f691ed221a9f0ca043fbeac")
	-- Transcript hash (CH..SH) from RFC 8448
	local transcript_hash = unhex("860c06edc07858ee8e78f0e7428c58edd6b43f2ca3e6e95f02ed063cf0e1cad8")

	local c_hs_traffic = litls.tls13_derive_secret_256(hs_secret, "c hs traffic", transcript_hash)
	-- RFC 8448: b3eddb126e067f35a780b3abf45e2d8f3b1a950738f52e9600746a0e27a55a21
	local expected = "b3eddb126e067f35a780b3abf45e2d8f3b1a950738f52e9600746a0e27a55a21"
	testimony.assert_equal(expected, hex(c_hs_traffic))
end)

testify:that("RFC 8448 server handshake traffic secret", function()
	local hs_secret = unhex("1dc826e93606aa6fdc0aadc12f741b01046aa6b99f691ed221a9f0ca043fbeac")
	local transcript_hash = unhex("860c06edc07858ee8e78f0e7428c58edd6b43f2ca3e6e95f02ed063cf0e1cad8")

	local s_hs_traffic = litls.tls13_derive_secret_256(hs_secret, "s hs traffic", transcript_hash)
	-- RFC 8448: b67b7d690cc16c4e75e54213cb2d37b4e9c912bcded9105d42befd59d391ad38
	local expected = "b67b7d690cc16c4e75e54213cb2d37b4e9c912bcded9105d42befd59d391ad38"
	testimony.assert_equal(expected, hex(s_hs_traffic))
end)

testify:that("RFC 8448 server handshake key + iv", function()
	local s_hs_traffic = unhex("b67b7d690cc16c4e75e54213cb2d37b4e9c912bcded9105d42befd59d391ad38")

	-- key = Expand-Label(secret, "key", "", 16)  (AES-128-GCM)
	local key = litls.tls13_expand_label_256(s_hs_traffic, "key", "", 16)
	-- RFC 8448: 3fce516009c21727d0f2e4e86ee403bc
	testimony.assert_equal("3fce516009c21727d0f2e4e86ee403bc", hex(key))

	-- iv = Expand-Label(secret, "iv", "", 12)
	local iv = litls.tls13_expand_label_256(s_hs_traffic, "iv", "", 12)
	-- RFC 8448: 5d313eb2671276ee13000b30
	testimony.assert_equal("5d313eb2671276ee13000b30", hex(iv))
end)

testify:that("RFC 8448 client handshake key + iv", function()
	local c_hs_traffic = unhex("b3eddb126e067f35a780b3abf45e2d8f3b1a950738f52e9600746a0e27a55a21")

	local key = litls.tls13_expand_label_256(c_hs_traffic, "key", "", 16)
	-- RFC 8448: dbfaa693d1762c5b666af5d950258d01
	testimony.assert_equal("dbfaa693d1762c5b666af5d950258d01", hex(key))

	local iv = litls.tls13_expand_label_256(c_hs_traffic, "iv", "", 12)
	-- RFC 8448: 5bd3c71b836e0b76bb73265f
	testimony.assert_equal("5bd3c71b836e0b76bb73265f", hex(iv))
end)

testify:that("RFC 8448 master secret derivation", function()
	local hs_secret = unhex("1dc826e93606aa6fdc0aadc12f741b01046aa6b99f691ed221a9f0ca043fbeac")
	local empty_hash = litls.sha256("")

	-- derived = Derive-Secret(hs_secret, "derived", Hash(""))
	local derived = litls.tls13_derive_secret_256(hs_secret, "derived", empty_hash)
	-- RFC 8448: 43de77e0c77713859a944db9db2590b53190a65b3ee2e4f12dd7a0bb7ce254b4
	testimony.assert_equal("43de77e0c77713859a944db9db2590b53190a65b3ee2e4f12dd7a0bb7ce254b4", hex(derived))

	-- master_secret = HKDF-Extract(derived, 0)
	local zeros = string.rep("\x00", 32)
	local master_secret = litls.hkdf_extract_sha256(derived, zeros)
	-- RFC 8448: 18df06843d13a08bf2a449844c5f8a478001bc4d4c627984d5a41da8d0402919
	testimony.assert_equal("18df06843d13a08bf2a449844c5f8a478001bc4d4c627984d5a41da8d0402919", hex(master_secret))
end)

testify:that("Derive-Secret determinism for application traffic", function()
	local master_secret = unhex("18df06843d13a08bf2a449844c5f8a478001bc4d4c627984d5a41da8d0402919")
	-- Use a known hash to verify Derive-Secret produces consistent output
	local transcript_hash = litls.sha256("test transcript")

	local c_ap_traffic = litls.tls13_derive_secret_256(master_secret, "c ap traffic", transcript_hash)
	testimony.assert_not_nil(c_ap_traffic)
	testimony.assert_equal(32, #c_ap_traffic)

	-- Same inputs produce same output
	local c_ap_traffic2 = litls.tls13_derive_secret_256(master_secret, "c ap traffic", transcript_hash)
	testimony.assert_equal(hex(c_ap_traffic), hex(c_ap_traffic2))

	-- Different labels produce different output
	local s_ap_traffic = litls.tls13_derive_secret_256(master_secret, "s ap traffic", transcript_hash)
	testimony.assert_equal(true, hex(c_ap_traffic) ~= hex(s_ap_traffic))
end)

testify:that("Expand-Label for key and iv extraction", function()
	-- Use a known traffic secret and verify key/iv derivation works
	local traffic_secret = unhex("b67b7d690cc16c4e75e54213cb2d37b4e9c912bcded9105d42befd59d391ad38")

	local key = litls.tls13_expand_label_256(traffic_secret, "key", "", 16)
	testimony.assert_not_nil(key)
	testimony.assert_equal(16, #key)

	local iv = litls.tls13_expand_label_256(traffic_secret, "iv", "", 12)
	testimony.assert_not_nil(iv)
	testimony.assert_equal(12, #iv)

	-- Deterministic
	local key2 = litls.tls13_expand_label_256(traffic_secret, "key", "", 16)
	testimony.assert_equal(hex(key), hex(key2))
end)

testify:conclude()