#include "cert_manager.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 <cstdlib>
#include <ctime>

#include <processthreadsapi.h>
#include <cstring>

std::string randomizeString(size_t length)
{
    const char charset[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
    std::string result;
    result.resize(length);
    for (size_t i = 0; i < length; ++i)
        result[i] = charset[rand() % (sizeof(charset) - 1)];
    return result;
}

CertManager::CertManager() : _sessionPkey(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)
    {
        EVP_PKEY_keygen_init(pctx);
        EVP_PKEY_CTX_set_rsa_keygen_bits(pctx, 2048);
        EVP_PKEY_keygen(pctx, &_sessionPkey);
        EVP_PKEY_CTX_free(pctx);
    }

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

    Log::verbose("No CA found. Generating new CA certificate.");
    return GenerateCA();
}

bool CertManager::LoadCA()
{
    BIO* keyBio = BIO_new_file("ca_key.pem", "r");
    if (!keyBio) return false;

    _caPkey = PEM_read_bio_PrivateKey(keyBio, nullptr, nullptr, nullptr);
    BIO_free(keyBio);

    if (!_caPkey) return false;

    BIO* certBio = BIO_new_file("ca_cert.pem", "r");
    if (!certBio) return false;

    _caCert = PEM_read_bio_X509(certBio, nullptr, nullptr, nullptr);
    BIO_free(certBio);

    if (!_caCert) return false;

    return true;
}

bool CertManager::GenerateCA()
{
    srand(static_cast<unsigned int>(time(nullptr)));

    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_keygen(pctx, &_caPkey);
    EVP_PKEY_CTX_free(pctx);

    _caCert = X509_new();
    X509_set_version(_caCert, 2);
    ASN1_INTEGER_set(X509_get_serialNumber(_caCert), 1);
    X509_gmtime_adj(X509_get_notBefore(_caCert), 0);
    X509_gmtime_adj(X509_get_notAfter(_caCert), 31536000L); // 1 year

    std::string org = randomizeString(16);
    std::string cn = randomizeString(16);

    X509_NAME* name = X509_get_subject_name(_caCert);
    X509_NAME_add_entry_by_txt(name, "C", MBSTRING_ASC, (unsigned char*)"US", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "O", MBSTRING_ASC, (unsigned char*)org.c_str(), -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)cn.c_str(), -1, -1, 0);
    X509_set_issuer_name(_caCert, name);
    X509_set_pubkey(_caCert, _caPkey);

    X509V3_CTX ctxV3;
    X509V3_set_ctx_nodb(&ctxV3);
    X509V3_set_ctx(&ctxV3, _caCert, _caCert, nullptr, nullptr, 0);
    X509_EXTENSION* extCA = X509V3_EXT_conf_nid(nullptr, &ctxV3, NID_basic_constraints, "critical,CA:TRUE");
    if (extCA)
    {
        X509_add_ext(_caCert, extCA, -1);
        X509_EXTENSION_free(extCA);
    }

    X509_sign(_caCert, _caPkey, EVP_sha256());

    BIO* keyBioOut = BIO_new_file("ca_key.pem", "w");
    if (keyBioOut)
    {
        PEM_write_bio_PrivateKey(keyBioOut, _caPkey, nullptr, nullptr, 0, nullptr, nullptr);
        BIO_free(keyBioOut);
    }

    BIO* certBioOut = BIO_new_file("ca_cert.pem", "w");
    if (certBioOut)
    {
        PEM_write_bio_X509(certBioOut, _caCert);
        BIO_free(certBioOut);
    }

    Log::info("Generated new CA key and certificate files. Installing to Windows Root CA store automatically...");

    STARTUPINFOA si;
    memset(&si, 0, sizeof(si));
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi;
    memset(&pi, 0, sizeof(pi));
    char cmd[] = "certutil.exe -user -addstore root ca_cert.pem";
    if (CreateProcessA(NULL, cmd, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi))
    {
        WaitForSingleObject(pi.hProcess, INFINITE);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    }

    return true;
}

SSL_CTX* CertManager::CreateHostContext(const std::string& host)
{
    std::lock_guard<std::mutex> lock(_mutex);

    auto it = _hostContexts.find(host);
    if (it != _hostContexts.end())
    {
        return it->second;
    }

    EVP_PKEY* pkey = _sessionPkey;
    if (!pkey) return nullptr;

    X509* cert = X509_new();
    X509_set_version(cert, 2);
    ASN1_INTEGER_set(X509_get_serialNumber(cert), static_cast<long>(std::hash<std::string>{}(host) & 0x7FFFFFFF));
    X509_gmtime_adj(X509_get_notBefore(cert), 0);
    X509_gmtime_adj(X509_get_notAfter(cert), 31536000L);

    std::string dynamicOrg = randomizeString(16);

    X509_NAME* name = X509_get_subject_name(cert);
    X509_NAME_add_entry_by_txt(name, "C", MBSTRING_ASC, (unsigned char*)"US", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "O", MBSTRING_ASC, (unsigned char*)dynamicOrg.c_str(), -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)(host.c_str()), -1, -1, 0);
    X509_set_issuer_name(cert, X509_get_subject_name(_caCert));
    X509_set_pubkey(cert, pkey);

    X509V3_CTX ctxV3;
    X509V3_set_ctx_nodb(&ctxV3);
    X509V3_set_ctx(&ctxV3, _caCert, cert, nullptr, nullptr, 0);
    std::string san = "DNS:" + host;
    X509_EXTENSION* extSAN = X509V3_EXT_conf_nid(nullptr, &ctxV3, NID_subject_alt_name, san.c_str());
    if (extSAN)
    {
        X509_add_ext(cert, extSAN, -1);
        X509_EXTENSION_free(extSAN);
    }

    X509_sign(cert, _caPkey, EVP_sha256());

    SSL_CTX* ctx = SSL_CTX_new(TLS_server_method());

    SSL_CTX_set_alpn_select_cb(
        ctx,
        [](SSL* /*ssl*/, const unsigned char** out, unsigned char* outlen, const unsigned char* in, unsigned int inlen,
           void* /*arg*/) -> int {
            for (unsigned int i = 0; i < inlen;)
            {
                unsigned int len = in[i];
                if (len == 8 && memcmp(&in[i + 1], "http/1.1", 8) == 0)
                {
                    *out = &in[i + 1];
                    *outlen = (unsigned char)len;
                    return SSL_TLSEXT_ERR_OK;
                }
                i += len + 1;
            }
            return SSL_TLSEXT_ERR_NOACK;
        },
        nullptr);

    SSL_CTX_use_certificate(ctx, cert);
    SSL_CTX_use_PrivateKey(ctx, pkey);

    X509_free(cert);

    _hostContexts[host] = ctx;
    return ctx;
}