#include "receiver.hpp"
#include "host.hpp"
#include "server_init_config.hpp"

#include <fstream>
#include <iterator>
#include <memory>
#include <vector>
#include "en_config.h"

extern "C" {
#include "postgres.h"
#include "miscadmin.h"
#include "utils/elog.h"
#include "lib/stringinfo.h"
}

namespace {

std::unique_ptr<ServerHost> g_tlsServer;

std::vector<uint8_t> LoadFile(const std::string& path)
{
    std::ifstream ifs(path, std::ios::binary);
    if (!ifs)
        elog(FATAL, "Cannot open %s", path.c_str());

    return std::vector<uint8_t>(std::istreambuf_iterator<char>(ifs), {});
}


void HandleSealedDek(const std::vector<uint8_t>& sealed)
{
    elog(LOG, "Received sealed DEK blob (size=%zu)", sealed.size());
    if (sealed.empty())
        return;

    const char* path = DEK_SEALED_PATH; // 定义DEK文件路径
    std::ofstream ofs(path,
                      std::ios::binary | std::ios::trunc);
    if (!ofs)
    {
        elog(WARNING, "Failed to persist sealed DEK to %s", path);
        return;
    }
    ofs.write(reinterpret_cast<const char*>(sealed.data()),
              sealed.size());
    if (!ofs)
        elog(WARNING, "Short write while persisting sealed DEK");
}

} // namespace

bool StartDekReceiverService(const std::string& listen_host,
                             const std::string& listen_port,
                             const std::string& enclave_signed_path)
{
    if (g_tlsServer)
        return true;

    ServerInitConfig cfg;
    cfg.allow_debug = 0;

    g_tlsServer = std::make_unique<ServerHost>();
    if (!g_tlsServer->Start(enclave_signed_path.c_str(),
                            listen_port.c_str(),
                            cfg))
    {
        elog(WARNING, "Start TLS server host failed");
        g_tlsServer.reset();
        return false;
    }

    elog(LOG, "DEK receiver listening on %s:%s",
         listen_host.c_str(), listen_port.c_str());
    return true;
}

bool ReceiveDekAndHandleOnce()
{
    if (!g_tlsServer)
    {
        elog(WARNING, "DEK receiver not started");
        return false;
    }

    std::vector<uint8_t> sealed_dek;
    if (!g_tlsServer->AcceptAndReceiveDek(sealed_dek))
    {
        elog(LOG, "AcceptAndReceiveDek failed");
        return false;
    }
    HandleSealedDek(sealed_dek);
    return true;
}

void StopDekReceiverService()
{
    if (!g_tlsServer)
        return;

    g_tlsServer->Stop();
    g_tlsServer.reset();
    elog(LOG, "DEK receiver stopped");
}

extern "C" bool KD_StartDekReceiverService(const char* listen_host,
                                           const char* listen_port,
                                           const char* enclave_signed_path)
{
    const std::string host  = listen_host    ? listen_host    : "";
    const std::string port  = listen_port    ? listen_port    : "";
    const std::string enclave = enclave_signed_path ? enclave_signed_path : "";

    return StartDekReceiverService(host, port, enclave);
}

extern "C" bool KD_ReceiveDekAndHandleOnce(void)
{
    return ReceiveDekAndHandleOnce();
}

extern "C" void KD_StopDekReceiverService(void)
{
    StopDekReceiverService();
}
