#include "ssl.h"
#include "utils.h"

#include <nerutils/log.h>

#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/pem.h>
#include <openssl/x509v3.h>
#include <openssl/bn.h>
#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/ssl.h>

#include <random>
#include <cstdlib>
#include <ctime>
#include <cstring>

#include <wincrypt.h>
#include <processthreadsapi.h>

template <typename T, void (*f)(T*)> struct Deleter
{
        void operator()(T* p) const { f(p); }
};

using EVP_PKEY_ptr = std::unique_ptr<EVP_PKEY, Deleter<EVP_PKEY, EVP_PKEY_free>>;
using X509_ptr = std::unique_ptr<X509, Deleter<X509, X509_free>>;
using SSL_CTX_ptr = std::unique_ptr<SSL_CTX, Deleter<SSL_CTX, SSL_CTX_free>>;
using BIO_ptr = std::unique_ptr<BIO, Deleter<BIO, BIO_vfree>>;

CertManager::CertManager() : _sessionPkey(nullptr), _caPkey(nullptr), _caCert(nullptr) {}

CertManager::~CertManager()
{
    if (_caPkey) EVP_PKEY_free(_caPkey);
    if (_caCert) X509_free(_caCert);
    if (_sessionPkey) EVP_PKEY_free(_sessionPkey);

    for (auto& pair : _hostContexts)
        SSL_CTX_free(pair.second);
}

bool CertManager::init()
{
    EVP_PKEY_CTX* pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr);
    if (!pctx) return false;

    EVP_PKEY_keygen_init(pctx);
    EVP_PKEY_CTX_set_rsa_keygen_bits(pctx, 2048);

    EVP_PKEY* rawPkey = nullptr;
    if (EVP_PKEY_keygen(pctx, &rawPkey) <= 0)
    {
        Log::error("Failed to generate session key");
        EVP_PKEY_CTX_free(pctx);
        return false;
    }
    _sessionPkey = rawPkey;
    EVP_PKEY_CTX_free(pctx);

    if (loadCA())
    {
        Log::verbose("Loaded existing CA certificate");
        return true;
    }

    Log::verbose("No CA found, generating");
    return generateCA();
}

bool CertManager::loadCA()
{
    std::string path = utils::getExePath();

    BIO_ptr keyBio(BIO_new_file((path + "/key.pem").c_str(), "r"));
    if (!keyBio) return false;
    _caPkey = PEM_read_bio_PrivateKey(keyBio.get(), nullptr, nullptr, nullptr);


    BIO_ptr certBio(BIO_new_file((path + "/cert.pem").c_str(), "r"));
    if (!certBio) return false;
    _caCert = PEM_read_bio_X509(certBio.get(), nullptr, nullptr, nullptr);

    return (_caPkey && _caCert);
}

void CertManager::installCert(X509* cert)
{
    if (!cert) return;

    /*
        X509 to DER
    */
    int derLen = i2d_X509(cert, nullptr);
    if (derLen < 0) return;

    unsigned char* derBuf = new unsigned char[derLen];
    unsigned char* p = derBuf;
    i2d_X509(cert, &p);

    PCCERT_CONTEXT certCtx = CertCreateCertificateContext(X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, derBuf, derLen);
    if (certCtx)
    {
        HCERTSTORE rootStore = CertOpenStore(CERT_STORE_PROV_SYSTEM, 0, NULL, CERT_SYSTEM_STORE_CURRENT_USER, L"Root");
        if (rootStore)
        {
            BOOL success =
                CertAddCertificateContextToStore(rootStore, certCtx, CERT_STORE_ADD_REPLACE_EXISTING, NULL);

            if (success)
                Log::info("CA certificate installed");
            else
                Log::error("Failed to install CA certificate");

            CertCloseStore(rootStore, 0);
        }
        CertFreeCertificateContext(certCtx);
    }

    delete[] derBuf;
}

bool CertManager::generateCA()
{
    std::random_device rd;
    std::mt19937 gen(rd());

    /*
        key
    */
    EVP_PKEY_CTX* pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA, nullptr);
    EVP_PKEY_keygen_init(pctx);
    EVP_PKEY_CTX_set_rsa_keygen_bits(pctx, 2048);
    EVP_PKEY* rawCaKey = nullptr;
    EVP_PKEY_keygen(pctx, &rawCaKey);
    _caPkey = rawCaKey;
    EVP_PKEY_CTX_free(pctx);

    /*
        cert
    */
    X509_ptr cert(X509_new());
    X509_set_version(cert.get(), 2);
    ASN1_INTEGER_set(X509_get_serialNumber(cert.get()), 1);
    X509_gmtime_adj(X509_get_notBefore(cert.get()), 0);
    X509_gmtime_adj(X509_get_notAfter(cert.get()), 31536000L); // 1 year

    X509_NAME* name = X509_get_subject_name(cert.get());
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)"Debug Proxy CA", -1, -1, 0);
    X509_set_issuer_name(cert.get(), name);
    X509_set_pubkey(cert.get(), _caPkey);

    /*
        CA constraints
    */
    X509V3_CTX v3ctx;
    X509V3_set_ctx(&v3ctx, cert.get(), cert.get(), nullptr, nullptr, 0);
    X509_EXTENSION* ext = X509V3_EXT_conf_nid(nullptr, &v3ctx, NID_basic_constraints, "critical,CA:TRUE");
    X509_add_ext(cert.get(), ext, -1);
    X509_EXTENSION_free(ext);

    if (X509_sign(cert.get(), _caPkey, EVP_sha256()) <= 0) return false;

    /*
        out   
    */
    std::string path = utils::getExePath();
    BIO_ptr kOut(BIO_new_file((path + "/key.pem").c_str(), "w"));
    PEM_write_bio_PrivateKey(kOut.get(), _caPkey, nullptr, nullptr, 0, nullptr, nullptr);

    BIO_ptr cOut(BIO_new_file((path + "/cert.pem").c_str(), "w"));
    PEM_write_bio_X509(cOut.get(), cert.get());

    /*
        install and release
    */
    installCert(cert.get());

    _caCert = cert.release();

    return true;
}

SSL_CTX* CertManager::createHostContext(const std::string& host)
{
    std::lock_guard<std::mutex> lock(_mutex);
    if (_hostContexts.count(host)) return _hostContexts[host];

    /*
        cert base
    */
    X509_ptr cert(X509_new());
    X509_set_version(cert.get(), 2);

    ASN1_INTEGER_set(X509_get_serialNumber(cert.get()), static_cast<long>(std::hash<std::string>{}(host) & 0x7FFFFFFF));

    X509_gmtime_adj(X509_get_notBefore(cert.get()), 0);
    X509_gmtime_adj(X509_get_notAfter(cert.get()), 31536000L);

    X509_NAME* name = X509_get_subject_name(cert.get());
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)host.c_str(), -1, -1, 0);
    X509_set_issuer_name(cert.get(), X509_get_subject_name(_caCert));
    X509_set_pubkey(cert.get(), _sessionPkey);

    /*
        SAN
    */
    X509V3_CTX v3ctx;
    X509V3_set_ctx(&v3ctx, _caCert, cert.get(), nullptr, nullptr, 0);
    std::string altName = "DNS:" + host;
    X509_EXTENSION* ext = X509V3_EXT_conf_nid(nullptr, &v3ctx, NID_subject_alt_name, altName.c_str());
    X509_add_ext(cert.get(), ext, -1);
    X509_EXTENSION_free(ext);

    /*
        sign & ctx load
    */
    if (X509_sign(cert.get(), _caPkey, EVP_sha256()) <= 0) return nullptr;

    SSL_CTX* ctx = SSL_CTX_new(TLS_server_method());
    if (!ctx) return nullptr;

    if (SSL_CTX_use_certificate(ctx, cert.get()) <= 0 || SSL_CTX_use_PrivateKey(ctx, _sessionPkey) <= 0)
    {
        SSL_CTX_free(ctx);
        return nullptr;
    }

    _hostContexts[host] = ctx;
    return ctx;
}
