// src/Services/EnrollmentService.h
#pragma once

#include <string>
#include <vector>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <regex>
#include <memory>

#include <spdlog/spdlog.h>
#include <httplib.h>
#include <stdio.h>
#include <sys/wait.h>

#include <openssl/x509.h>
#include <openssl/pem.h>
#include <openssl/asn1.h>

#include "ConfigService.h"
#include "Security/SslCertUtil.h"
#include "Security/TlsKeyUtil.h"

namespace snoop
{

    class EnrollmentService
    {
        std::shared_ptr<ConfigService> m_cfg;

        static std::string Exec(const std::string &cmd)
        {
            std::array<char, 4096> buf{};
            std::string out;
            FILE *pipe = popen((cmd + " 2>&1").c_str(), "r");
            if (!pipe)
            {
                throw std::runtime_error("popen failed: " + cmd);
            } 
            while (fgets(buf.data(), (int)buf.size(), pipe) != nullptr)
            {
                out.append(buf.data());
            }
            auto rc = pclose(pipe);
            int exitCode = WIFEXITED(rc) ? WEXITSTATUS(rc) : rc;
            if (exitCode != 0)
            {
                spdlog::warn("Command '{}' exited with code {}", cmd, exitCode);
            }    
            return out;
        }

        static void WriteFile(const std::filesystem::path &p, const std::string &data, bool createParents = true)
        {
            if (createParents)
            {
                std::filesystem::create_directories(p.parent_path());
            } 
            std::ofstream f(p, std::ios::binary);
            if (!f)
            {
                throw std::runtime_error("Cannot open file for write: " + p.string());
            }
            f.write(data.data(), (std::streamsize)data.size());
        }

        static std::string ReadFile(const std::filesystem::path &p)
        {
            std::ifstream f(p, std::ios::binary);
            if (!f)
            {
                throw std::runtime_error("Cannot open file for read: " + p.string());
            } 
            std::ostringstream ss;
            ss << f.rdbuf();
            return ss.str();
        }

        static std::string Trim(const std::string &s)
        {
            auto b = s.find_first_not_of(" \t\r\n");
            auto e = s.find_last_not_of(" \t\r\n");
            if (b == std::string::npos)
                return "";
            return s.substr(b, e - b + 1);
        }

        static std::string ParseCpuSerial()
        {
            // Default for “empty” case: 12 times '9'
            std::string fallback(12, '9');

            std::ifstream f("/proc/cpuinfo");
            if (!f)
            {
                spdlog::warn("/proc/cpuinfo not available, using fallback serial");
                return fallback;
            }
            std::string line;
            std::regex re(R"(^\s*Serial\s*:\s*([0-9A-Fa-f]+)\s*$)");
            while (std::getline(f, line))
            {
                std::smatch m;
                if (std::regex_match(line, m, re) && m.size() == 2)
                {
                    auto serial = Trim(m[1].str());
                    if (!serial.empty())
                    {
                        return serial;
                    }                     
                }
            }
            spdlog::warn("CPU Serial not found, using fallback");
            return fallback;
        }

    public:
        explicit EnrollmentService(std::shared_ptr<ConfigService> cfg) : m_cfg(std::move(cfg)) {}

        // Returns true if enrollment was performed now; false if it was already done
        bool EnsureEnrolled()
        {
            const std::string guid = m_cfg->GetGuid();
            const std::filesystem::path keystoreDir = "/etc/iot/keys";
            const auto encKeyPath = keystoreDir / "device.key.enc";
            const auto certPath = keystoreDir / "device.crt.pem";
            const auto caPath = keystoreDir / "issuing_ca.pem";
            const auto chainPath = keystoreDir / "ca_chain.pem";

            // If we already have encrypted key + cert, assume done
            if (std::filesystem::exists(encKeyPath) &&
                std::filesystem::exists(certPath))
            {
                spdlog::info("Enrollment already completed.");
                return false;
            }

            spdlog::info("Starting first-run enrollment...");

            // // 1) Run gen_device_csr.sh <GUID>
            // {
            //     // Assumes script is in the working dir or in PATH
            //     const std::string cmd = "bash ./gen_device_csr.sh " + guid;
            //     spdlog::info("Executing: {}", cmd);
            //     auto out = Exec(cmd);
            //     spdlog::debug("gen_device_csr.sh output:\n{}", out);
            // }
            const std::string keyName = "device_" + guid + ".key";
            const std::string csrName = "device_" + guid + ".csr";

            // if (!std::filesystem::exists(keyName) || !std::filesystem::exists(csrName))
            // {
            //     throw std::runtime_error("CSR or key was not generated by gen_device_csr.sh");
            // }

            // 1) Generate key + CSR via OpenSSL (no external script)
            // {
            //     spdlog::info("Generating device key + CSR via SslCertUtil");
            //     snoop::device_sec::SslCertUtil::GenerateEcKeyAndCsr(
            //         guid,
            //         std::filesystem::path(keyName),
            //         std::filesystem::path(csrName));
            // }

            // if (!std::filesystem::exists(keyName) || !std::filesystem::exists(csrName))
            // {
            //     throw std::runtime_error("CSR or key was not generated");
            // }

            // 1) CPU serial (awk '/Serial/ {print $3}' /proc/cpuinfo), if empty -> 9 * 12
            const std::string cpuSerial = ParseCpuSerial();
            spdlog::info("CPU_SERIAL = {}", cpuSerial);

            // 3) KEK = HMAC-SHA256(cpuSerial, key=GUID)  (bash equiv used)
            // std::string kek;
            // {
            //     // careful to avoid trailing newline
            //     const std::string cmd = "printf %s \"" + cpuSerial + "\""
            //                                                          " | openssl dgst -sha256 -hmac \"" +
            //                             guid + "\" | awk '{print $2}'";
            //     spdlog::info("Deriving KEK with HMAC-SHA256");
            //     kek = Trim(Exec(cmd));
            //     if (kek.empty())
            //         throw std::runtime_error("Failed to derive KEK");
            //     spdlog::debug("KEK (hex) = {}", kek);
            // }

            // 2) KEK = HMAC-SHA256(cpuSerial, key=GUID)
            std::string kek;
            {
                spdlog::info("Deriving KEK with HMAC-SHA256 (OpenSSL API)");
                kek = snoop::device_sec::SslCertUtil::ComputeHmacSha256Hex(guid, cpuSerial);
                if (kek.empty())
                {
                    throw std::runtime_error("Failed to derive KEK");
                }
                spdlog::debug("KEK (hex) = {}", kek);
            }

            // // 4) Encrypt the private key to /etc/iot/keys/device.key.enc using KEK; shred original
            // {
            //     std::filesystem::create_directories(keystoreDir);
            //     const std::string cmd =
            //         "openssl enc -aes-256-cbc -pbkdf2 -salt "
            //         "-pass pass:" +
            //         kek + " "
            //               "-in " +
            //         keyName + " "
            //                   "-out " +
            //         encKeyPath.string() + " "
            //                               "&& shred -u " +
            //         keyName;
            //     spdlog::info("Encrypting private key and shredding plaintext...");
            //     auto out = Exec(cmd);
            //     spdlog::debug("openssl enc output:\n{}", out);

            //     if (!std::filesystem::exists(encKeyPath))
            //     {
            //         throw std::runtime_error("Encrypted key not created: " + encKeyPath.string());
            //     }
            // }

            // 3) Generate key + CSR via OpenSSL (no external script)
            if (!std::filesystem::exists(encKeyPath))
            {
                // ---- Fresh key: generate + encrypt ----
                spdlog::info("No encrypted key yet -> generating new EC key + CSR");
                snoop::device_sec::SslCertUtil::GenerateEcKeyAndCsr(
                    guid,
                    std::filesystem::path(keyName),
                    std::filesystem::path(csrName));

                if (!std::filesystem::exists(keyName) || !std::filesystem::exists(csrName))
                {
                    throw std::runtime_error("CSR or key was not generated");
                }

                std::filesystem::create_directories(keystoreDir);

                spdlog::info("Encrypting private key via SslCertUtil (AES-256-CBC + PBKDF2)...");
                snoop::device_sec::SslCertUtil::EncryptFileAes256CbcPbkdf2(
                    std::filesystem::path(keyName),
                    encKeyPath,
                    kek);

                if (!std::filesystem::exists(encKeyPath))
                {
                    throw std::runtime_error("Encrypted key not created: " + encKeyPath.string());
                }
                // shred plaintext key
                const std::string shredCmd = "shred -u " + keyName;
                spdlog::info("Shredding plaintext private key with '{}'", shredCmd);
                auto out = Exec(shredCmd);
                spdlog::debug("shred output:\n{}", out);
            }
            else
            {
                // ---- Key already exists: reuse it ----
                spdlog::info("Encrypted key already exists -> reusing key, regenerating CSR only");

                // Decrypt to temp file (in /run/iot)
                snoop::device_sec::TempFile tf(std::filesystem::path("/run/iot"));
                snoop::device_sec::SslCertUtil::DecryptFileAes256CbcPbkdf2(
                    encKeyPath,
                    tf.path,
                    kek);

                // Generate CSR from existing key
                snoop::device_sec::SslCertUtil::GenerateCsrFromExistingKey(
                    guid,
                    tf.path,
                    std::filesystem::path(csrName));

                if (!std::filesystem::exists(csrName))
                {
                    throw std::runtime_error("CSR was not generated from existing key");
                }
                // TempFile destructor will securely delete the decrypted key file
            }

            // // 4) Encrypt the private key to /etc/iot/keys/device.key.enc using KEK; shred original
            // {
            //     std::filesystem::create_directories(keystoreDir);

            //     spdlog::info("Encrypting private key via SslCertUtil (AES-256-CBC + PBKDF2)...");
            //     snoop::device_sec::SslCertUtil::EncryptFileAes256CbcPbkdf2(
            //         std::filesystem::path(keyName),
            //         encKeyPath,
            //         kek);

            //     if (!std::filesystem::exists(encKeyPath))
            //     {
            //         throw std::runtime_error("Encrypted key not created: " + encKeyPath.string());
            //     }

            //     // still use shred to securely remove plaintext key, as requested
            //     const std::string shredCmd = "shred -u " + keyName;
            //     spdlog::info("Shredding plaintext private key with '{}'", shredCmd);
            //     auto out = Exec(shredCmd);
            //     spdlog::debug("shred output:\n{}", out);
            // }

            // 4.5) Ensure key is loaded into kernel keyring (first creation)
            try
            {
                bool loaded = snoop::device_sec::EnsureClientKeyInKernelKeyring(guid);
                if (loaded)
                {
                    spdlog::info("EnsureEnrolled: client key loaded into kernel keyring");
                }
                else
                {
                    spdlog::info("EnsureEnrolled: client key was already in keyring");
                }
            }
            catch (const std::exception &e)
            {
                spdlog::warn("EnsureEnrolled: failed to load key into keyring: {}", e.what());
            }

            // 5) Send CSR to /enroll/:guid as multipart form, field `json: {"csr":"<...>"}`
            std::string csrPem = ReadFile(csrName);
            {
                httplib::Client cli(m_cfg->GetBaseUrl());
                cli.enable_server_certificate_verification(false);

                httplib::MultipartFormDataItems items = {
                    // name, content, filename, content_type
                    {"csr", csrPem, "request.csr", "text/plain"}
                    // or "application/pkcs10" if you prefer:   "application/pkcs10"
                };

                const std::string path = "/api/enroll/" + guid;
                spdlog::info("POST {} (multipart)", path);
                auto res = cli.Post(path.c_str(), items);
                if (!res)
                {
                    throw std::runtime_error("Enroll request failed (no response)");
                }  
                spdlog::info("Enroll response status: {}", res->status);
                if (res->status != 200 && res->status != 201)
                {
                    spdlog::info("Enroll error: " + res->body);
                    throw std::runtime_error("Enroll failed: HTTP " + std::to_string(res->status));
                }

                // 6) Expect JSON: { "certificate": "...", "issuing_ca":"...", "ca_chain":"..." }
                nlohmann::json j;
                try
                {
                    j = nlohmann::json::parse(res->body);
                }
                catch (const std::exception &e)
                {
                    throw std::runtime_error(std::string("Enroll: invalid JSON: ") + e.what());
                }
                if (!j.contains("certificate") || !j["certificate"].is_string())
                {
                    throw std::runtime_error("Enroll response missing or invalid 'certificate'");
                }
                WriteFile(certPath, j["certificate"].get<std::string>());
                if (j.contains("issuing_ca"))
                {
                    if (!j["issuing_ca"].is_string())
                    {
                        throw std::runtime_error("'issuing_ca' must be a string PEM");
                    }                   
                    WriteFile(caPath, j["issuing_ca"].get<std::string>());
                }

                if (j.contains("ca_chain"))
                {
                    std::string chainPem;
                    if (j["ca_chain"].is_string())
                    {
                        chainPem = j["ca_chain"].get<std::string>();
                    }
                    else if (j["ca_chain"].is_array())
                    {
                        chainPem = JoinPemArray(j["ca_chain"]);
                    }
                    else
                    {
                        throw std::runtime_error("'ca_chain' must be string or array of strings");
                    }

                    // If server returned both issuing_ca and ca_chain that doesn't include it,
                    // you can choose to prepend issuing_ca to the chain file:
                    // if (std::filesystem::exists(caPath)) {
                    //     chainPem = ReadFileToString(caPath) + (chainPem.size() && chainPem.front()!='\n' ? "\n" : "") + chainPem;
                    // }

                    WriteFile(chainPath, chainPem);
                }
            }

            // (Optional) cleanup CSR after successful enrollment
            // try { std::filesystem::remove(csrName); } catch (...) {}

            return true;
        }

        // Renew cert either because it's about to expire (force=false) or
        // because server commanded us (force=true).
        // Returns true if renew was done and new files were written.
        bool RenewCertificate(bool force = false)
        {
            const std::string guid = m_cfg->GetGuid();
            const std::filesystem::path keystoreDir = "/etc/iot/keys";
            const auto certPath = keystoreDir / "device.crt.pem";
            const auto caPath = keystoreDir / "issuing_ca.pem";
            const auto chainPath = keystoreDir / "ca_chain.pem";

            // ----- 1) expiry check (case 1) -----
            if (!force)
            {
                if (!std::filesystem::exists(certPath))
                {
                    spdlog::warn("RenewCertificate: cert not found, skip (not forced)");
                    return false;
                }
                if (!IsCertExpiring(certPath, 1))
                {
                    spdlog::info("RenewCertificate: cert is fine, skip");
                    return false;
                }
                spdlog::info("RenewCertificate: cert expiring -> will renew");
            }
            else
            {
                spdlog::info("RenewCertificate: forced renew");
            }

            // ----- 2) re-generate CSR via your bash script -----
            // {
            //     const std::string cmd = "bash ./gen_device_csr.sh " + guid;
            //     spdlog::info("Executing (renew): {}", cmd);
            //     auto out = Exec(cmd);
            //     spdlog::debug("gen_device_csr.sh (renew) output:\n{}", out);
            // }
            // const std::string csrName = "device_" + guid + ".csr";
            // if (!std::filesystem::exists(csrName))
            // {
            //     throw std::runtime_error("Renew: CSR was not generated");
            // }
            // std::string csrPem = ReadFile(csrName);

            // ----- 2) re-generate CSR via OpenSSL (no bash) -----
            const std::string csrName = "device_" + guid + ".csr";
            {
                spdlog::info("RenewCertificate: generating new CSR via SslCertUtil");
                // We must reuse the existing key or generate a new one?
                // Current flow (bash script) always regenerated key+csr, so we mimic that:
                snoop::device_sec::SslCertUtil::GenerateEcKeyAndCsr(
                    guid,
                    std::filesystem::path("device_" + guid + ".key"),
                    std::filesystem::path(csrName));
            }

            if (!std::filesystem::exists(csrName))
            {
                throw std::runtime_error("Renew: CSR was not generated");
            }
            std::string csrPem = ReadFile(csrName);

            // ----- 3) build mTLS client using current device cert + keyring key -----
            const std::string base = m_cfg->GetBaseUrl();
            auto pu = ParseBaseUrl(base);
            auto cli = MakeMtlsClient(pu);

            // ----- 4) send multipart to /api/renew/:guid -----
            const std::string path = "/api/renew/" + guid;
            httplib::MultipartFormDataItems items = {
                {"csr", csrPem, "request.csr", "text/plain"}};

            spdlog::info("POST {} (multipart, mTLS)", path);

            auto res = cli->Post(path.c_str(), items);

            if (!res)
            {
                throw std::runtime_error("Renew request failed (no response)");
            }  
            spdlog::info("Renew response status: {}", res->status);
            if (res->status != 200)
            {
                throw std::runtime_error("Renew failed: HTTP " + std::to_string(res->status));
            }

            nlohmann::json j = nlohmann::json::parse(res->body);

            if (!j.contains("certificate") || !j["certificate"].is_string())
            {
                throw std::runtime_error("Renew response missing or invalid 'certificate'");
            }       
            WriteFile(certPath, j["certificate"].get<std::string>());

            if (j.contains("issuing_ca") && j["issuing_ca"].is_string())
            {
                WriteFile(caPath, j["issuing_ca"].get<std::string>());
            }

            if (j.contains("ca_chain"))
            {
                std::string chainPem;
                if (j["ca_chain"].is_string())
                {
                    chainPem = j["ca_chain"].get<std::string>();
                }
                else if (j["ca_chain"].is_array())
                {
                    chainPem = JoinPemArray(j["ca_chain"]);
                }
                else
                {
                    throw std::runtime_error("'ca_chain' must be string or array");
                }
                WriteFile(chainPath, chainPem);
            }

            if (j.contains("old_serial") && j["old_serial"].is_string())
            {
                spdlog::info("Renew: server says old_serial={}", j["old_serial"].get<std::string>());
            }

            spdlog::info("RenewCertificate: success");
            return true;
        }

    private:
        // very small helper to pack PEM into JSON string value (single-line, escaped quotes & newlines)
        static std::string EscapeJsonForOneLine(const std::string &in)
        {
            std::string out;
            out.reserve(in.size() * 2);
            for (char c : in)
            {
                switch (c)
                {
                case '\\':
                    out += "\\\\";
                    break;
                case '\"':
                    out += "\\\"";
                    break;
                case '\n':
                    out += "\\n";
                    break;
                case '\r': /* skip */
                    break;
                default:
                    out += c;
                    break;
                }
            }
            return out;
        }

        static std::string ReadFileToString(const std::string &path)
        {
            std::ifstream f(path, std::ios::binary);
            if (!f)
            {
                throw std::runtime_error("Failed to open file: " + path);
            } 
            std::ostringstream ss;
            ss << f.rdbuf();
            return ss.str();
        }

        static std::string JoinPemArray(const nlohmann::json &arr)
        {
            // arr is expected to be an array of strings; we join with a single '\n'
            std::string out;
            for (size_t i = 0; i < arr.size(); ++i)
            {
                const auto &v = arr.at(i);
                if (!v.is_string())
                {
                    throw std::runtime_error("ca_chain element is not string");
                }
                if (!out.empty() && out.back() != '\n')
                {
                    out.push_back('\n');
                }
                out += v.get<std::string>();
                if (out.back() != '\n')
                {
                    out.push_back('\n');
                }
            }
            return out;
        }

        static bool IsCertExpiring(const std::filesystem::path &certPath, int daysThreshold)
        {
            // we will use raw OpenSSL to read the PEM and check notAfter
            FILE *f = fopen(certPath.string().c_str(), "r");
            if (!f)
            {
                spdlog::warn("IsCertExpiring: cannot open cert {}", certPath.string());
                return false;
            }
            X509 *cert = PEM_read_X509(f, nullptr, nullptr, nullptr);
            fclose(f);
            if (!cert)
            {
                spdlog::warn("IsCertExpiring: cannot parse X509 from {}", certPath.string());
                return false;
            }
            const ASN1_TIME *notAfter = X509_get0_notAfter(cert);
            if (!notAfter)
            {
                spdlog::warn("IsCertExpiring: no notAfter in cert");
                X509_free(cert);
                return false;
            }

            // current time
            time_t now_t = time(nullptr);
            // convert ASN1_TIME → time_t diff
            int days = 0, secs = 0;
            // OpenSSL has ASN1_TIME_diff in newer versions
            if (ASN1_TIME_diff(&days, &secs, nullptr, notAfter) == 0)
            {
                spdlog::warn("IsCertExpiring: ASN1_TIME_diff failed");
                X509_free(cert);
                return false;
            }
            X509_free(cert);

            // days can be negative -> already expired
            spdlog::info("IsCertExpiring: certificate expires in {} days and {} seconds", days, secs);
            if (days < 0)
            {
                return true; // already expired
            }
            if (days == 0 && secs <= 0)
            {
                return true;
            }
            return days <= daysThreshold;
        }

        struct ParsedUrl
        {
            std::string scheme;
            std::string host;
            int port;
        };

        static ParsedUrl ParseBaseUrl(const std::string &base)
        {
            // same simple regex as in DeviceControlService
            std::regex re(R"(^\s*(https?)://([^/:]+)(?::(\d+))?\s*$)");
            std::smatch m;
            if (!std::regex_match(base, m, re))
            {
                throw std::runtime_error("EnrollmentService: invalid base URL: " + base);
            }
            ParsedUrl u;
            u.scheme = m[1].str();
            u.host = m[2].str();
            u.port = m[3].matched ? std::stoi(m[3].str()) : (u.scheme == "https" ? 443 : 80);
            return u;
        }

        // create SSL client with: CA from /etc/iot/keys, device cert from disk,
        // key from *kernel keyring* (decrypted to a temp file)
        static std::unique_ptr<httplib::SSLClient> MakeMtlsClient(const ParsedUrl &u)
        {
            // 1) choose CA
            std::filesystem::path ca = "/etc/iot/keys/issuing_ca.pem";
            if (!std::filesystem::exists(ca))
            {
                ca = "/etc/iot/keys/ca_chain.pem";
            }
            // 2) device cert
            const std::filesystem::path crt = "/etc/iot/keys/device.crt.pem";
            if (!std::filesystem::exists(crt))
            {
                throw std::runtime_error("MakeMtlsClient: device.crt.pem not found");
            }

            // 3) read key payload from kernel keyring and dump to temp (like DeviceControlService)
            auto payload = snoop::device_sec::ReadClientKeyPayloadFromKeyring();
            snoop::device_sec::TempFile tf(std::filesystem::temp_directory_path());
            tf.write_all(payload.data(), payload.size());

            // 4) build SSL client
            auto cli = std::make_unique<httplib::SSLClient>(u.host.c_str(),
                                                            u.port,
                                                            crt.string().c_str(),
                                                            tf.path.string().c_str(),
                                                            std::string() /*no password*/);

            // NOTE: we keep verification relaxed for now, like in your task loop
            cli->enable_server_certificate_verification(false);
            cli->set_ca_cert_path(ca.string().c_str());
            cli->set_connection_timeout(10);
            cli->set_read_timeout(60);
            cli->set_write_timeout(60);

            if (!cli->is_valid())
            {
                throw std::runtime_error("MakeMtlsClient: SSLClient not valid");
            }

            return cli;
        }
    };

} // namespace snoop
