#pragma once
#include "tlibc/collections/Array.h"
#include "tlibc/errors.h"
#include "bearssl_block.h"
#include "cryptography.h"

//////////////////////////////////////////////////////////////////////////////
//                                                                          //
//                                  AES.c                                   //
//                                                                          //
//////////////////////////////////////////////////////////////////////////////

#define AESBlockEncryptor_DEFAULT_CLASS (&br_aes_big_cbcenc_vtable)
#define AESBlockDecryptor_DEFAULT_CLASS (&br_aes_big_cbcdec_vtable)
#define AESStream_DEFAULT_CLASS (&br_aes_big_ctr_vtable)


//TODO: use PKS#7 instead of this garbage
typedef struct EncryptedBlockHeader {
    u8 padding_size;
} __attribute__((aligned(16))) EncryptedBlockHeader;

//////////////////////////////////////////////////////////////////////////////
//                            AESBlockEncryptor                             //
//////////////////////////////////////////////////////////////////////////////

#define __AES_BLOCK_IV_SIZE 16
// must be multiple of 16
#define __AES_BUFFER_SIZE 512

typedef struct AESBlockEncryptor {
    const br_block_cbcenc_class* enc_class;
    u8 enc_keys[sizeof(br_aes_big_cbcenc_keys)];
    u8 buf[__AES_BUFFER_SIZE];
    u8 iv[__AES_BLOCK_IV_SIZE];
    br_hmac_drbg_context rng_ctx;
} AESBlockEncryptor;

/// @param key supported sizes: 16, 24, 32
/// @param enc_class &br_aes_XXX_cbcenc_vtable
void AESBlockEncryptor_construct(AESBlockEncryptor* ptr, Array(u8) key, const br_block_cbcenc_class* enc_class);

/// @brief Encrypts a complete message. For part-by-part encryption use AESStreamEncryptor.
/// @param src array of any size
/// @param dst array of size >= AESBlockEncryptor_calcDstSize(src.size)
/// @return size of encrypted data
Result(u32) AESBlockEncryptor_encrypt(AESBlockEncryptor* ptr, Array(u8) src, Array(u8) dst);

#define AESBlockEncryptor_calcDstSize(SRC_SIZE) (__AES_BLOCK_IV_SIZE + sizeof(EncryptedBlockHeader) + ALIGN_TO(SRC_SIZE, 16))

//////////////////////////////////////////////////////////////////////////////
//                            AESBlockDecryptor                             //
//////////////////////////////////////////////////////////////////////////////

typedef struct AESBlockDecryptor {
    const br_block_cbcdec_class* dec_class;
    u8 dec_keys[sizeof(br_aes_big_cbcdec_keys)];
    u8 buf[__AES_BUFFER_SIZE];
    u8 iv[__AES_BLOCK_IV_SIZE];
} AESBlockDecryptor;

/// @param key supported sizes: 16, 24, 32
/// @param dec_class &br_aes_XXX_cbcdec_vtable
void AESBlockDecryptor_construct(AESBlockDecryptor* ptr, Array(u8) key, const br_block_cbcdec_class* dec_class);

/// @brief Decrypts a complete message. For part-by-part decryption use AESStreamEncryptor.
/// @param src array of size at least AESBlockEncryptor_calcDstSize(0). Size must be multiple of 16.
/// @param dst array of size >= src.size
/// @return size of decrypted data
Result(u32) AESBlockDecryptor_decrypt(AESBlockDecryptor* ptr, Array(u8) src, Array(u8) dst);

//////////////////////////////////////////////////////////////////////////////
//                            AESStreamEncryptor                            //
//////////////////////////////////////////////////////////////////////////////

#define __AES_STREAM_IV_SIZE 12

typedef struct AESStreamEncryptor {
    const br_block_ctr_class* ctr_class;
    u8 ctr_keys[sizeof(br_aes_big_ctr_keys)];
    u8 buf[__AES_BUFFER_SIZE];
    u8 iv[__AES_STREAM_IV_SIZE];
    u32 block_counter;
} AESStreamEncryptor;

/// @warning If you start from not first block in sequence (example: to continue a previously interrupted operation) set correct values for ptr->block_counter and ptr->iv or encryption/decryption will produce incorrect data.
/// @param key supported sizes: 16, 24, 32
/// @param dec_class &br_aes_XXX_ctr_vtable
void AESStreamEncryptor_construct(AESStreamEncryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class);

#define AESStreamEncryptor_calcDstSize(src_size) (src_size + __AES_BLOCK_IV_SIZE)

/// @brief If ptr->block_counter == 0, writes random IV to `dst`. After that writes encrypted data to dst.
/// @param src array of any size
/// @param dst array of size >= AESStreamEncryptor_calcDstSize(src.size)
/// @return size of encrypted data
Result(u32) AESStreamEncryptor_encrypt(AESStreamEncryptor* ptr, Array(u8) src, Array(u8) dst);

//////////////////////////////////////////////////////////////////////////////
//                            AESStreamDecryptor                            //
//////////////////////////////////////////////////////////////////////////////

typedef struct AESStreamDecryptor {
    const br_block_ctr_class* ctr_class;
    u8 ctr_keys[sizeof(br_aes_big_ctr_keys)];
    u8 buf[__AES_BUFFER_SIZE];
    u8 iv[__AES_STREAM_IV_SIZE];
    u32 block_counter;
} AESStreamDecryptor;

/// @warning If you start from not first block in sequence (example: to continue a previously interrupted operation) set correct values for ptr->block_counter and ptr->iv or encryption/decryption will produce incorrect data.
/// @param key supported sizes: 16, 24, 32
/// @param dec_class &br_aes_XXX_ctr_vtable
void AESStreamDecryptor_construct(AESStreamDecryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class);

/// @brief Reads IV from `src`, then decrypts data and writes it to dst
/// @param src array of size at least AESStreamEncryptor_calcDstSize(0).
/// @param dst array of size >= src.size
/// @return size of decrypted data
Result(u32) AESStreamDecryptor_decrypt(AESStreamDecryptor* ptr, Array(u8) src, Array(u8) dst);

#define __AESStreamDecryptor_calcDstSize(src_size) (src_size - __AES_BLOCK_IV_SIZE)