#pragma once
#include "tlibc/std.h"
#include "tlibc/collections/Array.h"
#include "tlibc/string/str.h"
#include "tlibc/errors.h"
#include "bearssl_block.h"
#include "bearssl_rand.h"
#include "bearssl_rsa.h"

//////////////////////////////////////////////////////////////////////////////
//                                  hash.c                                  //
//////////////////////////////////////////////////////////////////////////////

/// @brief hashes password multiple times using its own hash as salt 
/// @param password some byte array 
/// @param out_buffer u8[PASSWORD_HASH_SIZE]
/// @param iterations number of iterations
void hash_password(Array(u8) password, u8* out_buffer, i32 iterations);
#define PASSWORD_HASH_SIZE 32

#define __PASSWORD_HASH_LVL_ITERATIONS 1e5

//////////////////////////////////////////////////////////////////////////////
//                                  rng.c                                  //
//////////////////////////////////////////////////////////////////////////////

/// @brief Initialize prng context with sha256 hashing algorithm 
///         and seed from system-provided cryptographic random bytes source.
/// @param rng_vtable_ptr pointer to vtable field in prng context. The field must be initialized.
/// EXAMPLE:
/// ```
/// br_hmac_drbg_context rng_ctx = { .vtable = &br_hmac_drbg_vtable };
/// rng_init_sha256_seedFromTime(&rng_ctx.vtable);
/// ```
void rng_init_sha256_seedFromSystem(const br_prng_class** rng_vtable_ptr);

/// @brief Initialize prng context with sha256 hashing algorithm and seed from CLOCK_REALTIME.
/// @param rng_vtable_ptr pointer to vtable field in prng context. The field must be initialized.
/// EXAMPLE:
/// ```
/// br_hmac_drbg_context rng_ctx = { .vtable = &br_hmac_drbg_vtable };
/// rng_init_sha256_seedFromTime(&rng_ctx.vtable);
/// ```
void rng_init_sha256_seedFromTime(const br_prng_class** rng_vtable_ptr);


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

#define __AES_SESSION_KEY_SIZE 32
#define __AES_DB_KEY_SIZE 32

typedef struct EncryptedBlockHeader {
    u8 padding_size;
} __attribute__((aligned(16))) EncryptedBlockHeader;

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

typedef struct EncryptorAES {
    br_aes_ct64_cbcenc_keys enc_ctx;
    br_hmac_drbg_context rng_ctx;
    u8 iv[__AES_IV_SIZE];
    u8 buf[__AES_BUFFER_SIZE];
} EncryptorAES;

/// @param key Array<u8, 16 | 24 | 32>
void EncryptorAES_construct(EncryptorAES* ptr, Array(u8) key);

/// @brief Encrypts `src` and writes output to `dst`.
/// @param src array of any size
/// @param dst array of size >= EncryptorAES_calcDstSize(src.size)
Result(void) EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst);

#define EncryptorAES_calcDstSize(SRC_SIZE) (__AES_IV_SIZE + sizeof(EncryptedBlockHeader) + ALIGN_TO(SRC_SIZE, 16))


typedef struct DecryptorAES {
    br_aes_ct64_cbcdec_keys dec_ctx;
    u8 iv[__AES_IV_SIZE];
    u8 buf[__AES_BUFFER_SIZE];
} DecryptorAES;

/// @param key Array<u8, 16 | 24 | 32>
void DecryptorAES_construct(DecryptorAES* ptr, Array(u8) key);

/// @brief Decrypts `src` and writes output to `dst`.
/// @param src array of size at least EncryptorAES_calcDstSize(0). Size must be multiple of 16.
/// @param dst array of size >= src.size
/// @param decrypted_size size of original data without padding added by EncryptorAES_encrypt
Result(void) DecryptorAES_decrypt(DecryptorAES* ptr, Array(u8) src, Array(u8) dst, u32* decrypted_size);


//////////////////////////////////////////////////////////////////////////////
//                                  RSA.c                                   //
//////////////////////////////////////////////////////////////////////////////

#define __RSA_DEFAULT_KEY_SIZE 3072

/// @brief generate random key pair based on system time
/// @param key_size size of public key in bits (2048/3072/4096)
/// @param sk key for decryption
/// @param pk key for encryption
/// @param rng_vtable_ptr pointer to vtable field in prng context. The context must be initialized
Result(void) RSA_generateKeyPair(u32 key_size, 
    br_rsa_private_key* sk, br_rsa_public_key* pk,
    const br_prng_class** rng_vtable_ptr);

Result(void) RSA_generateKeyPairFromTime(u32 key_size, 
    br_rsa_private_key* sk, br_rsa_public_key* pk);

Result(void) RSA_generateKeyPairFromPassword(u32 key_size, 
    br_rsa_private_key* sk, br_rsa_public_key* pk, str password);

static inline void RSA_destroyPrivateKey(br_rsa_private_key* sk){
    free(sk->p);
}

static inline void RSA_destroyPublicKey(br_rsa_public_key* sk){
    free(sk->n);
}

/// @param sk some private key 
/// @param pk out public key. WARNING: .n is allocated on heap
Result(void) RSA_computePublicKey(const br_rsa_private_key* sk, br_rsa_public_key* pk);

/// @brief Encode key data in human-readable format
/// @param src some data
/// @return heap-allocated string
str RSA_serializePrivateKey_base64(const br_rsa_private_key* sk);

/// @param src serialized private key format "RSA-Private-%SIZE%:%DATA_BASE64%"
/// @param sk out private key. WARNING: .p is allocated on heap
Result(void) RSA_parsePrivateKey_base64(const str src, br_rsa_private_key* sk);

/// @brief Encode key data in human-readable format
/// @param src some data
/// @return heap-allocated string
str RSA_serializePublicKey_base64(const br_rsa_public_key* sk);

/// @param src serialized public key format "RSA-Public-%SIZE%:%DATA_BASE64%"
/// @param sk out public key. WARNING: .p is allocated on heap
Result(void) RSA_parsePublicKey_base64(const str src, br_rsa_public_key* sk);


typedef struct EncryptorRSA {
    const br_rsa_public_key* pk;
    br_hmac_drbg_context rng;
} EncryptorRSA;
void EncryptorRSA_construct(EncryptorRSA* ptr, const br_rsa_public_key* pk);
void EncryptorRSA_encrypt(EncryptorRSA* ptr, Array(u8) src, Array(u8) dst, u32* encrypted_size);

typedef struct DecryptorRSA {
    const br_rsa_private_key* sk;
} DecryptorRSA;

void DecryptorRSA_construct(DecryptorRSA* ptr, const br_rsa_private_key* sk);
void DecryptorRSA_decrypt(DecryptorRSA* ptr, Array(u8) buf, u32* decrypted_size);