

#ifdef __cplusplus
extern "C"
{
#endif
#include <libintl.h>
#include "postgres.h"
#ifdef __cplusplus
}
#endif

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

#include <string>
#include <vector>
#include <string>
#include <sstream>
#include <memory>
#include <pthread.h>
#include <fstream>
#include <iostream>
#include "utils.h"
#include "TimeCounter.h"
#include "Singleton.h"
#include "encryption.h"
#include "t.h"
#include "KeyManager.h"
// #include "kms/KeyManagementService.hpp
#include "kmsAdapter/dek_interface.hpp"
#include "en_config.h"

extern __thread bool Encrypt;
extern __thread bool SkipDecryptRes;
extern thread_local std::string dek_for_decrypt;

using std::shared_ptr;

extern "C" void decryptResult(int numberAttr, int numTuples, pAttrDescs pattDescs, pTuples ptuples)
__attribute__((visibility("default")));

void decryptResult(int numberAttr, int numTuples, pAttrDescs pattDescs, pTuples ptuples) {
	using namespace std;
	if (!Encrypt || SkipDecryptRes)
		return;

#ifndef NDEBUG
	// std::cout << "begin to decrypt: " << numberAttr << endl;
	// std::cout << retStuff << std::endl;
#endif

	TimeCounter *counter = Singleton<TimeCounter>::instance();
	Timer timer;
	PGresAttValue **tuples = (PGresAttValue **) ptuples;
	PGresAttDesc *attDescs = (PGresAttDesc *) pattDescs;
	for (int j = 0; j < numberAttr; ++j) {
		string cipherColName = attDescs[j].name;
		size_t oldAttSize = cipherColName.size();
		// int cipherColnameSZ = attDescs[i].name;
		string colnameSubfix = getMappedName(T_STRING_COLUMN, cipherColName.c_str(), NULL);
		if(colnameSubfix == attDescs[j].name) continue;
		string colname = colnameSubfix, subfix;
		bool unnamedFlag = false;

		auto ccit = cipherColName.find_first_of("_");

		if (ccit != cipherColName.npos) {
			auto prefix = cipherColName.substr(0, ccit);
			if (!prefix.compare("restmp")) {
				if (cipherColName.find("NOCRYPT") != string::npos) {
					continue;  // 如果是NOCRYPT后缀，跳过解密
				}
				unnamedFlag = true;
			}
		}

		if (!strcmp(cipherColName.c_str(), "?column?")) {
			subfix = "AES";
			colname = "unnamedCol";
			unnamedFlag = true;
		} else {
			auto it = colnameSubfix.find_last_of("_");
			if (it != colnameSubfix.npos) {
				colname = colnameSubfix.substr(0, it);
				subfix = colnameSubfix.substr(it + 1, colnameSubfix.size() - it);
			}
		}

		if (string2TypeCipher(subfix.c_str()) == CIPHER_COUNT) //别名中出现了下划线的情况。
			colname = colnameSubfix;
		strcpy(attDescs[j].name, colname.c_str()); // WARN
		if (strlen(attDescs[j].name) > oldAttSize)
			abort();
		shared_ptr<char> coltypename(new char[32](), [](char *p) { delete[] p; });
		getColumnType(colname.c_str(), coltypename.get());
		if (!*coltypename && !unnamedFlag) { // 这个column不在记录当中
			continue;
		}
		COLUMN_TYPE ctype = TYPE_STRING;
		if (!unnamedFlag)
			ctype = unifyColumnType(coltypename.get());
		else
			ctype = TYPE_INTEGER;

		void (*anyTypetoString)(const char *, char *const, void *) = NULL; // 把buffer转为相应类型的数据。
		switch (ctype) {
			// case TYPE_FLOAT:
			// 	anyTypetoString = &FloatTypetoString;
			// 	break;

			case TYPE_STRING:

				anyTypetoString = NULL; //无需转换
				break;
			case TYPE_INTEGER:
				anyTypetoString = &IntTypetoString;
				break;
		}

		if (!subfix.empty()) {
			T_Cipher type = string2TypeCipher(subfix.c_str());
			if (type == CIPHER_NOCRYPT)
				continue;
			//TODO 是否可以用OMP加速
			for (int i = 0; i < numTuples; ++i) {
				if (tuples[i][j].len == -1) // 数据列为Null
					continue;
				size_t oldDataLen = tuples[i][j].len;
				if(strncmp(tuples[i][j].value, "SGX:", 4) == 0) {
					// 是SGX环境下的AES加密数据
					// 使用AES解密
                    const char* encryptedData = tuples[i][j].value + 4;
                    // 获取正确的缓冲区大小
                    size_t buffSZ = strlen(encryptedData) / 2;
                    uint8_t *buff = new uint8_t[buffSZ]();
                    
                    // 转换十六进制字符串为二进制
                    t trans;
                    size_t actualSize = buffSZ;
                    trans.hex2buf(buff, &actualSize, encryptedData, strlen(encryptedData));   
                    
                    // 从文件中读取DEK
                    std::string dek;
                    std::ifstream dekFile(DEK_FILE_PATH, std::ios::binary);
                    if (dekFile.is_open()) {
                        // 获取文件大小
                        dekFile.seekg(0, std::ios::end);
                        size_t fileSize = dekFile.tellg();
                        dekFile.seekg(0, std::ios::beg);
                        
                        // 读取文件内容
                        dek.resize(fileSize);
                        dekFile.read(&dek[0], fileSize);
                        dekFile.close();
                    } else {
                        std::cout << "Failed to open DEK file, using default key" << std::endl;
                    }
                    
                    shared_ptr<uint8_t> aesOut(new uint8_t[BUF_SIZE], [](uint8_t *p) { delete[] p; });
                    encryption encryptCollection;
                    size_t outSize = encryptCollection.aeshmac_decrypt(
                        (char*)buff, 
                        actualSize,  // 使用实际的缓冲区大小
                        (char*)dek.c_str(),  // 使用DEK
                        (char*)aesOut.get()
                    );
                    // 确保输出大小正确
                    if (outSize >= sizeof(double)) {
                        double d = *(double*)aesOut.get();
                        std::ostringstream ss;
                        ss << d;
                        strcpy(tuples[i][j].value, ss.str().c_str());
                    } 
                    delete[] buff;
				} else {
					size_t buffSZ = tuples[i][j].len / 2 + 1; // 密文buffer长度， 应该大于hexstr长度的一半
					size_t plainBuffSZ = buffSZ;
					uint8_t *plainBuff = nullptr;
					uint8_t *pbuff = nullptr;
					shared_ptr<uint8_t> buff(new uint8_t[buffSZ](), [](uint8_t *p) { delete[] p; });
					if (type != CIPHER_SAHE && type != CIPHER_SMHE)  // true 为采用encode的加密列
					{
						hex2buf(buff.get(), &buffSZ, tuples[i][j].value, tuples[i][j].len); // 若要改变编码方式，主要换掉这个函
						pbuff = buff.get();
					} else {
						pbuff = (uint8_t *) tuples[i][j].value;
						buffSZ = tuples[i][j].len + 1;
					}
				// }
				// 当前这列的密文为float类型，那么就需要转换
				// if ((t == CIPHER_SAHE || t == CIPHER_SMHE) &&
				//     (pbuff[strlen((char *) pbuff) - 1] == '2' && anyTypetoString != FloatTypetoString)) {
				// 	anyTypetoString = FloatTypetoString;
				// }
				// if ((t == CIPHER_SAHE || t == CIPHER_SMHE) &&
				//     (pbuff[strlen((char *) pbuff) - 1] == '3' && anyTypetoString != FloatTypetoString)) {
				// 	anyTypetoString = FloatTypetoString2;
				// }
				// if ((t == CIPHER_SAHE || t == CIPHER_SMHE) &&
				//     (pbuff[strlen((char *) pbuff) - 1] == '4' && anyTypetoString != FloatTypetoString)) {
				// 	anyTypetoString = FloatTypetoString3;
				// }
				// if(CIPHER_NOCRYPT == t)
				//     buffSZ = tuples[i][j].len;

				//KMS test需要动态选择dek，rnd如果选错了dek直接报错
				// const char *p = dek_for_decrypt.c_str();

					std::string dek;
					// cmk_mapperGetDekByCol(cipherColName,dek);
					DekInterface::getDekColLevel(cipherColName, dek);
					plainBuff = encryptValue(type, pbuff, buffSZ, &plainBuffSZ, dek.c_str(), false);
					//KMS test end;
	//				plainBuff = encryptValue(t, pbuff, buffSZ, &plainBuffSZ, false);
					if (anyTypetoString) {
						//fixed by qxy for RND
						if (type == CIPHER_RND) {
							long l = *(long *)plainBuff;
							std::ostringstream ss;
							ss << l;
							strcpy(tuples[i][j].value, ss.str().c_str());
						} else {
							anyTypetoString((const char *) plainBuff, tuples[i][j].value, (void *) coltypename.get());
						}
						//anyTypetoString((const char *) plainBuff, tuples[i][j].value, (void *) coltypename.get());
					} else {
						memcpy(tuples[i][j].value, (const char *) plainBuff,
							min(strlen((const char *) plainBuff) + 1, plainBuffSZ));
					}
				}
				//edit  tuples[i][j].len = strlen((char *)tuples[i][j].value);以上原来全部注释
				/*
				shared_ptr<char> res(new char[plainBuffSZ]());
				if (ctype != TYPE_STRING) {
					if(pbuff[buffSZ-2] == '2') {    // 此密文是一个浮点数，需要除scale
						FloatTypetoString((const char *)plainBuff, tuples[i][j].value, (void *)coltypename.get());
					}
					else {
						IntTypetoString((const char *)plainBuff, tuples[i][j].value, (void *)coltypename.get());
					}
				}
				else {
					memcpy(tuples[i][j].value, (const char *)plainBuff, min(strlen((const char *)plainBuff) + 1, plainBuffSZ));
				}
				strcpy(tuples[i][j].value, (char*)plainBuff);
				strncpy(tuples[i][j].value, res.get(), min(strlen(res.get()), tuples[i][j].len) + 1);
				*/
				tuples[i][j].len = strlen((char *) tuples[i][j].value);
				if (tuples[i][j].len > oldDataLen)
					abort(); // 出现了内存越界
				// clean
				// delete[] buff;
				// delete[] res;
			}
		}
	}
#ifndef NDEBUG
	// std::cout << "decrypt done\n";
	// std::cout << retStuff << std::endl;
#endif

	counter->count("SQL Decryption", timer.passedTimeMicroSecond());

	cleanup();
}