

#include "KeyManager.h"
#include "LRUCache.hpp"
#include "SymAHE.h"
#include "SymMHE.h"
#include "Singleton.h"
#include "elogDefination.h"
#include "t.h"
#include "Aes_Hmac.h"  // 添加AESHMAC头文件

#include <stdlib.h>
#include <tlibc/mbusafecrt.h>
#include <memory>
#include "Enclave_t.h"
#include <sgx_tprotected_fs.h>
#include <sgx_trts.h>
#include <sgx_utils.h>
#include <sgx_tcrypto.h>


#include <math.h>
#include <string>
#include "../../../KeyDistribution/common/blob_format.h"
#include "en_config.h"


std::unique_ptr<SymMHE> psmheCipher(new SymMHE());
std::unique_ptr<SymAHE> psaheCipher(new SymAHE());
std::unique_ptr<encryption> aesEncryption(new encryption());
std::unique_ptr<KeyManager> keyManager(new KeyManager()); //? Singleton 不能用？

constexpr int64_t CacheNValue = INT64_MIN;


int ereport(int errlevel, const char *fmt, ...);

// 简单的日志函数
void Log(const char* fmt, ...) {
    char buf[BUFSIZ] = {'\0'};
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(buf, BUFSIZ, fmt, ap);
    va_end(ap);
    ocall_ereport(buf);  
}

auto 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 sgx_status_t decrypt_dek_blob(const uint8_t* blob,
                                     size_t blob_len,
                                     std::string& dek_out)
{
    if (!blob || blob_len < sizeof(blob_header_v2_t))
        return SGX_ERROR_INVALID_PARAMETER;

    const blob_header_v2_t* h = (const blob_header_v2_t*)blob;
    if (!blob_header_v2_check(h, blob_len))
        return SGX_ERROR_INVALID_PARAMETER;

    size_t off = sizeof(blob_header_v2_t);
    const uint8_t* key_info = blob + off; off += h->keyinfo_len;
    const uint8_t* iv = blob + off;       off += h->iv_len;
    const uint8_t* ct = blob + off;
    size_t ct_len = blob_len - off - h->tag_len;
    const uint8_t* tag = blob + blob_len - h->tag_len;

    if (h->policy != POLICY_MRSIGNER && h->policy != POLICY_MRENCLAVE)
        return SGX_ERROR_INVALID_PARAMETER;
    if (h->isv_prod_id != BLOB_ISV_PROD_ID || h->isv_svn < BLOB_ISV_SVN_MIN)
        return SGX_ERROR_INVALID_PARAMETER;

    sgx_key_request_t req{};
    size_t copy_len = (h->keyinfo_len > sizeof(req)) ? sizeof(req) : h->keyinfo_len;
    memcpy(&req, key_info, copy_len);

    sgx_key_128bit_t hw_key;
    sgx_status_t st = sgx_get_key(&req, &hw_key);
    if (st != SGX_SUCCESS)
        return st;

    std::vector<uint8_t> plain(ct_len);
    st = sgx_rijndael128GCM_decrypt(&hw_key,
                                    ct, (uint32_t)ct_len,
                                    plain.data(),
                                    iv, h->iv_len,
                                    nullptr, 0,
                                    reinterpret_cast<const sgx_aes_gcm_128bit_tag_t*>(tag));
    memset(&hw_key, 0, sizeof(hw_key));
    if (st != SGX_SUCCESS)
        return st;

    dek_out.assign(reinterpret_cast<const char*>(plain.data()), plain.size());
    std::fill(plain.begin(), plain.end(), 0);
    return SGX_SUCCESS;
}



std::string readDEKFromFile() {
    const char* dek_path = DEK_SEALED_PATH; // 定义DEK文件路径
    constexpr size_t max_blob = 2048;

    std::vector<uint8_t> blob(max_blob);
    size_t actual_size = 0;
    sgx_status_t status = ocall_read_dek(dek_path,
                                         reinterpret_cast<char*>(blob.data()),
                                         blob.size(),
                                         &actual_size);
    if (status != SGX_SUCCESS || actual_size == 0) {
        Log("Failed to read DEK blob, status=0x%x size=%zu", status, actual_size);
        return "";
    }

    Log("[DEK] ocall_read_dek ok: path=%s size=%zu", dek_path, actual_size);
    if (actual_size >= sizeof(blob_header_v2_t)) {
        const blob_header_v2_t* hdr = reinterpret_cast<const blob_header_v2_t*>(blob.data());
        if (!blob_header_v2_check(hdr, actual_size)) {
            Log("[DEK] header check fail");
            return "";
        }
        size_t off = sizeof(blob_header_v2_t);
        const uint8_t* iv  = blob.data() + off + hdr->keyinfo_len;
        const uint8_t* tag = blob.data() + actual_size - hdr->tag_len;
        size_t ct_len      = actual_size - off - hdr->keyinfo_len - hdr->iv_len - hdr->tag_len;
        Log("[DEK] header magic=0x%08x ver=%u policy=%u prod_id=%u svn=%u iv_len=%u tag_len=%u keyinfo_len=%u",
            hdr->magic, hdr->ver, hdr->policy, hdr->isv_prod_id, hdr->isv_svn, hdr->iv_len, hdr->tag_len, hdr->keyinfo_len);
        Log("[DEK] parsed iv_len=%u ct_len=%zu tag_len=%u", hdr->iv_len, ct_len, hdr->tag_len);
    } else {
        Log("[DEK] blob too small (%zu)", actual_size);
    }

    std::string dek_plain;
    status = decrypt_dek_blob(blob.data(), actual_size, dek_plain);
    if (status != SGX_SUCCESS) {
        Log("decrypt_dek_blob failed: 0x%x", status);
        return "";
    }

    return dek_plain;
}



int64_t decryptAESLong(char *hexstr) {
    int64_t ans = CacheNValue;
    if (ans == CacheNValue) {

        size_t asz = strlen(hexstr);

        t trans;
        size_t bufASz = asz / 2;

        char bufA[bufASz] = {0};
        char bufOut[aesEncryption->get_crypt_size(bufASz)];
        trans.hex2buf((uint8_t *) bufA, &bufASz, (const char *) hexstr, asz);
        
        // 使用统一的密钥获取函数
        string key = readDEKFromFile();
        
        aesEncryption->aeshmac_decrypt((char *) bufA, bufASz, (char *) key.c_str(), bufOut);
        ans = *(int64_t *) bufOut;
    }
    return ans;
}

double decryptAESDouble(char *hexstr) {
    double ans = 0.0;
    size_t asz = strlen(hexstr);

    t trans;
    size_t bufASz = asz / 2;
    char bufA[bufASz] = {0};
    char bufOut[aesEncryption->get_crypt_size(bufASz)];
    
    trans.hex2buf((uint8_t *) bufA, &bufASz, (const char *) hexstr, asz);
    
    // 使用统一的密钥获取函数
    string key = readDEKFromFile();
    
    size_t decryptedSize = aesEncryption->aeshmac_decrypt((char *) bufA, bufASz, (char *) key.c_str(), bufOut);
    // 解密后的数据强转成double，IEEE754标准
    ans = *(double *) bufOut;
    return ans;
}

void ecall_genMasterKey(const char *path) {
    keyManager->setPath(path, "");
    keyManager->genMasterKey();
}

void ecall_openMasterKey(const char *path, void *out, size_t outlen) {
    keyManager->setPath(path, "");
    keyManager->loadMasterKey();
    KeyManager::masterKey_t m;
    keyManager->masterKey(m);
    memcpy_s((void *) out, outlen, m, outlen);
}

void ecall_encryptKeyWithEnclave(const char *key, size_t keySize, const char *path) {
    keyManager->setPath(path, "");
    keyManager->setMasterKey(key);
    keyManager->saveAllKeysToFile();
}

void ecall_sahe_add(void *a, void *b, void *c, size_t outlen, int isLeftFloat, int isRightFloat) {
    // ! a and b must be buf rather than hex stream

    // TODO: 改成aes
    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num2 = decryptAESLong(static_cast<char *>(b));

    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num2 = decryptAESLong(static_cast<char *>(b));

    // Log("Decrypted values: %ld + %ld", num1, num2);
    // int64_t num3 = num1 + num2;

    // auto tmp = psaheCipher->encrypt(num3);

    // assert(tmp && strlen(tmp) > 5 && tmp[0]);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);
    // delete[] tmp;

    // 判断操作数是浮点数与否，选择相应解密逻辑
    double num1 = 0.0;
    double num2 = 0.0;
    if(isLeftFloat){
        num1 = decryptAESDouble(static_cast<char *>(a));
    }else{
        int64_t num1I = decryptAESLong(static_cast<char *>(a));
        num1 = (double)num1I;
    }
    if(isRightFloat){
        num2 = decryptAESDouble(static_cast<char *>(b));
    }else{
        int64_t num2I = decryptAESLong(static_cast<char *>(b));
        num2 = (double)num2I;
    }
    
    Log("Performing addition: %.15f + %.15f", num1, num2);
    double num3 = num1 + num2;
    Log("Result: %.15f", num3);
    
    // auto tmp = psaheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);

    // 替换原有的psaheCipher加密
    string key = readDEKFromFile();
    char bufIn[sizeof(double)];
    memcpy(bufIn, &num3, sizeof(double));
    
    // aes密文大小 + hmac密文大小
    size_t encryptedSize = aesEncryption->get_crypt_size(sizeof(double)) + 32;
    unsigned char bufOut[encryptedSize];
    
    char* keyBuf = const_cast<char*>(key.c_str());
    size_t outSize = aesEncryption->aeshmac_encrypt(bufIn, sizeof(double), keyBuf, (char*)bufOut);
    
    
    // 手动转换为十六进制字符串
    char hexstr[encryptedSize * 2 + 5];
    memset(hexstr, 0, sizeof(hexstr));
    memcpy(hexstr, "SGX:", 4);
    
    // 使用查找表进行十六进制转换
    static const char hex_chars[] = "0123456789abcdef";
    char* hexptr = hexstr + 4;
    for(size_t i = 0; i < outSize; i++) {
        hexptr[i * 2] = hex_chars[(bufOut[i] >> 4) & 0xF];
        hexptr[i * 2 + 1] = hex_chars[bufOut[i] & 0xF];
    }
    hexptr[outSize * 2] = '\0';  // 确保字符串结束
    
    
    // 验证加密结果
    char* testHexStr = hexstr + 4;
    double decryptedValue = decryptAESDouble(testHexStr);
    Log("Decryption verification - Original: %.15f, Decrypted: %.15f", num3, decryptedValue);
    
    if (fabs(num3 - decryptedValue) > 1e-10) {
        Log("Error: Encryption/Decryption mismatch!");
    }
    
    if (strlen(hexstr) >= outlen) {
        Log("Error: Output buffer too small");
    }
    
    memcpy(static_cast<char *>(c), hexstr, strlen(hexstr) + 1);
}

void ecall_sahe_mul(void *a, int64_t b, void *c, size_t outlen) {
    // ! a and b must be buf rather than hex stream

    // SymMHE *psaheCipher = new SymMHE();
    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // // int64_t num2 = decryptAESLong(static_cast<char *>(b));

    // int64_t num3 = num1 * b;

    // auto tmp = psaheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp), tmp);

    // delete[] tmp;
}

void ecall_sahe_sub(void *a, void *b, void *c, size_t outlen, int isLeftFloat, int isRightFloat) {
    // ! a and b must be buf rather than hex stream
    double num1 = 0.0;
    double num2 = 0.0;
    if(isLeftFloat){
        num1 = decryptAESDouble(static_cast<char *>(a));
    }else{
        int64_t num1I = decryptAESLong(static_cast<char *>(a));
        num1 = (double)num1I;
    }
    if(isRightFloat){
        num2 = decryptAESDouble(static_cast<char *>(b));
    }else{
        int64_t num2I = decryptAESLong(static_cast<char *>(b));
        num2 = (double)num2I;
    }

    Log("Performing subtraction: %.15f - %.15f", num1, num2);
    double num3 = num1 - num2;
    Log("Result: %.15f", num3);

    // 替换原有的psaheCipher加密
    string key = readDEKFromFile();
    char bufIn[sizeof(double)];
    memcpy(bufIn, &num3, sizeof(double));
    
    size_t encryptedSize = aesEncryption->get_crypt_size(sizeof(double)) + 32;
    unsigned char bufOut[encryptedSize];
    
    char* keyBuf = const_cast<char*>(key.c_str());
    size_t outSize = aesEncryption->aeshmac_encrypt(bufIn, sizeof(double), keyBuf, (char*)bufOut);

    // 手动转换为十六进制字符串
    char hexstr[encryptedSize * 2 + 5];
    memset(hexstr, 0, sizeof(hexstr));
    memcpy(hexstr, "SGX:", 4);
    
    // 使用查找表进行十六进制转换
    static const char hex_chars[] = "0123456789abcdef";
    char* hexptr = hexstr + 4;
    for(size_t i = 0; i < outSize; i++) {
        hexptr[i * 2] = hex_chars[(bufOut[i] >> 4) & 0xF];
        hexptr[i * 2 + 1] = hex_chars[bufOut[i] & 0xF];
    }
    hexptr[outSize * 2] = '\0';  // 确保字符串结束
    memcpy(static_cast<char *>(c), hexstr, strlen(hexstr) + 1);
}

void ecall_sahe_neg(void *a, void *c, size_t outlen) {
    // ! a and b must be buf rather than hex stream

    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num3 = -num1;
    // auto tmp = psaheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);
    // delete[] tmp;
}

void ecall_smhe_mul(void *a, void *b, void *c, size_t outlen, int isLeftFloat, int isRightFloat) {
    // ! a and b must be buf rather than hex stream
    double num1 = 0.0;
    double num2 = 0.0;
    if(isLeftFloat){
        num1 = decryptAESDouble(static_cast<char *>(a));
    }else{
        int64_t num1I = decryptAESLong(static_cast<char *>(a));
        num1 = (double)num1I;
    }
    if(isRightFloat){
        num2 = decryptAESDouble(static_cast<char *>(b));
    }else{
        int64_t num2I = decryptAESLong(static_cast<char *>(b));
        num2 = (double)num2I;
    }

    double num3 = num1 * num2;
    Log("Performing multiplication: %.15f * %.15f", num1, num2);
    Log("Result: %.15f", num3);

    string key = readDEKFromFile();
    char bufIn[sizeof(double)];
    memcpy(bufIn, &num3, sizeof(double));
    
    size_t encryptedSize = aesEncryption->get_crypt_size(sizeof(double)) + 32;
    unsigned char bufOut[encryptedSize];
    
    char* keyBuf = const_cast<char*>(key.c_str());
    size_t outSize = aesEncryption->aeshmac_encrypt(bufIn, sizeof(double), keyBuf, (char*)bufOut);
     // 手动转换为十六进制字符串
    char hexstr[encryptedSize * 2 + 5];
    memset(hexstr, 0, sizeof(hexstr));
    memcpy(hexstr, "SGX:", 4);
    
    // 使用查找表进行十六进制转换
    static const char hex_chars[] = "0123456789abcdef";
    char* hexptr = hexstr + 4;
    for(size_t i = 0; i < outSize; i++) {
        hexptr[i * 2] = hex_chars[(bufOut[i] >> 4) & 0xF];
        hexptr[i * 2 + 1] = hex_chars[bufOut[i] & 0xF];
    }
    hexptr[outSize * 2] = '\0';  // 确保字符串结束
    memcpy(static_cast<char *>(c), hexstr, strlen(hexstr) + 1);
}

void ecall_smhe_mulp(void *a, int64_t b, void *c, size_t outlen) {
    // ! a and b must be buf rather than hex stream

    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num3 = num1 * b;
    // auto tmp = psmheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);

    // delete[] tmp;
}

void ecall_smhe_div(void *a, void *b, void *c, size_t outlen, int isLeftFloat, int isRightFloat) {
    // ! a and b must be buf rather than hex stream
    double num1 = 0.0;
    double num2 = 0.0;
    if(isLeftFloat){
        num1 = decryptAESDouble(static_cast<char *>(a));
    }else{
        int64_t num1I = decryptAESLong(static_cast<char *>(a));
        num1 = (double)num1I;
    }
    if(isRightFloat){
        num2 = decryptAESDouble(static_cast<char *>(b));
    }else{
        int64_t num2I = decryptAESLong(static_cast<char *>(b));
        num2 = (double)num2I;
    }

    double num3 = num1 / num2;
    Log("Performing division: %.15f / %.15f", num1, num2);
    Log("Result: %.15f", num3);

    string key = readDEKFromFile();
    char bufIn[sizeof(double)];
    memcpy(bufIn, &num3, sizeof(double));
    
    size_t encryptedSize = aesEncryption->get_crypt_size(sizeof(double)) + 32;
    unsigned char bufOut[encryptedSize];
    
    char* keyBuf = const_cast<char*>(key.c_str());
    size_t outSize = aesEncryption->aeshmac_encrypt(bufIn, sizeof(double), keyBuf, (char*)bufOut);

    // 手动转换为十六进制字符串
    char hexstr[encryptedSize * 2 + 5];
    memset(hexstr, 0, sizeof(hexstr));
    memcpy(hexstr, "SGX:", 4);

    // 使用查找表进行十六进制转换
    static const char hex_chars[] = "0123456789abcdef";
    char* hexptr = hexstr + 4;
    for(size_t i = 0; i < outSize; i++) {
        hexptr[i * 2] = hex_chars[(bufOut[i] >> 4) & 0xF];
        hexptr[i * 2 + 1] = hex_chars[bufOut[i] & 0xF];
    }
    hexptr[outSize * 2] = '\0';  // 确保字符串结束
    memcpy(static_cast<char *>(c), hexstr, strlen(hexstr) + 1); 
}

void ecall_smhe_pow(void *a, int64_t b, void *c, size_t outlen) {
    // ! a and b must be buf rather than hex stream

    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num3 = pow(num1, b);

    // auto tmp = psmheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);
    // delete[] tmp;
}

void ecall_smhe_inv(void *a, void *c, size_t outlen) {
    // ! a and b must be buf rather than hex stream

    // int64_t num1 = decryptAESLong(static_cast<char *>(a));
    // int64_t num3 = 1 / num1;
    // auto tmp = psmheCipher->encrypt(num3);
    // strcpy_s(static_cast<char *>(c), strlen(tmp) + 1, tmp);
    // delete[] tmp;
}

void ecall_ore_compare(void *a, void *b, int *c) {

    int64_t num1 = decryptAESLong(static_cast<char *>(a));
    int64_t num2 = decryptAESLong(static_cast<char *>(b));
    if (num1 == num2)
        *c = 0;
    else if (num1 > num2)
        *c = 1;
    else
        *c = -1;
}

int ereport(int errlevel, const char *fmt, ...) {
    char buf[BUFSIZ] = {'\0'};
    va_list ap;
    va_start(ap, fmt);
    vsnprintf(buf, BUFSIZ, fmt, ap);
    va_end(ap);
    ocall_ereport(buf);
    return (int) strnlen(buf, BUFSIZ - 1) + 1;
}


// template<typename T>
// void swap(T& a, T& b) {
//     T c = a;
//     a = b;
//     b = c;
// }

size_t bsearch(int64_t *arr, size_t size, int64_t x) {
    size_t l = 0, r = size - 1, mid;
    while (l < r) {
        mid = (r - l) / 2 + l;
        if (arr[mid] == x) {
            return mid;
        } else if (arr[mid] > x) {
            r = mid - 1;
            if (mid == 0)
                return -1; // 造成r数据无符号溢出
        } else {
            l = mid + 1;
        }
    }
    return arr[l] == x ? l : -1;
}

void quickSort(int64_t *nums, int64_t l, int64_t r) {
    if (l >= r)
        return;
    int64_t i = l, j = r;
    int64_t x = nums[l];
    while (i < j) {
        while (i < j && nums[l] <= nums[j])
            j--;
        while (i < j && nums[l] >= nums[i])
            i++;
        swap(nums[i], nums[j]);
    }

    swap(nums[i], nums[l]);
    quickSort(nums, l, i - 1);
    quickSort(nums, i + 1, r);
}

size_t ecall_bsearch(int64_t *arr, size_t size, int64_t x) {
    return bsearch(arr, size, x);
}

void ecall_qsort(int64_t *arr, size_t size) {
    quickSort(arr, 0, size - 1);
}


void ecall_init(void *pTaskQueue) {
#ifdef ECALL_BATCH
    static taskpool::TaskPoolLoop taskPoolLoop((taskpool::TaskQueue *)pTaskQueue);
    taskPoolLoop.loop();
#endif
}