#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/engine.h>
#include <openssl/rsa.h>

// Function to generate RSA key
EVP_PKEY* generate_RSA_key() {
    EVP_PKEY* pkey = EVP_PKEY_new();
    RSA* rsa = RSA_new();
    BIGNUM* e = BN_new();

    BN_set_word(e, RSA_F4);

    if (RSA_generate_key_ex(rsa, 2048, e, NULL) != 1) {
        printf("Error generating RSA key.\n");
        return NULL;
    }

    if (!EVP_PKEY_assign_RSA(pkey, rsa)) {
        printf("Error assigning RSA key to EVP_PKEY.\n");
        return NULL;
    }

    BN_free(e);
    return pkey;
}

static int HasNoPrivateKey(const RSA* rsa)
{
    if (rsa == NULL)
        return 1;
    
    printf("Got RSA!\n");
    // Shared pointer, don't free.
    const RSA_METHOD* meth = RSA_get_method(rsa);

    // The method has described itself as having the private key external to the structure.
    // That doesn't mean it's actually present, but we can't tell.
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcast-qual"
    if (RSA_test_flags(rsa, RSA_FLAG_EXT_PKEY) || RSA_meth_get_flags((RSA_METHOD*)meth) & RSA_FLAG_EXT_PKEY)
#pragma clang diagnostic pop
    {
        return 0;
    }
    printf("RSA_FLAG_EXT_PKEY:%d\n",RSA_FLAG_EXT_PKEY);

    // In the event that there's a middle-ground where we report failure when success is expected,
    // one could do something like check if the RSA_METHOD intercepts all private key operations:
    //
    // * meth->rsa_priv_enc
    // * meth->rsa_priv_dec
    // * meth->rsa_sign (in 1.0.x this is only respected if the RSA_FLAG_SIGN_VER flag is asserted)
    //
    // But, for now, leave it at the EXT_PKEY flag test.

    // The module is documented as accepting either d or the full set of CRT parameters (p, q, dp, dq, qInv)
    // So if we see d, we're good. Otherwise, if any of the rest are missing, we're public-only.
    const BIGNUM* d;
    RSA_get0_key(rsa, NULL, NULL, &d);

    printf("d:%d\n", (d == NULL));
    if (d != NULL)
    {
        return 0;
    }

    const BIGNUM* p;
    const BIGNUM* q;
    const BIGNUM* dmp1;
    const BIGNUM* dmq1;
    const BIGNUM* iqmp;

    RSA_get0_factors(rsa, &p, &q);
    RSA_get0_crt_params(rsa, &dmp1, &dmq1, &iqmp);

    printf("Factors - p:%d q:%d\n", (p == NULL), (q == NULL));
    printf("CRT Params - dmp1:%d dmq1:%d iqmp:%d\n", (dmp1 == NULL), (dmq1 == NULL), (iqmp == NULL));
    if (p == NULL || q == NULL || dmp1 == NULL || dmq1 == NULL || iqmp == NULL)
    {
        return 1;
    }

    return 0;
}

static EVP_PKEY* LoadKeyFromEngine(
    const char* engineName,
    const char* keyName,
    ENGINE_LOAD_KEY_PTR load_func)
{
    ERR_clear_error();

    EVP_PKEY* ret = NULL;
    ENGINE* engine = NULL;

    // Per https://github.com/openssl/openssl/discussions/21427
    // using EVP_PKEY after freeing ENGINE is correct.
    engine = ENGINE_by_id(engineName);

    if (!engine) {
        fprintf(stderr,
                "Could not load %s Engine: %s\n", engineName,
                ERR_error_string(ERR_peek_last_error(), NULL));
    }

    if (engine != NULL)
    {
        ENGINE_ctrl_cmd_string(engine, "SO_PATH", "/usr/lib64/engines-1.1/pkcs11.so", 0);
        ENGINE_ctrl_cmd_string(engine, "MODULE_PATH", "/opt/nfast/toolkits/pkcs11/libcknfast.so", 0);

        if (ENGINE_init(engine))
        {
            ret = load_func(engine, keyName, NULL, NULL);
            
            ENGINE_finish(engine);
        }

        ENGINE_free(engine);
    }

    return ret;
}

EVP_PKEY* CryptoNative_LoadPrivateKeyFromEngine(const char* engineName, const char* keyName)
{
    return LoadKeyFromEngine(engineName, keyName, ENGINE_load_private_key);
}

EVP_PKEY* CryptoNative_LoadPublicKeyFromEngine(const char* engineName, const char* keyName)
{
    return LoadKeyFromEngine(engineName, keyName, ENGINE_load_public_key);
}

// Function to sign and verify using EVP_PKEY
int sign_and_verify(EVP_PKEY* pkey, unsigned char* data, size_t data_len) {
    EVP_MD_CTX* mdctx = EVP_MD_CTX_new();
    unsigned char* signature = NULL;
    size_t sig_len;

    // Signing
    if (EVP_DigestSignInit(mdctx, NULL, EVP_sha256(), NULL, pkey) != 1) {
        printf("Error initializing DigestSign.\n");
        return -1;
    }

    if (EVP_DigestSignUpdate(mdctx, data, data_len) != 1) {
        printf("Error updating DigestSign.\n");
        return -1;
    }

    if (EVP_DigestSignFinal(mdctx, NULL, &sig_len) != 1) {
        printf("Error finalizing DigestSign.\n");
        return -1;
    }

    signature = (unsigned char*)malloc(sig_len);

    if (EVP_DigestSignFinal(mdctx, signature, &sig_len) != 1) {
        printf("Error getting signature.\n");
        return -1;
    }

    // Verification
    EVP_MD_CTX_reset(mdctx);

    if (EVP_DigestVerifyInit(mdctx, NULL, EVP_sha256(), NULL, pkey) != 1) {
        printf("Error initializing DigestVerify.\n");
        return -1;
    }

    if (EVP_DigestVerifyUpdate(mdctx, data, data_len) != 1) {
        printf("Error updating DigestVerify.\n");
        return -1;
    }

    if (EVP_DigestVerifyFinal(mdctx, signature, sig_len) != 1) {
        printf("Signature verification failed.\n");
        return -1;
    }

    printf("Signature verification succeeded.\n");
    free(signature);
    EVP_MD_CTX_free(mdctx);
    return 0;
}

EVP_PKEY* get_RSA_key() {
    //return generate_RSA_key();
    return CryptoNative_LoadPrivateKeyFromEngine("pkcs11", "pkcs11:type=private;object=my-private-key");
}

int main() {
    ENGINE_load_builtin_engines();

    EVP_PKEY* pkey = get_RSA_key();
    if (pkey == NULL) {
        return -1;
    }
    printf("Main - Got Private Key!\n");

    RSA *rsa_key = EVP_PKEY_get1_RSA(pkey);
    int r = HasNoPrivateKey(rsa_key);
    printf("Main - HasNoPrivateKey Result:%d\n",r);

    unsigned char data[] = "Hello, world!";
    size_t data_len = strlen((char*)data);

    if (sign_and_verify(pkey, data, data_len) != 0) {
        return -2;
    }
    printf("Main - Got Signature!\n");

    EVP_PKEY_free(pkey);
    return 0;
}