local ffi = require("@lune/ffi")

-- load Windows crypto provider
local bcrypt = ffi.open("bcrypt.dll")

-- define function signatures (Lune style)
local BCryptOpenAlgorithmProvider = bcrypt:func(
	"BCryptOpenAlgorithmProvider",
	{ "pointer", "pointer", "pointer", "u32" },
	"i32"
)

local BCryptSetProperty = bcrypt:func(
	"BCryptSetProperty",
	{ "pointer", "pointer", "pointer", "u32", "u32" },
	"i32"
)

local BCryptGenerateSymmetricKey = bcrypt:func(
	"BCryptGenerateSymmetricKey",
	{ "pointer", "pointer", "pointer", "u32", "pointer", "u32", "u32" },
	"i32"
)

local BCryptEncrypt = bcrypt:func(
	"BCryptEncrypt",
	{ "pointer", "pointer", "u32", "pointer", "pointer", "u32", "pointer", "u32", "pointer", "u32" },
	"i32"
)

local BCryptDecrypt = bcrypt:func(
	"BCryptDecrypt",
	{ "pointer", "pointer", "u32", "pointer", "pointer", "u32", "pointer", "u32", "pointer", "u32" },
	"i32"
)

local BCryptDestroyKey = bcrypt:func(
	"BCryptDestroyKey",
	{ "pointer" },
	"i32"
)

local BCryptCloseAlgorithmProvider = bcrypt:func(
	"BCryptCloseAlgorithmProvider",
	{ "pointer", "u32" },
	"i32"
)

-- constants
local AES = "AES"
local CBC = "ChainingModeCBC"

-- PKCS7 padding
local function pad(data)
	local pad = 16 - (#data % 16)
	return data .. string.rep(string.char(pad), pad)
end

local function unpad(data)
	local n = string.byte(data, -1)
	return data:sub(1, -1 - n)
end

-- AES CBC encrypt
local function aes_encrypt(key: string, iv: string, plaintext: string): string
	assert(#key == 16 or #key == 32, "key must be 16 or 32 bytes")
	assert(#iv == 16, "iv must be 16 bytes")

	local arena = ffi.arena()

	local alg = arena:alloc(ffi.sizeof("pointer"))
	local keyh = arena:alloc(ffi.sizeof("pointer"))

	local keybuf = arena:alloc(#key)
	ffi.copy(keybuf, key, #key)

	local ivbuf = arena:alloc(#iv)
	ffi.copy(ivbuf, iv, #iv)

	local data = pad(plaintext)
	local inbuf = arena:alloc(#data)
	ffi.copy(inbuf, data, #data)

	-- open AES provider
	assert(BCryptOpenAlgorithmProvider(alg, AES, nil, 0) == 0)

	-- set CBC mode
	assert(BCryptSetProperty(ffi.readpointer(alg), CBC, CBC, #CBC, 0) == 0)

	-- create key
	assert(BCryptGenerateSymmetricKey(
		ffi.readpointer(alg),
		keyh,
		nil,
		0,
		keybuf,
		#key,
		0
	) == 0)

	-- query output size
	local outlenbuf = arena:alloc(4)
	BCryptEncrypt(
		ffi.readpointer(keyh),
		inbuf,
		#data,
		nil,
		ivbuf,
		#iv,
		nil,
		0,
		outlenbuf,
		0
	)

	local outlen = ffi.readu32(outlenbuf)
	local outbuf = arena:alloc(outlen)

	-- encrypt
	assert(BCryptEncrypt(
		ffi.readpointer(keyh),
		inbuf,
		#data,
		nil,
		ivbuf,
		#iv,
		outbuf,
		outlen,
		outlenbuf,
		0
	) == 0)

	BCryptDestroyKey(ffi.readpointer(keyh))
	BCryptCloseAlgorithmProvider(ffi.readpointer(alg), 0)

	return ffi.string(outbuf, outlen)
end

-- AES CBC decrypt
local function aes_decrypt(key: string, iv: string, cipher: string): string
	assert(#key == 16 or #key == 32)
	assert(#iv == 16)

	local arena = ffi.arena()

	local alg = arena:alloc(ffi.sizeof("pointer"))
	local keyh = arena:alloc(ffi.sizeof("pointer"))

	local keybuf = arena:alloc(#key)
	ffi.copy(keybuf, key, #key)

	local ivbuf = arena:alloc(#iv)
	ffi.copy(ivbuf, iv, #iv)

	local inbuf = arena:alloc(#cipher)
	ffi.copy(inbuf, cipher, #cipher)

	assert(BCryptOpenAlgorithmProvider(alg, AES, nil, 0) == 0)
	assert(BCryptSetProperty(ffi.readpointer(alg), CBC, CBC, #CBC, 0) == 0)

	assert(BCryptGenerateSymmetricKey(
		ffi.readpointer(alg),
		keyh,
		nil,
		0,
		keybuf,
		#key,
		0
	) == 0)

	local outlenbuf = arena:alloc(4)
	BCryptDecrypt(
		ffi.readpointer(keyh),
		inbuf,
		#cipher,
		nil,
		ivbuf,
		#iv,
		nil,
		0,
		outlenbuf,
		0
	)

	local outlen = ffi.readu32(outlenbuf)
	local outbuf = arena:alloc(outlen)

	assert(BCryptDecrypt(
		ffi.readpointer(keyh),
		inbuf,
		#cipher,
		nil,
		ivbuf,
		#iv,
		outbuf,
		outlen,
		outlenbuf,
		0
	) == 0)

	BCryptDestroyKey(ffi.readpointer(keyh))
	BCryptCloseAlgorithmProvider(ffi.readpointer(alg), 0)

	return unpad(ffi.string(outbuf, outlen))
end

-- test
local key = "1234567890123456"
local iv  = "abcdefghijklmnop"

local enc = aes_encrypt(key, iv, "hello lune")
local dec = aes_decrypt(key, iv, enc)

print(enc)
print(dec)