#include "encryptsql.h"
#include <iostream>
#include <regex>
#include <string>
#include <pthread.h>

#ifdef __cplusplus
extern "C"
{
#endif
#include "postgres.h" // 与 #include <iomanip> 冲突

#include "utils/palloc.h"
#include "utils/memutils.h"

#include "nodes/nodes.h"
#include "encryptsql/node2str.h"

#include "utils/memutils.h"

#ifdef __cplusplus
}
#endif

#include <unistd.h>
#include <cstring>
#include <cstdlib>

#include "utils.h"
#include "encryptsql/fieldmap.h"

#include "Singleton.h"
#include "LRUCache.hpp"
#include "KeyManager.h"
#include "SymAHE.h"
#include "TimeCounter.h"


// #include "kms/KeyManagementService.hpp"
// #include "kms/PGconnectForDEK.h"

#include "en_config.h"
#include "kms/kms_interface.hpp"
#include "kms/kms_core_local.hpp"
#include "kmsAdapter/dek_interface.hpp"

// #ifndef MAPPERPATH
// #define MAPPERPATH "/etc/encryptsql/map.json"
// #endif
// #define CMKFILEPATH "/new_enc/kms_test/output.json"
// #define CMKKEYPATH "/new_enc/kms_test/domainkey"
// // 持久化文件路径
// #define CMK_AUTO_ROTATE_STATUS_PATH "/new_enc/kms_test/CMK_auto_rotate_status.json"

extern void LogInfo(const char *fmt, ...);

constexpr int MEMORY_DELETE_LIMIT = 20000;

__thread size_t sql_processed_counter = 0;
#ifdef USE_LRU
// extern LRUCache<long, long> HECache;           // for RandNum(id)  
extern LRUCache<long, std::string> AESCache;   // for AES decrypt
#endif


extern "C" List *raw_parser(const char *str);

extern Node *encryptStmt(Node *parseTree, EncryptInfo *info, char **err_msg_internal);

extern const char *parseTree2str(Node *parseTree);

extern "C" const char *encryptOneSql(const char *sql, char **err_msg, const char *user_name,
                                     const char *db_name) __attribute__((visibility("default")));

extern "C" void disableAsNeeded() __attribute__((visibility("default")));

extern void recovery_backup(bool is_transaction);

__thread MemoryContext EncryptContext = NULL;

static __thread bool Init = false;
__thread bool Encrypt = false;
__thread bool SkipDecryptRes = false;
thread_local std::string dek_for_decrypt = "";


void disableAsNeeded() {
    if (!Init)
        std::cout << "Encryption Library Linked\n";
    return;
}


bool needDecryptResult(Node *node) {

    if (node->type == T_ExplainStmt) {
        return false;
    }
    return true;
}

bool needSkipEncryption(const char *sql, Node *node = nullptr) {
    if (sql) {
        if (isFilteredSql(sql)) {
            return true;
        }
    }
    if (node) {
        if (node->type == T_VariableShowStmt) {
            return true;
        }
    }

    return false;
}

bool isCmkCreateCmd(const std::string &input) {
    std::string cmd = input;
    std::transform(cmd.begin(), cmd.end(), cmd.begin(), ::tolower);
    if (cmd.substr(0, 10) == "create cmk") {
        return true;
    }
    return false;
}

bool isCorrectCmkCreateCmd(const std::string &input, int &rotate_period) {
    std::regex cmd_regex(R"(create cmk -rotate_period=(\d+);)", std::regex::icase);
    std::smatch match;
    if (std::regex_match(input, match, cmd_regex)) {
        std::string str_rotate_period = match[1].str();
        rotate_period = std::stoi(str_rotate_period);
        return true;
    }
    rotate_period = -1;
    return false;
}

bool isCmkDeleteCmd(const std::string &input) {
    std::string cmd = input;
    std::transform(cmd.begin(), cmd.end(), cmd.begin(), ::tolower);
    if (cmd.substr(0, 10) == "delete cmk") {
        return true;
    }
    return false;
}

bool isRoleDropCmd(const std::string &input) {
    std::string cmd = input;
    std::transform(cmd.begin(), cmd.end(), cmd.begin(), ::tolower);
    if (cmd.substr(0, 9) == "drop role" || cmd.substr(0, 9) == "drop user") {
        return true;
    }
    return false;
}

//检查手动轮转命令
bool isManualRotateCmd(const std::string &input) {
    std::string cmd = input;
    std::transform(cmd.begin(), cmd.end(), cmd.begin(), ::tolower);
    if (cmd.substr(0, 14) == "rotate cmk now") {
        return true;
    }
    return false;
}

bool isDescribeCmkCmd(const std::string &input) {
    std::string cmd = input;
    std::transform(cmd.begin(), cmd.end(), cmd.begin(), ::tolower);
    return cmd.substr(0, 12) == "describe cmk";
}

bool parseDescribeCmkCmd(const std::string& input, bool& decrypt) {
    std::regex cmd_regex(R"(describe cmk(?:\s+(-d|-decrypt))?\s*;?\s*$)", std::regex::icase);
    std::smatch match;
    if (std::regex_match(input, match, cmd_regex)) {
        decrypt = match[1].matched;   // 只要匹配到-d或者-decrypt存在就说明要解密，返回明文
        return true;
    }
    return false;
}

const char *
encryptOneSql(const char *sql, char **err_msg, const char *user_name, const char *db_name) {
    TimeCounter *counter = Singleton<TimeCounter>::instance();
    Timer globalTimer;
#ifndef USE_ENCRYPT
    Encrypt = false;
    return sql;
#endif
    // std::cout << "main start\n";
    SkipDecryptRes = false;
    if (checkCmd(sql, &Encrypt)) {
        SkipDecryptRes = true;
        return Encrypt ? "select 1 as enc_status;" : "select 0 as enc_status;"; // 这是设置transform的语句
    }
    if (!Encrypt) {
        return sql;
    } else {
        recovery_backup(false);
    }

    Node *parseTree = NULL;
    Node *encParseTree = NULL;
    const char *res = NULL;

    long long tid = gettid();
    __pid_t pid = getpid();
#ifndef NDEBUG
    std::cout << "PID:" << pid << std::endl;
    debug2file(sql, true);
    std::cout << sql << std::endl;
#endif

    if (!Init) {
        MemoryContextInit();
        MessageContext = AllocSetContextCreate(TopMemoryContext,
                                               "MessageContext",
                                               ALLOCSET_DEFAULT_MINSIZE,
                                               ALLOCSET_DEFAULT_INITSIZE,
                                               ALLOCSET_DEFAULT_MAXSIZE);
        EncryptContext = AllocSetContextCreate(TopMemoryContext,
                                               "EncryptContext",
                                               ALLOCSET_DEFAULT_MINSIZE,
                                               ALLOCSET_DEFAULT_INITSIZE,
                                               ALLOCSET_DEFAULT_MAXSIZE);
        //初始化映射表
        mapperInit(MAP_JSON_PATH);

        //注册kms
        LocalKmsCore::registerLocalKms();
        LocalKmsConfig config(CMK_FILE_PATH, CMK_KEY_PATH, CMK_ROTATE_STATUS_PATH, user_name, db_name);
        KMSInterface::initialize(config);

#if 0
        HuaweiKmsCore::registerHuaweiKms();
        HuaweiKmsConfig hwConfig(
            "HPUAI1XPW4U0CPJCSE7C",      // AK
            "qSOZjbqKxffaqFm3FG7KrqV4MtAKYgsAVACB0ZLs",  // SK
            "d8aadaf24bde4314bf2789fb87f4ad83",  // projectId
            "cn-north-4",                 // region
            "",                           // endpoint (可选)
            "my-app-key"                  // keyAlias (可选)
        );
        KMSInterface::initialize(hwConfig);
#endif 

        //初始化DEK
        DekInterface::initialize(user_name, db_name);

#ifdef USE_LRU

        // HECache.setName("Client HE");           // for RandNum(id)
        AESCache.setName("Client AES");
#endif
        Init = true;
    }
    sql_processed_counter++;

    if (sql_processed_counter < MEMORY_DELETE_LIMIT) {
        MemoryContextResetChildren(TopMemoryContext);
        // MemoryContextReset(MessageContext);
    } else {
        sql_processed_counter = 0;
        MemoryContextDeleteChildren(TopMemoryContext);
        // MemoryContextDelete(MessageContext);
    }

    if (needSkipEncryption(sql)) {
        SkipDecryptRes = true;
        return sql;
    }

    MemoryContextSwitchTo(MessageContext);

    std::string user_in(sql);
    if (isCmkCreateCmd(user_in)) {
        SkipDecryptRes = true;
        int rotate_period = -1;
        if (isCorrectCmkCreateCmd(user_in, rotate_period)) {
            // if (cmk_mapperHasCmk(user_name)) {
            if (KMSInterface::hasCmk(user_name)){
                const char *err_msg_tmp = "The user already has a CMK.";
                *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
                strcpy(*err_msg, err_msg_tmp);
                return "";
            }
            CMK cmk;
            // cmk_mapperCreateCmk(user_name, rotate_period, cmk);
            // cmk_mapperStoreCmk(cmk);
            KMSInterface::createCmk(user_name, rotate_period);
            std::string ans = "select " + to_string(cmk._key_id) + " as CMK_ID;";
            char *ans_c = (char *) palloc(ans.size() + 5);
            sprintf(ans_c, "%s", ans.c_str());
            return ans_c;
        } else {
            const char *err_msg_tmp = "Wrong CMK create command.";
            *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
            strcpy(*err_msg, err_msg_tmp);
            return "";
        }
    }   

    if (isCmkDeleteCmd(user_in)) {
        SkipDecryptRes = true;
        // if (cmk_mapperDeleteCmk(user_name)) {
        if (KMSInterface::deleteCmk(user_name)) {
            std::string ans = "select 'CMK deleted successfully' as CMK_delete_status;";
            char *ans_c = (char *) palloc(ans.size() + 5);
            sprintf(ans_c, "%s", ans.c_str());
            return ans_c;
        } else {
            const char *err_msg_tmp = "Failed to delete CMK.";
            *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
            strcpy(*err_msg, err_msg_tmp);
            return "";
        }
    }

    if (isRoleDropCmd(user_in)) {
        SkipDecryptRes = true;
        // if (cmk_mapperHasCmk(user_name)) {
        if(KMSInterface::hasCmk(user_name)){
            const char *err_msg_tmp = "The user still has a CMK. Please delete it first.";
                *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
                strcpy(*err_msg, err_msg_tmp);
                return "";
        }
    }

    //是否有cmk
    // if (!cmk_mapperHasCmk(user_name)) {
    if(! KMSInterface::hasCmk(user_name)){
        const char *err_msg_tmp = "The user do not have CMK now.";
        *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
        strcpy(*err_msg, err_msg_tmp);
        return "";
    }

    
    if (isDescribeCmkCmd(user_in)) {
        SkipDecryptRes = true;   

        bool decrypt = false;
        if (!parseDescribeCmkCmd(user_in, decrypt)) {
            const char *err_msg_tmp = "Invalid describe cmk syntax.";
            *err_msg = (char *) palloc(strlen(err_msg_tmp) + 1);
            strcpy(*err_msg, err_msg_tmp);
            return "";
        }

        std::string ans;
        if (!KMSInterface::describeCmk(user_name, ans, decrypt)) {
            const char *err_msg_tmp = "No CMK found for user.";
            *err_msg = (char *) palloc(strlen(err_msg_tmp) + 1);
            strcpy(*err_msg, err_msg_tmp);
            return "";
        }


        char *ans_c = (char *) palloc(ans.size() + 1);
        strcpy(ans_c, ans.c_str());
        return ans_c;
    }

    //手动轮转部分
    if (isManualRotateCmd(user_in)) {
        SkipDecryptRes = true;
        // if (cmk_mapperRotateCmkNow(user_name)) {
        if (KMSInterface::rotateCmkNow(user_name)) {
            std::string ans = "select 'Manually rotate successfully' as CMK_rotation_status;";
            char *ans_c = (char *) palloc(ans.size() + 5);
            sprintf(ans_c, "%s", ans.c_str());
            return ans_c;
        }
        else {
            const char *err_msg_tmp = "Failed to rotate CMK manually.";
            *err_msg = (char *) malloc(strlen(err_msg_tmp) + 1);
            strcpy(*err_msg, err_msg_tmp);
            return "";
        }
    }

    //自动轮转部分
    std::string auto_rotate_action = "";
    if (isAutoRotateCmd(sql, &auto_rotate_action)) {
        SkipDecryptRes = true;
        // if (cmk_mapperHandleAutoRotateCmd(&auto_rotate_action, user_name)) {
        if(KMSInterface::handleAutoRotateCmd(&auto_rotate_action, user_name)) {
            // return cmk_mapperGetCmkAutoRotateStatusByUsername(user_name) ? "select 'ON' as CMK_rotate_status;" : "select 'OFF' as CMK_rotate_status;";
            return KMSInterface::getCmkAutoRotateStatusByUsername(user_name) ? "select 'ON' as CMK_rotate_status;" : "select 'OFF' as CMK_rotate_status;";
        }  
    }

    // RotateCommandResult result = parseRotateCommand(user_in,user_name,db_name);
    DekInterface::RotateCommandResult result = DekInterface::parseRotateCommand(user_in,user_name,db_name);
    
    // if(result.valid){
    //     connectionUpdateDek(result);
    //     cmk_mapperClearDekForUpdate();
    //     return "";
    // }
    if(result.valid){
        DekInterface::connectionUpdateDek(result);
        DekInterface::clearDekForUpdate();
        return "";
    }

    Timer t;
    List *AST = (List *) raw_parser(sql);

    if (AST)
        parseTree = (Node *) linitial(AST);
    else {
        SkipDecryptRes = true;
        errmsg("error when parse SQL:\n %s\n", sql);
    }
    if (parseTree) {
        MemoryContextSwitchTo(EncryptContext);
        Timer t;
        if (needSkipEncryption(nullptr, parseTree)) {
            SkipDecryptRes = true;
            return sql;
        }
        EncryptInfo *info = (EncryptInfo *) palloc0(sizeof(EncryptInfo));
        info->sql = sql;
        //新增错误处理机制
        char *err_msg_internal = NULL;
        std::string current_cmk;
        //得到当前用户的cmk
        // cmk_mapperGetCmkByUsername(user_name, current_cmk);
        //存储当前cmk
        // cmk_mapperSetCurrentCmk(current_cmk);
        //set user and db
        // cmk_mapperSetInfoUser(user_name);
        // cmk_mapperSetInfoDb(db_name);
        DekInterface::setInfoUser(user_name);
        DekInterface::setInfoDb(db_name);
        
        encParseTree = encryptStmt(parseTree, info, &err_msg_internal);
        if (err_msg_internal != NULL) {
            *err_msg = (char *) malloc(strlen(err_msg_internal) + 1); //手动分配内存，因为要保持c-style
            strcpy(*err_msg, err_msg_internal);
        }
        counter->count("SQL Encryption", t.passedTimeMicroSecond());
        if (!needDecryptResult(parseTree)) {
            SkipDecryptRes = true;
        }
        // mapperSave();
    }

    if (encParseTree) {
        Timer t;
        res = parseTree2str(encParseTree);
        counter->count("SQL AST To String", t.passedTimeMicroSecond());
    }

    if (!res)
        SkipDecryptRes = true;

    const char *retStuff = res ? res : sql;
#ifndef NDEBUG

    // std::cout << "PID:" << tid << std::endl;
    debug2file(sql, true);
    // std::cout << sql << std::endl;

    debug2file(retStuff, false);
    std::cout << retStuff << std::endl;
    std::cout << "main done\n";
#endif
    return retStuff;
}
