#include <openenclave/host.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>
#include "tls_server_u.h"
#include "host.hpp"
#include "server_init_config.hpp"
#include <array>
#include <vector>
#include <errno.h>
#include <string.h>

extern "C" {
#include "postgres.h"
#include "miscadmin.h"
#include "postmaster/interrupt.h"
#include "utils/elog.h"
}


namespace {

/* 监听上下文 */
int g_listen_fd = -1;
constexpr int kSocketWaitTimeoutMs = 500;

static inline const char* errno_str(int err)
{
    return err ? strerror(err) : "success";
}

static bool shutdown_requested(void)
{
    return ShutdownRequestPending || ProcDiePending;
}

static bool wait_for_socket_event(int fd, short events, const char* action)
{
    pollfd pfd{};
    pfd.fd = fd;
    pfd.events = events;

    for (;;)
    {
        int rc = poll(&pfd, 1, kSocketWaitTimeoutMs);
        if (rc > 0)
            return true;

        if (rc == 0)
        {
            if (shutdown_requested())
                return false;
            continue;
        }

        if (errno == EINTR)
        {
            if (shutdown_requested())
                return false;
            continue;
        }

        int err = errno;
        elog(WARNING,
             "%s poll(fd=%d, events=%hd) failed: errno=%d (%s)",
             action,
             fd,
             events,
             err,
             errno_str(err));
        return false;
    }
}

bool init_listener(const char* port)
{
    g_listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (g_listen_fd < 0)
        return false;

    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(static_cast<uint16_t>(atoi(port)));

    int opt = 1;
    setsockopt(g_listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    return bind(g_listen_fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) ==
               0 &&
           listen(g_listen_fd, 1) == 0;
}

} // namespace

/* OCALL 实现 */
oe_result_t ocall_wait_for_client(uint64_t* session_handle)
{
    sockaddr_in client{};
    socklen_t len = sizeof(client);

    for (;;)
    {
        if (!wait_for_socket_event(g_listen_fd, POLLIN, "ocall_wait_for_client"))
            return OE_FAILURE;

        int fd = accept(g_listen_fd, reinterpret_cast<sockaddr*>(&client), &len);
        if (fd < 0)
        {
            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
            {
                if (shutdown_requested())
                    return OE_FAILURE;
                continue;
            }

            int err = errno;
            elog(WARNING,
                 "accept(fd=%d) failed: errno=%d (%s)",
                 g_listen_fd,
                 err,
                 errno_str(err));
            return OE_FAILURE;
        }

        *session_handle = static_cast<uint64_t>(fd);
        return OE_OK;
    }
}

oe_result_t ocall_session_send(uint64_t session_handle,
                   const uint8_t* buf,
                   size_t len,
                   size_t* bytes_written)
{
    const int fd = static_cast<int>(session_handle);
    size_t total = 0;

    elog(LOG,
         "ocall_session_send(fd=%d, len=%zu) enter",
         fd,
         len);

    while (total < len)
    {
        if (!wait_for_socket_event(fd, POLLOUT, "ocall_session_send"))
            return OE_FAILURE;

        ssize_t sent = send(fd,
                            buf + total,
                            len - total,
#ifdef MSG_NOSIGNAL
                            MSG_NOSIGNAL);
#else
                            0);
#endif
        if (sent < 0)
        {
            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
                continue;
            int err = errno;

            elog(WARNING,
                 "ocall_session_send(fd=%d) failed: errno=%d (%s) "
                 "already_sent=%zu of %zu",
                 fd,
                 err,
                 errno_str(err),
                 total,
                 len);
            return OE_FAILURE;
        }

        if (sent == 0)
        {
            elog(WARNING,
                 "ocall_session_send(fd=%d) peer closed while sending "
                 "(already_sent=%zu of %zu)",
                 fd,
                 total,
                 len);
            break;
        }

        total += static_cast<size_t>(sent);
    }

    *bytes_written = total;
    elog(LOG,
         "ocall_session_send(fd=%d) done, wrote=%zu",
         fd,
         total);
    return OE_OK;
}

oe_result_t ocall_session_recv(uint64_t session_handle,
                   uint8_t* buf,
                   size_t len,
                   size_t* bytes_read)
{
    const int fd = static_cast<int>(session_handle);

    elog(LOG,
         "ocall_session_recv(fd=%d, len=%zu) enter",
         fd,
         len);

    for (;;)
    {
        if (!wait_for_socket_event(fd, POLLIN, "ocall_session_recv"))
            return OE_FAILURE;

        ssize_t recvd = recv(fd, buf, len, 0);
        if (recvd < 0)
        {
            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
                continue;
            int err = errno;

            elog(WARNING,
                 "ocall_session_recv(fd=%d) failed: errno=%d (%s)",
                 fd,
                 err,
                 errno_str(err));
            return OE_FAILURE;
        }

        if (recvd == 0)
        {
            elog(WARNING,
                 "ocall_session_recv(fd=%d) peer closed socket",
                 fd);
            return OE_FAILURE;
        }

        *bytes_read = static_cast<size_t>(recvd);
        elog(LOG,
             "ocall_session_recv(fd=%d) got=%zu bytes",
             fd,
             *bytes_read);
        return OE_OK;
    }
}

void ocall_session_close(uint64_t session_handle)
{
    close(static_cast<int>(session_handle));
}

void ocall_log(const char* msg)
{
    fprintf(stderr, "[ServerHost] %s\n", msg ? msg : "(null)");
}

/* ServerHost成员函数 */
ServerHost::~ServerHost()
{
    Stop();
}

bool ServerHost::Start(const char* enclave_path,
            const char* listen_port,
            const ServerInitConfig& cfg)
{
    if (!init_listener(listen_port))
    {
        elog(WARNING, "init_listener(%s) failed: %m", listen_port);
        return false;
    }
        

    oe_enclave_t* enc = nullptr;
    if (oe_create_tls_server_enclave(enclave_path,
                                        OE_ENCLAVE_TYPE_SGX,
                                        OE_ENCLAVE_FLAG_DEBUG,
                                        nullptr,
                                        0,
                                        &enc) != OE_OK)
    {
        elog(WARNING, "oe_create_tls_server_enclave failed");
        return false;
    }

    enclave_ = enc;

    server_init_args_t args{};
    args.allow_debug = cfg.allow_debug;

    oe_result_t ret = ecall_server_init(enclave_, &result_, &args);
    if (ret != OE_OK || result_ != OE_OK)
    {
        elog(WARNING,
             "ecall_server_init failed: ret=%s, oe=%s",
             oe_result_str(ret),
             oe_result_str(result_));
        return false;
    }
    return true;
}

// bool ServerHost::AcceptAndReceiveDek(std::vector<uint8_t>& out_dek)
// {
//     oe_result_t rpc_ret = OE_OK;
//     bool session_active = false;
//     bool success = false;

//     auto cleanup_session = [&](void) {
//         if (!session_active)
//             return;

//         oe_result_t cleanup_oe = OE_OK;
//         oe_result_t cleanup_rpc =
//             ecall_server_reset_session(enclave_, &cleanup_oe);
//         if (cleanup_rpc != OE_OK || cleanup_oe != OE_OK)
//         {
//             elog(WARNING,
//                  "AcceptAndReceiveDek: ecall_server_reset_session failed rpc=%s oe=%s",
//                  oe_result_str(cleanup_rpc),
//                  oe_result_str(cleanup_oe));
//         }
//         session_active = false;
//     };

//     rpc_ret = ecall_server_accept(enclave_, &result_);
//     if (rpc_ret != OE_OK || result_ != OE_OK)
//     {
//         elog(LOG,
//              "AcceptAndReceiveDek: ecall_server_accept failed rpc=%s oe=%s",
//              oe_result_str(rpc_ret),
//              oe_result_str(result_));
//         cleanup_session();
//         return false;
//     }
//     elog(LOG, "AcceptAndReceiveDek: client accepted");
//     session_active = true;

//     uint32_t tag = 0;
//     rpc_ret = ecall_server_receive_dek(enclave_, &result_, &tag);
//     if (rpc_ret != OE_OK || result_ != OE_OK)
//     {
//         elog(LOG,
//              "AcceptAndReceiveDek: ecall_server_receive_dek failed rpc=%s oe=%s",
//              oe_result_str(rpc_ret),
//              oe_result_str(result_));
//         cleanup_session();
//         return false;
//     }
//     elog(LOG, "AcceptAndReceiveDek: DEK received (tag=%u)", tag);

//     size_t len = 0;
//     std::array<uint8_t, 1024> buf{};

//     rpc_ret = ecall_server_get_dek(
//         enclave_, &result_, buf.data(), buf.size(), &len);
//     if (rpc_ret != OE_OK || result_ != OE_OK)
//     {
//         elog(LOG,
//              "AcceptAndReceiveDek: ecall_server_get_dek failed rpc=%s oe=%s",
//              oe_result_str(rpc_ret),
//              oe_result_str(result_));
//         cleanup_session();
//         return false;
//     }
//     elog(LOG, "AcceptAndReceiveDek: DEK length=%zu", len);

//     out_dek.assign(buf.begin(), buf.begin() + len);
//     success = true;
//     cleanup_session();
//     return success;
// }
bool ServerHost::AcceptAndReceiveDek(std::vector<uint8_t>& out_sealed_dek)
{
    oe_result_t rpc_ret = OE_OK;
    bool session_active = false;
    bool success = false;

    auto cleanup_session = [&]() {
        if (!session_active)
            return;

        oe_result_t cleanup_oe = OE_OK;
        oe_result_t cleanup_rpc =
            ecall_server_reset_session(enclave_, &cleanup_oe);
        if (cleanup_rpc != OE_OK || cleanup_oe != OE_OK)
        {
            elog(WARNING,
                 "AcceptAndReceiveDek: ecall_server_reset_session failed rpc=%s oe=%s",
                 oe_result_str(cleanup_rpc),
                 oe_result_str(cleanup_oe));
        }
        session_active = false;
    };

    rpc_ret = ecall_server_accept(enclave_, &result_);
    if (rpc_ret != OE_OK || result_ != OE_OK)
    {
        elog(LOG,
             "AcceptAndReceiveDek: ecall_server_accept failed rpc=%s oe=%s",
             oe_result_str(rpc_ret),
             oe_result_str(result_));
        cleanup_session();
        return false;
    }
    elog(LOG, "AcceptAndReceiveDek: client accepted");
    session_active = true;

    uint32_t tag = 0;
    rpc_ret = ecall_server_receive_dek(enclave_, &result_, &tag);
    if (rpc_ret != OE_OK || result_ != OE_OK)
    {
        elog(LOG,
             "AcceptAndReceiveDek: ecall_server_receive_dek failed rpc=%s oe=%s",
             oe_result_str(rpc_ret),
             oe_result_str(result_));
        cleanup_session();
        return false;
    }
    elog(LOG, "AcceptAndReceiveDek: sealed DEK produced (tag=%u)", tag);

    size_t sealed_len = 0;
    std::array<uint8_t, 2048> sealed_buf{};
    rpc_ret = ecall_server_get_sealed_dek(
        enclave_, &result_, sealed_buf.data(), sealed_buf.size(), &sealed_len);
    if (rpc_ret != OE_OK || result_ != OE_OK)
    {
        elog(LOG,
             "AcceptAndReceiveDek: ecall_server_get_sealed_dek failed rpc=%s oe=%s",
             oe_result_str(rpc_ret),
             oe_result_str(result_));
        cleanup_session();
        return false;
    }
    elog(LOG, "AcceptAndReceiveDek: sealed blob size=%zu", sealed_len);

    out_sealed_dek.assign(sealed_buf.begin(),
                          sealed_buf.begin() + sealed_len);
    success = true;
    cleanup_session();
    return success;
}


void ServerHost::Stop()
{
    if (enclave_)
    {
        ecall_server_shutdown(enclave_);
        oe_terminate_enclave(enclave_);
    }
    enclave_ = nullptr;
    if (g_listen_fd >= 0)
        close(g_listen_fd);
    g_listen_fd = -1;
}
