#include "server_session.hpp"
#include "tls_server_t.h"

#include <arpa/inet.h>
#include <cstring>
#include <cstdarg>
#include <cstdio>
#include <openenclave/enclave.h>
#include <openenclave/sgx/seal.h>
#include <openenclave/corelibc/stdlib.h>
#include <mbedtls/md.h>
#include <mbedtls/gcm.h> 
#include <openssl/err.h>
#include "../../common/openssl_utility.h"
#include "../../common/blob_format.h"
#include <string>

static ServerSession g_session;
static bool g_seal_plugin_registered = false;

static void LogToHost(const char* fmt, ...)
{
    char buf[256];
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(buf, sizeof(buf), fmt, ap);
    va_end(ap);
    ocall_log(buf);
}

static std::string bytes_to_hex(const uint8_t* data, size_t len)
{
    static const char* h = "0123456789abcdef";
    std::string s;
    s.reserve(len * 2);
    for (size_t i = 0; i < len; ++i)
    {
        s.push_back(h[data[i] >> 4]);
        s.push_back(h[data[i] & 0xF]);
    }
    return s;
}

static std::string hex_head(const uint8_t* p, size_t n) {
    static const char* h = "0123456789abcdef";
    n = (n > 32 ? 32 : n);
    std::string s; s.reserve(n*2);
    for (size_t i=0;i<n;i++){ s.push_back(h[p[i]>>4]); s.push_back(h[p[i]&0xF]); }
    return s;
}


static void TlsInfoCallback(const SSL* ssl, int where, int ret)
{
    LogToHost("[DEK TLS] info callback fired: where=0x%x ret=0x%x", where, ret);
    if (where & SSL_CB_ALERT)
    {
        LogToHost("[DEK TLS] %s alert %s:%s",
                  (where & SSL_CB_WRITE) ? "send" : "recv",
                  SSL_alert_type_string_long(ret),
                  SSL_alert_desc_string_long(ret));
    }
}

// SGX Sealing 
static void secure_memzero(void* ptr, size_t len)
{
    if (!ptr || !len)
        return;
    volatile uint8_t* p = reinterpret_cast<volatile uint8_t*>(ptr);
    while (len--)
        *p++ = 0;
}

static oe_result_t seal_dek(const uint8_t* dek,
                            size_t dek_len,
                            std::vector<uint8_t>& sealed_out)
{
    oe_result_t rc = OE_OK;
    uint8_t iv[BLOB_IV_LEN];
    uint8_t tag[BLOB_TAG_LEN];
    uint8_t* hw_key   = nullptr; size_t hw_key_len   = 0;
    uint8_t* key_info = nullptr; size_t key_info_len = 0;
    std::vector<uint8_t> cipher;

    rc = oe_random(iv, sizeof(iv));
    if (rc != OE_OK)
        return rc;

    rc = oe_get_seal_key_by_policy(OE_SEAL_POLICY_PRODUCT,
                                   &hw_key, &hw_key_len,
                                   &key_info, &key_info_len);
    if (rc != OE_OK)
        goto cleanup;
    if (key_info_len > UINT16_MAX) { rc = OE_OUT_OF_MEMORY; goto cleanup; }

    cipher.resize(dek_len);

    {
        mbedtls_gcm_context gcm;
        mbedtls_gcm_init(&gcm);
        int mrc = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES,
                                     hw_key, static_cast<int>(hw_key_len * 8));
        if (mrc != 0) { rc = OE_FAILURE; mbedtls_gcm_free(&gcm); goto cleanup; }

        mrc = mbedtls_gcm_crypt_and_tag(&gcm, MBEDTLS_GCM_ENCRYPT,
                                        dek_len, iv, sizeof(iv),
                                        nullptr, 0, dek, cipher.data(),
                                        sizeof(tag), tag);
        if (mrc != 0) { rc = OE_FAILURE; mbedtls_gcm_free(&gcm); goto cleanup; }
        mbedtls_gcm_free(&gcm);
    }

    {
        blob_header_v2_t h = {};
        h.magic = BLOB_MAGIC;
        h.ver = BLOB_VER_2;
        h.policy = POLICY_MRSIGNER;
        h.isv_prod_id = BLOB_ISV_PROD_ID;
        h.isv_svn = BLOB_ISV_SVN_MIN;
        h.iv_len = BLOB_IV_LEN;
        h.tag_len = BLOB_TAG_LEN;
        h.keyinfo_len = static_cast<uint16_t>(key_info_len);

        size_t total = sizeof(h) + key_info_len + sizeof(iv) + dek_len + sizeof(tag);
        sealed_out.resize(total);
        uint8_t* p = sealed_out.data();
        memcpy(p, &h, sizeof(h));          p += sizeof(h);
        memcpy(p, key_info, key_info_len); p += key_info_len;
        memcpy(p, iv, sizeof(iv));         p += sizeof(iv);
        memcpy(p, cipher.data(), dek_len); p += dek_len;
        memcpy(p, tag, sizeof(tag));

        LogToHost("[DEK-OE] hw_key=%s", bytes_to_hex(hw_key, hw_key_len).c_str());
        LogToHost("[DEK-OE] sealed size=%zu header: magic=0x%08x ver=%u policy=%u prod_id=%u svn=%u iv_len=%u tag_len=%u keyinfo_len=%u",
                  sealed_out.size(), h.magic, h.ver, h.policy,
                  h.isv_prod_id, h.isv_svn, h.iv_len, h.tag_len, h.keyinfo_len);
    }

cleanup:
    if (hw_key) { secure_memzero(hw_key, hw_key_len); oe_free(hw_key); }
    if (key_info) oe_free(key_info);
    secure_memzero(iv, sizeof(iv));
    return rc;
}




static bool ssl_write_all(SSL* ssl, const uint8_t* data, size_t len)
{
    size_t written = 0;
    while (written < len)
    {
        int rc = SSL_write(ssl,
                           reinterpret_cast<const char*>(data + written),
                           static_cast<int>(len - written));
        if (rc <= 0)
        {
            int err = SSL_get_error(ssl, rc);
            LogToHost("[DEK TLS] SSL_write failed: err=%d", err);
            return false;
        }
        written += static_cast<size_t>(rc);
    }
    return true;
}


static int bio_ocall_write(BIO* bio, const char* buf, int len)
{
    const uint64_t handle = reinterpret_cast<uint64_t>(BIO_get_data(bio));
    size_t written = 0;
    oe_result_t oe_ret = OE_OK;

    if (ocall_session_send(
            &oe_ret,
            handle,
            reinterpret_cast<const uint8_t*>(buf),
            static_cast<size_t>(len),
            &written) != OE_OK ||
        oe_ret != OE_OK)
    {
        return -1;
    }

    return static_cast<int>(written);
}

static int bio_ocall_read(BIO* bio, char* buf, int len)
{
    const uint64_t handle = reinterpret_cast<uint64_t>(BIO_get_data(bio));
    size_t read = 0;
    oe_result_t oe_ret = OE_OK;

    if (ocall_session_recv(
            &oe_ret,
            handle,
            reinterpret_cast<uint8_t*>(buf),
            static_cast<size_t>(len),
            &read) != OE_OK ||
        oe_ret != OE_OK)
    {
        return -1;
    }

    return static_cast<int>(read);
}

static int bio_ocall_create(BIO* bio)
{
    BIO_set_init(bio, 1);
    BIO_set_data(bio, nullptr);
    BIO_set_flags(bio, 0);
    return 1;
}

static int bio_ocall_destroy(BIO* bio)
{
    if (!bio)
        return 0;
    BIO_set_init(bio, 0);
    BIO_set_data(bio, nullptr);
    return 1;
}

static long bio_ocall_ctrl(BIO* bio, int cmd, long num, void* ptr)
{
    switch (cmd)
    {
        case BIO_CTRL_FLUSH:
            return 1;                     // socket无需刷新
        case BIO_CTRL_GET_CLOSE:
            return BIO_get_shutdown(bio);
        case BIO_CTRL_SET_CLOSE:
            BIO_set_shutdown(bio, (int)num);
            return 1;
        case BIO_CTRL_PENDING:
        case BIO_CTRL_WPENDING:
        case BIO_CTRL_DUP:
            return 0;
        default:
            return 0;
    }
}

static void close_session(uint64_t handle)
{
    ocall_session_close(handle);
}

static BIO_METHOD* init_bio_method()
{
    static BIO_METHOD* method = nullptr;
    if (!method)
    {
        BIO_METHOD* base = (BIO_METHOD*)BIO_s_socket();
        method = BIO_meth_new(BIO_get_new_index() |
                              BIO_TYPE_SOURCE_SINK, "oe_socket");
        BIO_meth_set_write(method, bio_ocall_write);
        BIO_meth_set_read(method, bio_ocall_read);
        BIO_meth_set_ctrl(method, bio_ocall_ctrl);
        BIO_meth_set_create(method, bio_ocall_create);
        BIO_meth_set_destroy(method, bio_ocall_destroy);
        BIO_meth_set_puts(method, BIO_meth_get_puts(base));
        BIO_meth_set_gets(method, BIO_meth_get_gets(base));
        BIO_meth_set_callback_ctrl(
            method, BIO_meth_get_callback_ctrl(base));
    }
    return method;
}

extern "C" {

oe_result_t ecall_server_init(server_init_args_t* args)
{
    if (g_session.initialized)
        return OE_OK;

    if (load_oe_modules() != OE_OK)
        return OE_FAILURE;

    // SGX Seal
    if (!g_seal_plugin_registered)  
    {
        oe_result_t reg = oe_register_seal_plugin(
            &oe_seal_plugin_gcm_aes,
            true /* 设为默认插件 */);
        if (reg != OE_OK)
        {
            LogToHost("[DEK TLS] 注册密封插件失败: %s",
                      oe_result_str(reg));
            return OE_FAILURE;
        }
        g_seal_plugin_registered = true;
    }

    OpenSSL_add_ssl_algorithms();
    SSL_load_error_strings();

    SSL_CONF_CTX* conf = SSL_CONF_CTX_new();
    g_session.ssl_ctx = SSL_CTX_new(TLS_server_method());
    if (!conf || !g_session.ssl_ctx)
        return OE_FAILURE;

    if (initalize_ssl_context(conf, g_session.ssl_ctx) != OE_OK)
        return OE_FAILURE;

    X509* cert = nullptr;
    EVP_PKEY* pkey = nullptr;
    if (load_tls_certificates_and_keys(g_session.ssl_ctx, cert, pkey) != OE_OK)
    {
        LogToHost("[DEK TLS] failed to generate attested certificate");
        return OE_FAILURE;
    }
    SSL_CTX_set_info_callback(g_session.ssl_ctx, TlsInfoCallback);
    SSL_CTX_set_verify(g_session.ssl_ctx, SSL_VERIFY_NONE, nullptr);

    g_session.initialized = true;
    return OE_OK;
}


oe_result_t ecall_server_accept()
{
    if (!g_session.initialized)
    {
        LogToHost("[DEK TLS] ecall_server_accept: session not initialized");
        return OE_FAILURE;
    }
    if (g_session.has_session)
    {
        LogToHost("[DEK TLS] ecall_server_accept: busy (session active)");
        return OE_BUSY;
    }

    uint64_t handle = OE_UINT64_MAX;
    oe_result_t oe_ret = OE_OK;

    if (ocall_wait_for_client(&oe_ret, &handle) != OE_OK || oe_ret != OE_OK)
    {
        LogToHost("[DEK TLS] ecall_server_accept: ocall_wait_for_client failed rpc=%s oe=%s",
                  oe_result_str(oe_ret),
                  oe_result_str(oe_ret));
        return OE_FAILURE;
    }
    LogToHost("[DEK TLS] ecall_server_accept: client handle=%lu", handle);

    BIO* bio = BIO_new(init_bio_method());
    BIO_set_data(bio, reinterpret_cast<void*>(handle));
    BIO_set_init(bio, 1);

    g_session.ssl = SSL_new(g_session.ssl_ctx);
    SSL_set_bio(g_session.ssl, bio, bio);

    int ret = SSL_accept(g_session.ssl);
    LogToHost("[DEK TLS] state=%s", SSL_state_string_long(g_session.ssl));

    if (ret <= 0)
    {
        int why = SSL_get_error(g_session.ssl, ret);
        unsigned long err = ERR_get_error();
        LogToHost("[DEK TLS] ecall_server_accept: SSL_accept failed (code=%d)", why);
        while (err != 0)
        {
            const char* reason = ERR_reason_error_string(err);
            if (reason)
                LogToHost("[DEK TLS] SSL error: %s", reason);
            err = ERR_get_error();
        }
        ERR_print_errors_fp(stderr);  // <-- 新增：输出完整 OpenSSL 错误栈
        close_session(handle);
        g_session.Reset();
        return OE_FAILURE;
    }

    LogToHost("[DEK TLS] ecall_server_accept: TLS session established");
    g_session.session_handle = handle;
    g_session.has_session = true;
    return OE_OK;
}



oe_result_t ecall_server_receive_dek(uint32_t* dek_hash)
{
    if (!g_session.has_session)
        return OE_FAILURE;

    uint32_t dek_len_be = 0;
    if (SSL_read(g_session.ssl, &dek_len_be, sizeof(dek_len_be)) !=
        sizeof(dek_len_be))
        return OE_FAILURE;

    uint32_t dek_len = ntohl(dek_len_be);
    if (dek_len == 0 || dek_len > 4096)
        return OE_FAILURE;

    g_session.dek_buffer.resize(dek_len);
    size_t received = 0;
    while (received < dek_len)
    {
        int ret = SSL_read(
            g_session.ssl,
            g_session.dek_buffer.data() + received,
            static_cast<int>(dek_len - received));
        if (ret <= 0)
            return OE_FAILURE;
        received += static_cast<size_t>(ret);
    }

    // g_session.sealed_dek_ready = true;
    // seal
    oe_result_t seal_rc = seal_dek(g_session.dek_buffer.data(),
                               g_session.dek_buffer.size(),
                               g_session.sealed_dek);
    if (seal_rc != OE_OK)
    {
        LogToHost("[DEK TLS] sealing DEK failed: %s", oe_result_str(seal_rc));
        return OE_FAILURE;
    }

    secure_memzero(g_session.dek_buffer.data(), g_session.dek_buffer.size());
    g_session.dek_buffer.clear();
    g_session.dek_ready = false;
    g_session.sealed_dek_ready = true;

    uint8_t digest[32];
    mbedtls_md_context_t md;
    mbedtls_md_init(&md);
    mbedtls_md_setup(
        &md, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 0);
    mbedtls_md_starts(&md);
    mbedtls_md_update(
        &md, g_session.dek_buffer.data(), g_session.dek_buffer.size());
    mbedtls_md_finish(&md, digest);
    mbedtls_md_free(&md);

    if (dek_hash)
        memcpy(dek_hash, digest, sizeof(uint32_t));

    if (!ssl_write_all(g_session.ssl, digest, sizeof(uint32_t)))
        return OE_FAILURE;

    return OE_OK;
}

// SGX Seal
oe_result_t ecall_server_get_sealed_dek(uint8_t* buf,
                                        size_t buf_len,
                                        size_t* out_len)
{
    if (!g_session.sealed_dek_ready || !buf || !out_len)
        return OE_FAILURE;

    if (buf_len < g_session.sealed_dek.size())
        return OE_BUFFER_TOO_SMALL;

    memcpy(buf, g_session.sealed_dek.data(), g_session.sealed_dek.size());
    *out_len = g_session.sealed_dek.size();

    g_session.sealed_dek_ready = false;
    g_session.sealed_dek.clear();
    return OE_OK;
}

// 对非enclave暴露，仅用于测试
oe_result_t ecall_server_get_dek(
    uint8_t* buf, size_t buf_len, size_t* out_len)
{
    if (!g_session.dek_ready || !buf || !out_len)
        return OE_FAILURE;

    if (buf_len < g_session.dek_buffer.size())
        return OE_BUFFER_TOO_SMALL;

    memcpy(buf, g_session.dek_buffer.data(), g_session.dek_buffer.size());
    *out_len = g_session.dek_buffer.size();
    return OE_OK;
}

void ecall_server_shutdown()
{
    if (g_session.has_session)
    {
        SSL_shutdown(g_session.ssl);
        close_session(g_session.session_handle);
    }
    g_session.Shutdown();
    if (g_seal_plugin_registered)
    {
        oe_unregister_seal_plugin(&oe_seal_plugin_gcm_aes.id);
        g_seal_plugin_registered = false;
    }
    oe_verifier_shutdown();
}

oe_result_t ecall_server_reset_session()
{
    if (!g_session.has_session)
        return OE_OK;

    SSL_shutdown(g_session.ssl);
    close_session(g_session.session_handle);
    g_session.Reset();
    return OE_OK;
}

} // extern "C"
