// src/Security/SslCertUtil.h
#pragma once

#include <string>
#include <vector>
#include <filesystem>
#include <fstream>
#include <stdexcept>

#include <sys/stat.h> // chmod

#include <openssl/evp.h>
#include <openssl/ec.h>
#include <openssl/pem.h>
#include <openssl/x509v3.h>
#include <openssl/rand.h>
#include <openssl/hmac.h>

#include <spdlog/spdlog.h>

namespace snoop
{
    namespace device_sec
    {

        class SslCertUtil
        {
        public:
            // Generates:
            //   - EC P-256 private key at keyPath (PEM, chmod 600)
            //   - CSR at csrPath with:
            //       CN = guid
            //       subjectAltName = URI:urn:device:<guid>
            static void GenerateEcKeyAndCsr(const std::string &guid,
                                            const std::filesystem::path &keyPath,
                                            const std::filesystem::path &csrPath)
            {
                spdlog::info("SslCertUtil: generating EC key + CSR for GUID={}", guid);

                EVP_PKEY *pkey = nullptr;
                EC_KEY *ec = nullptr;

                try
                {
                    // ----- 1) Generate EC key (prime256v1) -----
                    ec = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
                    if (!ec)
                    {
                        throw std::runtime_error("EC_KEY_new_by_curve_name failed");
                    }
                    if (EC_KEY_generate_key(ec) != 1)
                    {
                        throw std::runtime_error("EC_KEY_generate_key failed");
                    }

                    pkey = EVP_PKEY_new();
                    if (!pkey)
                    {
                        throw std::runtime_error("EVP_PKEY_new failed");
                    }
                    if (EVP_PKEY_assign_EC_KEY(pkey, ec) != 1)
                    {
                        throw std::runtime_error("EVP_PKEY_assign_EC_KEY failed");
                    }
                    // pkey now owns ec
                    ec = nullptr; // prevent double free

                    // ----- 5) Write key + CSR to disk -----
                    SavePrivateKeyPem(pkey, keyPath);
                    BuildAndSaveCsr(pkey, guid, csrPath);

                    spdlog::info("SslCertUtil: key written to '{}', csr written to '{}'",
                                 keyPath.string(), csrPath.string());
                }
                catch (...)
                {
                    if (pkey)
                    {
                        EVP_PKEY_free(pkey);
                    }
                    if (ec)
                    {
                        EC_KEY_free(ec);
                    }
                    throw; // rethrow
                }

                EVP_PKEY_free(pkey);
            }

            // CSR from an existing key
            static void GenerateCsrFromExistingKey(const std::string &guid,
                                                   const std::filesystem::path &keyPath,
                                                   const std::filesystem::path &csrPath)
            {
                spdlog::info("SslCertUtil: generating CSR from existing key '{}'", keyPath.string());

                FILE *kf = fopen(keyPath.string().c_str(), "rb");
                if (!kf)
                {
                    throw std::runtime_error("GenerateCsrFromExistingKey: cannot open key file: " +
                                             keyPath.string());
                }

                EVP_PKEY *pkey = PEM_read_PrivateKey(kf, nullptr, nullptr, nullptr);
                fclose(kf);

                if (!pkey)
                {
                    throw std::runtime_error("GenerateCsrFromExistingKey: PEM_read_PrivateKey failed");
                }

                try
                {
                    BuildAndSaveCsr(pkey, guid, csrPath);
                }
                catch (...)
                {
                    EVP_PKEY_free(pkey);
                    throw;
                }

                EVP_PKEY_free(pkey);
                spdlog::info("SslCertUtil: CSR written to '{}'", csrPath.string());
            }

            // Encrypt file with AES-256-CBC + PBKDF2 + salt.
            //
            // Equivalent to:
            //   openssl enc -aes-256-cbc -pbkdf2 -salt -pass pass:<password> -in in -out out
            //
            // NOTE: Does NOT shred or delete the input file; caller should do that
            //       (you’ll still call `shred -u` from EnrollmentService).
            static void EncryptFileAes256CbcPbkdf2(const std::filesystem::path &inPath,
                                                   const std::filesystem::path &outPath,
                                                   const std::string &password)
            {
                spdlog::info("SslCertUtil: encrypting '{}' -> '{}' (AES-256-CBC + PBKDF2)",
                             inPath.string(), outPath.string());

                // ----- 1) Read input file -----
                std::ifstream fin(inPath, std::ios::binary);
                if (!fin)
                {
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: cannot open input file: " +
                                             inPath.string());
                }
                std::vector<unsigned char> plaintext(
                    (std::istreambuf_iterator<char>(fin)),
                    std::istreambuf_iterator<char>());
                fin.close();

                // ----- 2) Prepare salt, key, iv -----
                unsigned char salt[8];
                if (RAND_bytes(salt, sizeof(salt)) != 1)
                {
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: RAND_bytes failed");
                }

                const EVP_CIPHER *cipher = EVP_aes_256_cbc();
                const int keyLen = EVP_CIPHER_key_length(cipher);
                const int ivLen = EVP_CIPHER_iv_length(cipher);

                std::vector<unsigned char> keyiv(keyLen + ivLen);

                const int iterations = 10000; // PBKDF2 rounds
                if (PKCS5_PBKDF2_HMAC(password.c_str(),
                                      static_cast<int>(password.size()),
                                      salt, sizeof(salt),
                                      iterations,
                                      EVP_sha256(),
                                      keyLen + ivLen,
                                      keyiv.data()) != 1)
                {
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: PKCS5_PBKDF2_HMAC failed");
                }

                unsigned char *key = keyiv.data();
                unsigned char *iv = keyiv.data() + keyLen;

                // ----- 3) Encrypt -----
                EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
                if (!ctx)
                {
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: EVP_CIPHER_CTX_new failed");
                }

                std::vector<unsigned char> ciphertext(plaintext.size() + EVP_CIPHER_block_size(cipher));
                int outLen1 = 0;
                int outLen2 = 0;

                if (EVP_EncryptInit_ex(ctx, cipher, nullptr, key, iv) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: EVP_EncryptInit_ex failed");
                }

                if (EVP_EncryptUpdate(ctx,
                                      ciphertext.data(), &outLen1,
                                      plaintext.data(), static_cast<int>(plaintext.size())) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: EVP_EncryptUpdate failed");
                }

                if (EVP_EncryptFinal_ex(ctx,
                                        ciphertext.data() + outLen1, &outLen2) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: EVP_EncryptFinal_ex failed");
                }

                EVP_CIPHER_CTX_free(ctx);
                ciphertext.resize(outLen1 + outLen2);

                // ----- 4) Write output: [salt][ciphertext] -----
                std::ofstream fout(outPath, std::ios::binary | std::ios::trunc);
                if (!fout)
                {
                    throw std::runtime_error("EncryptFileAes256CbcPbkdf2: cannot open output file: " +
                                             outPath.string());
                }

                // Just raw salt followed by ciphertext
                fout.write(reinterpret_cast<const char *>(salt), sizeof(salt));
                fout.write(reinterpret_cast<const char *>(ciphertext.data()),
                           static_cast<std::streamsize>(ciphertext.size()));
                fout.close();

                spdlog::info("SslCertUtil: encryption finished, wrote {}", outPath.string());
            }

            // Compute HMAC-SHA256 and return lowercase hex string.
            // Equivalent to:
            //   printf %s "<data>" | openssl dgst -sha256 -hmac "<key>" | awk '{print $2}'
            static std::string ComputeHmacSha256Hex(const std::string &key,
                                                    const std::string &data)
            {
                unsigned char mac[EVP_MAX_MD_SIZE];
                unsigned int macLen = 0;

                if (!HMAC(EVP_sha256(),
                          key.data(), static_cast<int>(key.size()),
                          reinterpret_cast<const unsigned char *>(data.data()),
                          data.size(),
                          mac, &macLen))
                {
                    throw std::runtime_error("ComputeHmacSha256Hex: HMAC(EVP_sha256) failed");
                }

                return ToHexLower(mac, macLen);
            }

            // Decrypt file with AES-256-CBC + PBKDF2 + salt.
            // Layout: [8-byte salt][ciphertext]
            // Password is KEK (hex string from HMAC).
            static void DecryptFileAes256CbcPbkdf2(const std::filesystem::path &inPath,
                                                   const std::filesystem::path &outPath,
                                                   const std::string &password)
            {
                spdlog::info("SslCertUtil: decrypting '{}' -> '{}' (AES-256-CBC + PBKDF2)",
                             inPath.string(), outPath.string());

                // ----- 1) Read input file -----
                std::ifstream fin(inPath, std::ios::binary);
                if (!fin)
                {
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: cannot open input file: " +
                                             inPath.string());
                }
                std::vector<unsigned char> enc(
                    (std::istreambuf_iterator<char>(fin)),
                    std::istreambuf_iterator<char>());
                fin.close();

                if (enc.size() < 8)
                {
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: encrypted file too small");
                }

                // First 8 bytes = salt, rest = ciphertext
                unsigned char salt[8];
                std::copy(enc.begin(), enc.begin() + 8, salt);
                std::vector<unsigned char> ciphertext(enc.begin() + 8, enc.end());

                // ----- 2) Derive key/iv -----
                const EVP_CIPHER *cipher = EVP_aes_256_cbc();
                const int keyLen = EVP_CIPHER_key_length(cipher);
                const int ivLen = EVP_CIPHER_iv_length(cipher);

                std::vector<unsigned char> keyiv(keyLen + ivLen);

                const int iterations = 10000;
                if (PKCS5_PBKDF2_HMAC(password.c_str(),
                                      static_cast<int>(password.size()),
                                      salt, sizeof(salt),
                                      iterations,
                                      EVP_sha256(),
                                      keyLen + ivLen,
                                      keyiv.data()) != 1)
                {
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: PKCS5_PBKDF2_HMAC failed");
                }

                unsigned char *key = keyiv.data();
                unsigned char *iv = keyiv.data() + keyLen;

                // ----- 3) Decrypt -----
                EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
                if (!ctx)
                {
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: EVP_CIPHER_CTX_new failed");
                }

                std::vector<unsigned char> plaintext(ciphertext.size() + EVP_CIPHER_block_size(cipher));
                int outLen1 = 0;
                int outLen2 = 0;

                if (EVP_DecryptInit_ex(ctx, cipher, nullptr, key, iv) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: EVP_DecryptInit_ex failed");
                }

                if (EVP_DecryptUpdate(ctx,
                                      plaintext.data(), &outLen1,
                                      ciphertext.data(), static_cast<int>(ciphertext.size())) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: EVP_DecryptUpdate failed");
                }

                if (EVP_DecryptFinal_ex(ctx,
                                        plaintext.data() + outLen1, &outLen2) != 1)
                {
                    EVP_CIPHER_CTX_free(ctx);
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: EVP_DecryptFinal_ex failed");
                }

                EVP_CIPHER_CTX_free(ctx);
                plaintext.resize(outLen1 + outLen2);

                // ----- 4) Write plaintext -----
                std::filesystem::create_directories(outPath.parent_path());
                std::ofstream fout(outPath, std::ios::binary | std::ios::trunc);
                if (!fout)
                {
                    throw std::runtime_error("DecryptFileAes256CbcPbkdf2: cannot open output file: " +
                                             outPath.string());
                }

                fout.write(reinterpret_cast<const char *>(plaintext.data()),
                           static_cast<std::streamsize>(plaintext.size()));
                fout.close();

                spdlog::info("SslCertUtil: decryption finished, wrote {}", outPath.string());
            }

        private:
            static void BuildAndSaveCsr(EVP_PKEY *pkey,
                                        const std::string &guid,
                                        const std::filesystem::path &csrPath)
            {
                X509_REQ *req = X509_REQ_new();
                if (!req)
                {
                    throw std::runtime_error("BuildAndSaveCsr: X509_REQ_new failed");
                }

                try
                {
                    // Subject: CN = GUID
                    X509_NAME *name = X509_NAME_new();
                    if (!name)
                    {
                        throw std::runtime_error("BuildAndSaveCsr: X509_NAME_new failed");
                    }
                    if (X509_NAME_add_entry_by_NID(
                            name,
                            NID_commonName,
                            MBSTRING_ASC,
                            reinterpret_cast<const unsigned char *>(guid.c_str()),
                            -1, -1, 0) != 1)
                    {
                        X509_NAME_free(name);
                        throw std::runtime_error("BuildAndSaveCsr: X509_NAME_add_entry_by_NID(CN) failed");
                    }
                    if (X509_REQ_set_subject_name(req, name) != 1)
                    {
                        X509_NAME_free(name);
                        throw std::runtime_error("BuildAndSaveCsr: X509_REQ_set_subject_name failed");
                    }
                    X509_NAME_free(name);

                    // Public key
                    if (X509_REQ_set_pubkey(req, pkey) != 1)
                    {
                        throw std::runtime_error("BuildAndSaveCsr: X509_REQ_set_pubkey failed");
                    }

                    // subjectAltName = URI:urn:device:<GUID>
                    X509V3_CTX ctx;
                    X509V3_set_ctx_nodb(&ctx);
                    X509V3_set_ctx(&ctx, nullptr, nullptr, req, nullptr, 0);

                    std::string sanStr = "URI:urn:device:" + guid;
                    X509_EXTENSION *ext = X509V3_EXT_conf_nid(
                        nullptr, &ctx, NID_subject_alt_name,
                        const_cast<char *>(sanStr.c_str()));
                    if (!ext)
                    {
                        throw std::runtime_error("BuildAndSaveCsr: X509V3_EXT_conf_nid(subjectAltName) failed");
                    }

                    STACK_OF(X509_EXTENSION) *exts = sk_X509_EXTENSION_new_null();
                    if (!exts)
                    {
                        X509_EXTENSION_free(ext);
                        throw std::runtime_error("BuildAndSaveCsr: sk_X509_EXTENSION_new_null failed");
                    }
                    sk_X509_EXTENSION_push(exts, ext);
                    if (X509_REQ_add_extensions(req, exts) != 1)
                    {
                        sk_X509_EXTENSION_pop_free(exts, X509_EXTENSION_free);
                        throw std::runtime_error("BuildAndSaveCsr: X509_REQ_add_extensions failed");
                    }
                    sk_X509_EXTENSION_pop_free(exts, X509_EXTENSION_free);

                    // Sign with SHA-256
                    if (X509_REQ_sign(req, pkey, EVP_sha256()) <= 0)
                    {
                        throw std::runtime_error("BuildAndSaveCsr: X509_REQ_sign failed");
                    }

                    SaveCsrPem(req, csrPath);
                }
                catch (...)
                {
                    X509_REQ_free(req);
                    throw;
                }

                X509_REQ_free(req);
            }

            static void SavePrivateKeyPem(EVP_PKEY *pkey, const std::filesystem::path &keyPath)
            {
                const auto parent = keyPath.parent_path();
                if (!parent.empty())
                {
                    std::filesystem::create_directories(parent);
                }

                FILE *f = fopen(keyPath.string().c_str(), "wb");
                if (!f)
                {
                    throw std::runtime_error("SavePrivateKeyPem: fopen failed: " + keyPath.string());
                }

                if (PEM_write_PrivateKey(f, pkey, nullptr, nullptr, 0, nullptr, nullptr) != 1)
                {
                    fclose(f);
                    throw std::runtime_error("SavePrivateKeyPem: PEM_write_PrivateKey failed");
                }

                fclose(f);

                // chmod 600
                ::chmod(keyPath.string().c_str(), 0600);
            }

            static void SaveCsrPem(X509_REQ *req, const std::filesystem::path &csrPath)
            {
                const auto parent = csrPath.parent_path();
                if (!parent.empty())
                {
                    std::filesystem::create_directories(parent);
                }

                FILE *f = fopen(csrPath.string().c_str(), "wb");
                if (!f)
                {
                    throw std::runtime_error("SaveCsrPem: fopen failed: " + csrPath.string());
                }

                if (PEM_write_X509_REQ(f, req) != 1)
                {
                    fclose(f);
                    throw std::runtime_error("SaveCsrPem: PEM_write_X509_REQ failed");
                }

                fclose(f);
            }

            static std::string ToHexLower(const unsigned char *buf, size_t len)
            {
                static const char *hex = "0123456789abcdef";
                std::string out;
                out.reserve(len * 2);

                for (size_t i = 0; i < len; ++i)
                {
                    unsigned char b = buf[i];
                    out.push_back(hex[b >> 4]);
                    out.push_back(hex[b & 0x0F]);
                }
                return out;
            }
        };

    } // namespace device_sec
} // namespace snoop
