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

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

#define RSA_DEFAULT_KEY_SIZE 3072

/// @brief generate random key pair
/// @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);

/// @brief generate random key pair using system crypto-rng provider
Result(void) RSA_generateKeyPairFromSystemRandom(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);

//////////////////////////////////////////////////////////////////////////////
//                               RSAEncryptor                               //
//////////////////////////////////////////////////////////////////////////////

typedef struct RSAEncryptor {
    const br_rsa_public_key* pk;
    br_hmac_drbg_context rng;
} RSAEncryptor;

/// RSA OAEP encryption with SHA256 hashing algorithm
void RSAEncryptor_construct(RSAEncryptor* ptr, const br_rsa_public_key* pk);

/*
https://crypto.stackexchange.com/a/42100
+---------+----------+-------------------------------------------+
|         | overhead |         max message size (bytes)          |
| Hash    | (bytes)  | RSA-1024 | RSA-2048 | RSA-3072 | RSA-4096 |
+---------+----------+----------+----------+----------+----------+
| SHA-1   | 42       | 86       | 214      | 342      | 470      |
| SHA-224 | 58       | 70       | 198      | 326      | 454      |
| SHA-256 | 66       | 62       | 190      | 318      | 446      |
| SHA-384 | 98       | 30       | 158      | 286      | 414      |
| SHA-512 | 130      | N/A      | 126      | 254      | 382      |
+---------+----------+----------+----------+----------+----------+
*/
#define RSAEncryptor_calcMaxSrcSize(KEY_LEN_BITS, HASH_LEN_BITS) (KEY_LEN_BITS / 8 - 2 * HASH_LEN_BITS / 8 - 2)

/// @param src buffer with size <= `RSAEncryptor_calcMaxSrcSize(key_size_bits, 256)`
/// @param dst buffer with size >= key size in bytes
/// @return size of encrypted data
Result(u32) RSAEncryptor_encrypt(RSAEncryptor* ptr, Array(u8) src, Array(u8) dst);

//////////////////////////////////////////////////////////////////////////////
//                               RSADecryptor                               //
//////////////////////////////////////////////////////////////////////////////

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

/// RSA OAEP encryption with SHA256 hashing algorithm
void RSADecryptor_construct(RSADecryptor* ptr, const br_rsa_private_key* sk);

/// @param src buffer with size == key size in bytes
/// @param dst buffer with size >= `RSAEncryptor_calcMaxSrcSize(key_size_bits, 256)`
/// @return size of decrypted data
Result(u32) RSADecryptor_decrypt(RSADecryptor* ptr, Array(u8) src, Array(u8) dst);