#pragma once
#include <filesystem>
#include <string>
#include <array>
#include <vector>
#include <stdexcept>
#include <cstdio>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <spdlog/spdlog.h>
#include "SslCertUtil.h"

namespace snoop
{
    namespace device_sec
    {
        // --- helpers ---
        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 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());
            }
            int rc = pclose(pipe);
            int exitCode = WIFEXITED(rc) ? WEXITSTATUS(rc) : rc;
            if (exitCode != 0)
            {
                spdlog::warn("Command '{}' exited with code {}", cmd, exitCode);
            }
            return out;
        }

        // dumps the client key from keyring to a temp file and returns its path
        static std::filesystem::path ExtractClientKeyFromKernelKeyring()
        {
            std::string id = Trim(Exec("keyctl search @s user iot-client-key | tail -n1"));
            if (id.empty())
            {
                throw std::runtime_error("iot-client-key not found in keyring");
            }
            // Create a secure temp file
            char tmpl[] = "/run/iot/iot-keyXXXXXX";
            int fd = mkstemp(tmpl);
            if (fd < 0)
            {
                throw std::runtime_error("mkstemp failed for client key");
            }
            close(fd);
            std::filesystem::path p(tmpl);

            // Pipe the key payload into the temp file
            std::string cmd = "keyctl pipe " + id + " > " + p.string();
            Exec(cmd);

            // quick sanity
            if (std::filesystem::file_size(p) == 0)
            {
                std::error_code ec;
                std::filesystem::remove(p, ec);
                throw std::runtime_error("keyctl pipe produced empty client key");
            }
            return p;
        }

        struct TempFile
        {
            std::filesystem::path path;
            int fd{-1};
            explicit TempFile(const std::filesystem::path &dir, const char *pattern = "iot-keyXXXXXX")
            {
                std::string tmpl = (dir / pattern).string();
                std::vector<char> name(tmpl.begin(), tmpl.end());
                name.push_back('\0');
                fd = mkstemp(name.data());
                if (fd < 0)
                {
                    throw std::runtime_error("mkstemp failed");
                }   
                fchmod(fd, S_IRUSR | S_IWUSR);
                path = name.data();
            }
            void write_all(const void *data, size_t n)
            {
                const uint8_t *p = static_cast<const uint8_t *>(data);
                size_t off = 0;
                while (off < n)
                {
                    ssize_t w = ::write(fd, p + off, n - off);
                    if (w <= 0)
                    {
                        throw std::runtime_error("write failed");
                    }
                    off += (size_t)w;
                }
                fsync(fd);
            }
            ~TempFile()
            {
                if (fd >= 0)
                {
                    ::close(fd);
                } 
                std::error_code ec;
                std::filesystem::remove(path, ec);
            }
        };

        inline std::vector<uint8_t> ReadClientKeyPayloadFromKeyring()
        {
            // 1) get key id
            std::string id = Trim(Exec("keyctl search @s user iot-client-key | tail -n1"));
            if (id.empty())
            {
                throw std::runtime_error("iot-client-key not found in keyring");
            }
                
            // 2) capture payload (no redirection to file)
            std::string bytes = Exec("keyctl pipe " + id);
            if (bytes.empty())
            {
                throw std::runtime_error("keyctl pipe returned empty payload");
            } 

            return std::vector<uint8_t>(bytes.begin(), bytes.end());
        }

        static std::string GetCpuSerial()
        {
            std::ifstream f("/proc/cpuinfo");
            if (!f)
            {
                spdlog::warn("GetCpuSerial: /proc/cpuinfo not available, using fallback");
                return std::string(12, '9');
            }

            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("GetCpuSerial: Serial not found, using fallback");
            return std::string(12, '9');
        }

        inline bool IsClientKeyInKernelKeyring()
        {
            std::string id = Trim(Exec("keyctl search @s user iot-client-key | tail -n1"));
            return !id.empty();
        }
        
        // Ensure that iot-client-key exists in kernel keyring.
        // - If already present: no-op, returns false (nothing loaded).
        // - If not present: decrypts /etc/iot/keys/device.key.enc using KEK
        //   derived from (GUID, CPU_SERIAL), padds into keyring, returns true.
        inline bool EnsureClientKeyInKernelKeyring(const std::string &guid)
        {
            if (IsClientKeyInKernelKeyring())
            {
                spdlog::info("EnsureClientKeyInKernelKeyring: key already present, nothing to do");
                return false;
            }

            const std::filesystem::path encPath = "/etc/iot/keys/device.key.enc";
            if (!std::filesystem::exists(encPath))
            {
                throw std::runtime_error("EnsureClientKeyInKernelKeyring: encrypted key not found at " +
                                         encPath.string());
            }

            // Derive KEK = HMAC-SHA256(cpuSerial, key=GUID)
            const std::string cpuSerial = GetCpuSerial();
            spdlog::info("EnsureClientKeyInKernelKeyring: CPU_SERIAL = {}", cpuSerial);

            const std::string kek =
                snoop::device_sec::SslCertUtil::ComputeHmacSha256Hex(guid, cpuSerial);

            if (kek.empty())
            {
                throw std::runtime_error("EnsureClientKeyInKernelKeyring: empty KEK");
            }
            spdlog::debug("EnsureClientKeyInKernelKeyring: KEK (hex) = {}", kek);

            // Decrypt into tmpfs: /run/iot/device.key
            const std::filesystem::path runDir = "/run/iot";
            const std::filesystem::path plainPath = runDir / "device.key";
            std::filesystem::create_directories(runDir);

            snoop::device_sec::SslCertUtil::DecryptFileAes256CbcPbkdf2(
                encPath,
                plainPath,
                kek);

            if (!std::filesystem::exists(plainPath))
            {
                throw std::runtime_error("EnsureClientKeyInKernelKeyring: decrypted key not created");
            }

            // Load into kernel keyring
            const std::string cmd = "keyctl padd user iot-client-key @s < " + plainPath.string();
            spdlog::info("EnsureClientKeyInKernelKeyring: loading key into kernel keyring...");
            auto out = Exec(cmd);
            spdlog::debug("keyctl padd output:\n{}", out);

            // Securely erase plaintext
            const std::string shredCmd = "shred -u " + plainPath.string();
            spdlog::info("EnsureClientKeyInKernelKeyring: shredding plaintext key with '{}'", shredCmd);
            auto shredOut = Exec(shredCmd);
            spdlog::debug("shred output:\n{}", shredOut);

            return true;
        }
    }
}