// 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 "ConfigService.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");
            }

            // 2) 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);
            }

            // 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());
                }
            }

            // 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)
                    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;
        }

    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;
        }
    };

} // namespace snoop
