#pragma once

#include <vector>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <openssl/evp.h>
#include <openssl/pem.h>


template<std::size_t size>
std::pair<std::array<std::uint8_t, size>, std::array<std::uint8_t, size>> newRsaKeys() {
    // create the context
    const auto context = std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)>(
        EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr),
        EVP_PKEY_CTX_free
    );

    if (context == nullptr)
        throw std::runtime_error("Could not create an EVP context.");

    if (EVP_PKEY_keygen_init(context.get()) <= 0)
        throw std::runtime_error("Could not initialize the EVP context.");

    // configure the context
    if (EVP_PKEY_CTX_set_rsa_keygen_bits(context.get(), size) <= 0)
        throw std::runtime_error("Error setting RSA key size.");

    // create the private key
    EVP_PKEY* rawKeyPair = nullptr;
    if (EVP_PKEY_keygen(context.get(), &rawKeyPair) <= 0)
        throw std::runtime_error("Could not generate RSA private key.");

    const std::unique_ptr<EVP_PKEY, decltype(&EVP_PKEY_free)> keyPair(
        rawKeyPair,
        EVP_PKEY_free
    );

    // extract the private and public key
    const std::shared_ptr<BIO> privateBio(
        BIO_new(BIO_s_mem()),
        BIO_free
    );

    if (!PEM_write_bio_PrivateKey(
        privateBio.get(),
        keyPair.get(),
        nullptr,
        nullptr,
        0,
        nullptr,
        nullptr
    ))
        throw std::runtime_error("Could not generate RSA private key.");

    std::array<std::uint8_t, size> privateKey;
    BIO_read(privateBio.get(), privateKey.data(), BIO_pending(privateBio.get()));

    const std::shared_ptr<BIO> publicBio(
        BIO_new(BIO_s_mem()),
        BIO_free
    );
    if (!PEM_write_bio_PUBKEY(publicBio.get(), keyPair.get()))
        throw std::runtime_error("Could not generate RSA public key.");

    std::array<std::uint8_t, size> publicKey;
    BIO_read(publicBio.get(), publicKey.data(), BIO_pending(publicBio.get()));

    return {privateKey, publicKey};
}