local Stream = require(".crypto.util.stream")
local Hex = require(".crypto.util.hex")
local Array = require(".crypto.util.array")

-- Ciphers
local AES128Cipher = require(".crypto.cipher.aes128")
local AES192Cipher = require(".crypto.cipher.aes192")
local AES256Cipher = require(".crypto.cipher.aes256")

-- Modes
local CBCMode = require(".crypto.cipher.mode.cbc")
local ECBMode = require(".crypto.cipher.mode.ecb")
local CFBMode = require(".crypto.cipher.mode.cfb")
local OFBMode = require(".crypto.cipher.mode.ofb")
local CTRMode = require(".crypto.cipher.mode.ctr")

-- Padding
local ZeroPadding = require(".crypto.padding.zero")

local public = {}

local getBlockCipher = function(keyLength)
    if keyLength == 128 then
        return AES128Cipher
    elseif keyLength == 192 then
        return AES192Cipher
    elseif keyLength == 256 then
        return AES256Cipher
    elseif keyLength == nil then
        return AES128Cipher
    else
        return nil
    end
end

local getMode = function(mode)
    if mode == "CBC" then
        return CBCMode
    elseif mode == "ECB" then
        return ECBMode
    elseif mode == "CFB" then
        return CFBMode
    elseif mode == "OFB" then
        return OFBMode
    elseif mode == "CTR" then
        return CTRMode
    else
        return nil
    end
end


--- Encrypts the given data using AES encryption.
--- @param data string - The data to be encrypted.
--- @param key string - The key to use for encryption.
--- @param iv? string (optional) - The initialization vector to use for encryption. Defaults to 16 null bytes.
--- @param mode? string (optional) - The mode to use for encryption. Defaults to "CBC".
--- @param keyLength? number (optional) - The length of the key to use for encryption. Defaults to 128.
--- @returns table - A table containing the encrypted data in bytes, hex, and string formats.
public.encrypt = function(data, key, iv, mode, keyLength)
    local d = Array.fromString(data)
    local k = Array.fromString(key)
    local _iv = iv ~= nil and Array.fromString(iv) or Array.fromHex("00000000000000000000000000000000")

    local cipherMode = getMode(mode) or CBCMode
    local blockCipher = getBlockCipher(keyLength) or AES128Cipher

    local cipher = cipherMode.Cipher()
        .setKey(k)
        .setBlockCipher(blockCipher)
        .setPadding(ZeroPadding);


    local cipherOutput = cipher
        .init()
        .update(Stream.fromArray(_iv))
        .update(Stream.fromArray(d))
        .finish()

    local results = {}

    results.asBytes = function()
        return cipherOutput.asBytes()
    end

    results.asHex = function()
        return cipherOutput.asHex()
    end

    results.asString = function()
        return cipherOutput.asString()
    end

    return results
end

--- Decrypts the given data using AES decryption.
--- @param cipher string - The hex encoded cipher to be decrypted.
--- @param key string - The key to use for decryption.
--- @param iv? string (optional) - The initialization vector to use for decryption. Defaults to 16 null bytes.
--- @param mode? string (optional) - The mode to use for decryption. Defaults to "CBC".
--- @param keyLength? number (optional) - The length of the key to use for decryption. Defaults to 128.
public.decrypt = function(cipher, key, iv, mode, keyLength)
    local cipherText = Array.fromHex(cipher)
    local k = Array.fromString(key)
    local _iv = iv ~= nil and Array.fromString(iv) or Array.fromHex("00000000000000000000000000000000")

    local cipherMode = getMode(mode) or CBCMode
    local blockCipher = getBlockCipher(keyLength) or AES128Cipher


    local decipher = cipherMode.Decipher()
            .setKey(k)
            .setBlockCipher(blockCipher)
            .setPadding(ZeroPadding);


    local plainOutput = decipher
        .init()
        .update(Stream.fromArray(_iv))
        .update(Stream.fromArray(cipherText))
        .finish()

    local results = {}

    results.asBytes = function()
        return plainOutput.asBytes()
    end

    results.asHex = function()
        return plainOutput.asHex()
    end

    results.asString = function()
        return plainOutput.asString()
    end

    return results
end


return public