#include "RsaPublicKey.hpp"

#include <ranges>
#include <openssl/rsa.h>
#include <openssl/x509.h>

#include "utils/serialize/basics.hpp"


namespace drp::util::crypto {


RsaPublicKey::RsaPublicKey() = default;

RsaPublicKey::RsaPublicKey(const std::vector<uint8_t>& data, const int padMode) {
    this->_data = data;
    this->padMode = padMode;
}


[[nodiscard]] std::vector<uint8_t> RsaPublicKey::encrypt(const std::vector<std::uint8_t>& plainData) const {
    const auto key = this->getOpenSslKey();

    // initialize the encryption context
    const std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)> context(
        EVP_PKEY_CTX_new(key.get(), nullptr),
        EVP_PKEY_CTX_free
    );

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

    // initialize the encryption operation
    if (EVP_PKEY_encrypt_init(context.get()) <= 0)
        throw std::runtime_error("Could not initialize encryption.");

    // set the padding
    if (EVP_PKEY_CTX_set_rsa_padding(context.get(), this->padMode) <= 0)
        throw std::runtime_error("Could not set RSA padding.");

    // get the size of the output buffer
    std::size_t encryptedDataLength;
    if (EVP_PKEY_encrypt(
        context.get(),
        nullptr,
        &encryptedDataLength,
        plainData.data(),
        plainData.size()
    ) <= 0)
        throw std::runtime_error("Could not determine output length.");

    std::vector<std::uint8_t> encryptedData(encryptedDataLength);

    // encrypt the data
    if (EVP_PKEY_encrypt(
        context.get(),
        encryptedData.data(),
        &encryptedDataLength,
        plainData.data(),
        plainData.size()
    ) <= 0)
        throw std::runtime_error("Could not encrypt data.");

    encryptedData.resize(encryptedDataLength);

    return encryptedData;
}


std::vector<std::uint8_t> RsaPublicKey::serialize() const {
    // serialize the members
    const auto serializedData = serialize::serializeVector(this->_data);
    const auto serializedPadMode = serialize::serializeObject<std::uint8_t>(static_cast<std::uint8_t>(this->padMode));

    // create a buffer to store our members
    std::vector<std::uint8_t> data;

    // store our members
    data.insert(data.end(), serializedData.begin(), serializedData.end());
    data.insert(data.end(), serializedPadMode.begin(), serializedPadMode.end());

    return data;
}

RsaPublicKey RsaPublicKey::deserialize(std::vector<std::uint8_t>& data) {
    // deserialize the members
    const auto keyData = serialize::deserializeVector<std::uint8_t>(data);
    const auto keyPadding = static_cast<int>(serialize::deserializeObject<std::uint8_t>(data));

    return RsaPublicKey(keyData, keyPadding);
}


[[nodiscard]] std::unique_ptr<EVP_PKEY, decltype(&EVP_PKEY_free)> RsaPublicKey::getOpenSslKey() const {
    // get the bio from the public key data
    const std::unique_ptr<BIO, decltype(&BIO_free)> bio(
        BIO_new_mem_buf(
            this->_data.data(),
            static_cast<int>(this->_data.size())
        ),
        BIO_free
    );
    if (bio == nullptr)
        throw std::runtime_error("Could not create BIO for public key.");

    // get the key from the bio
    EVP_PKEY* rawKey = nullptr;
    if (!d2i_PUBKEY_bio(
        bio.get(),
        &rawKey
    ))
        throw std::runtime_error("Could not deserialize RSA public key.");

    return {
        rawKey,
        EVP_PKEY_free
    };
}


}