

extern "C"
{

#include "postgres.h" // 与 #include <iomanip> 冲突

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

#include "ore_crypto.h"
#include "ore_blk.h"
#include "errors.h"
#include "ORE_STR22STRUCT.h"
}

#include <random>
#include "crypto/crypto.h"

#include <cstring>
#include <memory>
#include <string>
#include <unordered_set>

#include "MathUtils.h"
#include "Symcipher.h"
#include "ArraySymCipher.h"
#include "str22struct.h"
#include "SymPHE.h"
#include "SymAHE.h"
#include "SymMHE.h"
#include "t.h"
#include "KeyManager.h"
#include "TimeCounter.h"
#include "Singleton.h"
#include <iostream>
// #include "aes_cbc_iv_128_append.h"
#include <openssl/rand.h>
//#include "kms/KeyManagementService.hpp"

#include "sm4.h"
#include "Aes_Hmac.h"

using std::shared_ptr;
using std::string;
using std::unordered_set;


extern string ctext_to_str(char **dest, ore_blk_ciphertext cipher);
extern void cmk_mapperSetDekOnlyForOre(const std::string &dek);
extern void cmk_mapperClearDekOnlyForOre();

uint8_t *encryptValue(
		T_Cipher encryptCipher, uint8_t *in_text, size_t in_size, size_t *out_size, const char *dek, bool isEncrypt,
		bool isFloat) // this function will alloc memory
{
	*out_size = 0;

	TimeCounter *counter = Singleton<TimeCounter>::instance();

	shared_ptr<uint8_t> outtext(nullptr, [](uint8_t *p) { delete[] p; });
	shared_ptr<uint8_t> aesOut(new uint8_t[BUF_SIZE], [](uint8_t *p) { delete[] p; });

	char *aesKey = nullptr;
	char *aesKey1 = nullptr;

	string aesK = getAESKey();
	aesKey = (char *) aesK.c_str();
	aesKey1 = (char *) dek;

    //test only
//    std::string test_only = "oO^EH@27[AsHO1;I";
//    aesKey1 = (char *) test_only.c_str();

	// #if defined(SGX_ORE) || defined(SGX_MHE)
	// #else
	// #endif
	if (isEncrypt && (encryptCipher == CIPHER_SAHE || encryptCipher == CIPHER_SMHE || encryptCipher == CIPHER_ORE)) {
		if (in_size != sizeof(long)) { // maybe a float
			double tmp = atof((const char *) in_text);
			long longTmp = static_cast<long>(tmp);
			in_text = (uint8_t *) &longTmp;
			in_size = sizeof(long);
		}
	}

	if ((encryptCipher == CIPHER_AES || encryptCipher == CIPHER_DET||encryptCipher == CIPHER_AESHMAC) &&
	    (in_size % AES_BLOCK_SIZE != 0 || in_size < AES_BLOCK_SIZE)) {
		int cur_size = 0; // AES扩展明文成BLOCK倍

		int tmp_size = (in_size < AES_BLOCK_SIZE) ? AES_BLOCK_SIZE : in_size;
		while (tmp_size > 0) {
			cur_size += AES_BLOCK_SIZE;
			tmp_size -= AES_BLOCK_SIZE;
		}
		char *tmp_text = (char *) palloc0(cur_size);
		memcpy(tmp_text, (char *) in_text, in_size);
		in_text = (uint8_t *) tmp_text;
		in_size = cur_size;
	}

	if (isEncrypt) {
		switch (encryptCipher) {
			case CIPHER_SAHE: {
				SymAHE symAHE = SymAHE();
				if (in_size != sizeof(long))
					throw "in size does not match with long size in sahe";
				Timer t;
				outtext.reset((uint8_t *) symAHE.encrypt(*(long *) in_text, aesKey1));
				counter->count("SAHE Encryption", t.passedTimeMicroSecond());

				*out_size = strlen((const char *) outtext.get());
				break;
			}
			case CIPHER_SMHE: {
				SymMHE symMHE = SymMHE();
				if (in_size != sizeof(long))
					throw "in size does not match with long size in smhe";

				Timer t;
				outtext.reset((uint8_t *) symMHE.encrypt(*(long *) in_text, aesKey1));

				counter->count("SMHE Encryption", t.passedTimeMicroSecond());
				*out_size = strlen((const char *) outtext.get());
				break;
// #ifdef SGX_MHE
// 				encryption encryptCollection;
// 				if (!aesKey)
// 					throw "cannot get AES key\n";
// 				Timer t;

// 				*out_size = encryptCollection.aes_encrypt((char *)in_text, (char *)aesKey, (char *)aesOut.get(), in_size);

// 				counter->count("ORE_AES Encryption", t.passedTimeMicroSecond());

// 				outtext = aesOut;
// 				break;
// #else
// 				SymMHE symMHE = SymMHE();
// 				if (in_size != sizeof(long))
// 					throw "in size does not match with long size in smhe";

// 				Timer t;
// 				outtext.reset((uint8_t *) symMHE.encrypt(*(long *) in_text, aesKey1));

// 				counter->count("SMHE Encryption", t.passedTimeMicroSecond());
// 				*out_size = strlen((const char *) outtext.get());
// 				break;
// #endif
			}

			case CIPHER_DET:

			// 	encryption encryptCollection;
			// 	if (!aesKey1)
			// 		throw "cannot get AES key\n";
			// 	Timer t;

			// 	// *out_size = encryptCollection.aes_encrypt((char *) in_text, (char *) aesKey, (char *) aesOut.get(),
			// 	//                                           in_size);
			// 	*out_size = encryptCollection.aes_encrypt((char *) in_text, (char *) aesKey1, (char *) aesOut.get(),
			// 	                                          in_size);

			// 	counter->count("AES Encryption", t.passedTimeMicroSecond());

			// 	outtext = aesOut;
			// 	break;
			// }

			case CIPHER_AES:
			case CIPHER_AESHMAC: {
				// if (!aesKey1)
				// 	throw "cannot get AES key\n";
				// Timer t;
				// uint8_t *out=nullptr;
				// encrypt_aeshmack(in_text,in_size,out,*out_size,(uint8_t *)aesKey1);

				// outtext.reset(out);
				// counter->count("AES_HMAC Encryption", t.passedTimeMicroSecond());

				// break;
				
				encryption encryptCollection;
				if (!aesKey1)
					throw "cannot get AES key\n";
				Timer t;
				// 确保aesOut缓冲区足够大（新增）
				int required_size = in_size + 16 + 32; // 原始数据 + AES块填充 + HMAC
				if (required_size > BUF_SIZE) {
					aesOut.reset(new uint8_t[required_size], [](uint8_t *p) { delete[] p; });
				}

				*out_size = encryptCollection.aeshmac_encrypt((char *) in_text, in_size, (char *) aesKey1, (char *) aesOut.get());
				counter->count("AES Encryption", t.passedTimeMicroSecond());
				outtext = aesOut;
				break;
				
			}

			// case CIPHER_SM4: {
			// 	sm4 sm4((uint8_t *)aesKey1);
			// 	uint8_t	*out;
			// 	sm4.encrypt_ecb(in_text, in_size, out, *out_size);
			// 	outtext.reset(out);
			// 	break;
			// }
			case CIPHER_SM4:
			case CIPHER_SM4CK: {
				sm4 sm4((uint8_t*)aesKey1);
				uint8_t	*out = nullptr;

				int ret = sm4.encrypt_ecb_with_checksum(in_text, in_size, out, *out_size);
				if(ret == -1){
					string err_msg = "eqsm4 encrypt data failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}else if(ret == -2){
					string err_msg = "eqsm4 sm3-hmac calculation failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}
				outtext.reset(out);
				break;
			}

			case CIPHER_ORE: {

// #if defined(SGX_ORE) || defined(SGX_MHE)
#if defined(SGX_ORE)
				encryption encryptCollection;
				if (!aesKey)
					throw "cannot get AES key\n";
				Timer t;

				*out_size = encryptCollection.aes_encrypt((char *)in_text, (char *)aesKey, (char *)aesOut.get(), in_size);

				counter->count("ORE_AES Encryption", t.passedTimeMicroSecond());

				outtext = aesOut;
				break;
#else
				Timer t;
				ore_blk_params params;
//                cmk_mapperSetDekOnlyForOre(aesKey1);
				init_ore_blk_params(params, ORE_nbits, ORE_block_len);

				ore_blk_secret_key sk;
				ore_blk_setup(sk, params, aesKey1);

				ore_blk_ciphertext ctxt1;
				init_ore_blk_ciphertext(ctxt1, params);

				ore_blk_encrypt_ui(ctxt1, sk, *(long *) in_text);

				char *__unused = nullptr;
				// ore_blk_ciphertext cipher2;
				string str = ctext_to_str(&__unused, ctxt1);

				outtext.reset(new uint8_t[str.size() + 1]());

				counter->count("ORE Encryption", t.passedTimeMicroSecond());

				memcpy(outtext.get(), str.c_str(), str.size());

				*out_size = str.size();
//                cmk_mapperClearDekOnlyForOre();
				break;
#endif
			}

				//新增RND
// 			case CIPHER_RND: {
// //				 const unsigned char *key = (const unsigned char *) aesKey; // 使用上文生成的aesKey
// 				const unsigned char *key = (const unsigned char *) aesKey1; // 使用上文生成的aesKey

// 				const char *plaintext = (char *) in_text;
// 				int plaintext_len = strlen(plaintext);
// 				unsigned char encrypted[plaintext_len + AES_BLOCK_LEN * 2];
// 				int padded_plaintext_len = 0;
// 				padded_plaintext_len = aes_encrypt((const unsigned char *) plaintext, plaintext_len, key, encrypted);
// 				uint8_t *newBuffer = new uint8_t[padded_plaintext_len + AES_BLOCK_LEN];
// 				memcpy(newBuffer, encrypted, padded_plaintext_len + AES_BLOCK_LEN);
// 				outtext.reset(newBuffer);
// 				*out_size = padded_plaintext_len + AES_BLOCK_LEN;
// 				//test_only
// //				unsigned char test[plaintext_len + AES_BLOCK_LEN * 2];
// //				aes_decrypt((const unsigned char *) outtext.get(), padded_plaintext_len, key, test, padded_plaintext_len);
// 				break;
// 			}
			case CIPHER_RND:
			case CIPHER_RNDSM4CK: {
				sm4 sm4((uint8_t *) aesKey1);
				uint8_t *out = nullptr;
				int ret = sm4.encrypt_gcm(in_text, in_size, out, *out_size);
				if(ret == -1){
					string err_msg = "rnd gmac calculation failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
				}
				outtext.reset(out);
				break;
			}

			case CIPHER_NOCRYPT:

				outtext.reset(new uint8_t[in_size]());
				memcpy(outtext.get(), in_text, in_size);
				*out_size = in_size;
				break;

				// case CIPHER_NOCRYPT:
				//     outtext = in_text;
				//     *out_size = strlen((const char *)outtext);
				//     break;
		}
	} else { // decrypt

		switch (encryptCipher) {
			case CIPHER_SAHE: {
				in_text[in_size - 1] = 0; // 将浮点数标识位取消
				SymAHE symAHE = SymAHE();
				long *tmp = new long;
				Timer t;
				*tmp = symAHE.decrypt((char *) in_text, aesKey1);
				counter->count("SAHE Decryption", t.passedTimeMicroSecond());

				outtext.reset((uint8_t *) tmp);
				*out_size = sizeof(long);
				break;
			}

			case CIPHER_SMHE: {
				in_text[in_size - 1] = 0; // 将浮点数标识位取消
				SymMHE symMHE = SymMHE();
				long long *tmp = new long long;
				Timer t;
				*tmp = symMHE.decrypt((char *) in_text, aesKey1);
				counter->count("SMHE Decryption", t.passedTimeMicroSecond());

				outtext.reset((uint8_t *) tmp);
				*out_size = sizeof(long);

				break;
			}

			case CIPHER_DET:

			// case CIPHER_AES: {
			// 	encryption encryptCollection;
			// 	Timer t;
			// 	// *out_size = encryptCollection.aes_decrypt((char *) in_text, in_size, (char *) aesKey,
			// 	//                                           (char *) aesOut.get());
			// 	*out_size = encryptCollection.aes_decrypt((char *) in_text, in_size, (char *) aesKey1,
			// 	                                          	(char *) aesOut.get());
			// 	counter->count("AES Decryption", t.passedTimeMicroSecond());

			// 	outtext = aesOut;
			// 	break;
			// }
			case CIPHER_AES:
			case CIPHER_AESHMAC: {
				
				// Timer t;
				// uint8_t *out=nullptr;
				// int ret = decrypt_aeshmack(in_text,in_size,out,*out_size,(uint8_t *)aesKey1);
				// if(ret == -1){
				// 	string err_msg = "eq decrypt data failed!";
				// 	errmsg(err_msg.c_str());
				// 	(*out_size) = 0;
				// 	break;
				// }else if(ret == -2){
				// 	string err_msg = "eq hmac verification failed!";
				// 	errmsg(err_msg.c_str());
				// 	(*out_size) = 0;
				// 	break;
				// }else if(ret == -3){
				// 	string err_msg = "eq invalid input encrypted data!";
				// 	errmsg(err_msg.c_str());
				// 	(*out_size) = 0;
				// 	break;
				// }

				// counter->count("AES_HMAC Decryption", t.passedTimeMicroSecond());
				
				// outtext.reset(out);

				// (*out_size)+= 1;
				// break;
				/*
					//原版
					// encryption encryptCollection;
					// Timer t;
					// *out_size = encryptCollection.aeshmac_decrypt((char *) in_text, in_size, (char *) aesKey1,
					//                                           (char *) aesOut.get());
					// counter->count("AES Decryption", t.passedTimeMicroSecond());

					// outtext = aesOut;
					// break;
				*/

				
				encryption encryptCollection;
				Timer t;
				// 先计算需要的缓冲区大小，确保足够大
				int required_size = in_size; // 最坏情况下，解密后的数据大小可能接近加密数据大小
				if (required_size > BUF_SIZE) {
					aesOut.reset(new uint8_t[required_size], [](uint8_t *p) { delete[] p; });
				}
				// 调用解密函数
				int decrypt_result = encryptCollection.aeshmac_decrypt((char *) in_text, in_size, (char *) aesKey1,
				                                          (char *) aesOut.get());
				if (decrypt_result < 0) {
					// 解密失败处理
					*out_size = 0;
					break;
				}
				// 设置正确的输出大小
				*out_size = decrypt_result;
				outtext = aesOut;
				counter->count("AES Decryption", t.passedTimeMicroSecond());
				break;
			}

			// case CIPHER_SM4: {
			// 	sm4 sm4((uint8_t *)aesKey1);
			// 	uint8_t *out;
			// 	sm4.decrypt_ecb(in_text, in_size, out, *out_size);
			// 	outtext.reset(out);
			// 	(*out_size) += 1;
			// 	break;
			// }
			case CIPHER_SM4:
			case CIPHER_SM4CK: {
				sm4 sm4((uint8_t *)aesKey1);
				uint8_t *out;
				int ret = sm4.decrypt_ecb_with_checksum(in_text, in_size, out, *out_size);
				if(ret == -1) {
					string err_msg = "eqsm4 decrypt data failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}else if(ret == -2){
					string err_msg = "eqsm4 sm3-hmac verification failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}else if(ret == -3){
					string err_msg = "eqsm4 invalid input encrypted data!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}
				outtext.reset(out);
				(*out_size) += 1;
				break;
			}

			case CIPHER_NOCRYPT:

				// outtext = encryptCollection. in_text;
				outtext.reset(new uint8_t[in_size]());
				memcpy(outtext.get(), in_text, in_size);
				*out_size = in_size;
				// outtext = in_text;
				break;
			case CIPHER_ORE:
				// ! ORE has No Decrypt!
				ereport(FATAL,
				        (errcode(ERRCODE_CANNOT_CONNECT_NOW),
						        errmsg("Can not decrypt ORE!\n")));
				exit(-1);
				break;
// 			case CIPHER_RND: {
// 				//TODO 新增解密功能
// //				 const unsigned char *key = (const unsigned char *) aesKey;
// 				const unsigned char *key = (const unsigned char *) aesKey1;
// 				unsigned char out_tmp[in_size + AES_BLOCK_LEN];
// 				auto plain_size = aes_decrypt((const unsigned char *) in_text, in_size, key, out_tmp, in_size);
// 				if (plain_size < AES_BLOCK_SIZE) {//这里应该是啥，先用16保证测试顺利
// 					*out_size = AES_BLOCK_SIZE;
// 				} else {
// 					*out_size = ((plain_size / AES_BLOCK_LEN) + 1) * AES_BLOCK_SIZE;
// 				}
// 				outtext.reset(new uint8_t[*out_size]());
// 				std::memcpy(outtext.get(), out_tmp, plain_size);
// 				break;
// 			}
			case CIPHER_RND:
			case CIPHER_RNDSM4CK: {
				sm4 sm4((uint8_t *)aesKey1);
				uint8_t *out;
				int ret = sm4.decrypt_gcm(in_text, in_size, out, *out_size);
				if(ret == -1) {
					string err_msg = "rnd gmac verification failed!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}else if(ret == -3){
					string err_msg = "rnd invalid input encrypted data!";
					errmsg(err_msg.c_str());
					(*out_size) = 0;
					break;
				}
				outtext.reset(out);
				(*out_size) += 1;
				break;
		}
		}

	}

	char *pallocbuf = (char *) palloc0(*out_size + 2);
	memset(pallocbuf, 0, *out_size + 2);
	memcpy(pallocbuf, (const char *) outtext.get(), *out_size);
	
	if (encryptCipher == CIPHER_SAHE || encryptCipher == CIPHER_SMHE) {
		size_t cipherSize = strlen((char *) outtext.get());
		//将密文的最后一位用作标识是否是浮点数
		if (isEncrypt) {
			pallocbuf[cipherSize] = isFloat ? '2' : '1';
		}
		(*out_size) += 2;
	}

	return (uint8_t *) pallocbuf;
}

size_t getCipherFormatSize(const char *fmt, T_Cipher t) {
	if (!strcmp("hexstr", fmt))
		return getCipherBufSize(t) * 2 + 1;
	if (!strcmp("base64", fmt))
		return getCipherBufSize(t) * 1.4 + 1;
	return -1;
}
