#include <stdio.h>
#include <time.h>
#include "kem_api.h"
#include "piglet1.h"
#include "ffi_field.h"
#include "ffi_elt.h"
#include "ffi_vec.h"
#include "ffi_mat.h"
#include "parameters.h"
#include "rng.h"
#include "parsing.h"
#include<windows.h>

#pragma warning(disable : 4996)

FILE *fpVector;

void write_buff(FILE *fp, unsigned char *buff, long len)
{
	unsigned char len_buff[4];
	len_buff[0] = (unsigned char)(len >> 24);
	len_buff[1] = (unsigned char)(len >> 16);
	len_buff[2] = (unsigned char)(len >> 8);
	len_buff[3] = (unsigned char)(len) ;
	fwrite(len_buff, 1, 4, fp);
	fwrite(buff, 1, len, fp);
}

int main(int argc, char *argv[]) {
    unsigned char pk[KEM_PUBLICKEYBYTES]; // public key
    unsigned char sk[KEM_SECRETKEYBYTES]; // secret key，pk appended
    unsigned char ss1[KEM_BYTES]; // shared secret
    unsigned char ss2[KEM_BYTES]; // shared secret
	unsigned char ct[CIPHERTEXT_BYTES];

	ffi_vec u[PARAM_K], v;
	int sum = 0;
	clock_t start[3],end[3];
	unsigned long total_time[3]={0,0,0};
	unsigned int experiment = 1;

	char filepath[1024] = "";
	strcpy(filepath, argv[0]);
	//获取exe执行文件名字(去掉后缀".exe")
	*strrchr(filepath, '\\') = '\0';   // 从最左边开始最后一次出现"."
	strcpy(filepath + strlen(filepath), "\\KEM_VEC_INFO.dat");
	fpVector = fopen(filepath, "wb");
	if (!fpVector) {
		printf("\n测试向量文件创建失败!");
		return -1;
	}

	//做好域初始化
	ffi_field_init();

	setbuf(stdin, NULL);
	printf("\n请输入加解密测试次数: ");
	scanf_s("%d",&experiment);
	printf("\n正在进行加解密测试....(%d次)\n", experiment);
	// 加解密时间测试
	for (unsigned int i = 0; i < experiment; i++) {
		// 密钥生成，保存在 pk and sk.
		start[0] = clock(); 
		piglet1_pke_keygen(pk,sk);
		end[0] = clock();
		total_time[0] = total_time[0] + end[0] - start[0];

		// 随机选择明文
		ffi_vec m;
		m.SetLength(PARAM_L);
		ffi_vec_set_random_using_rng(m, PARAM_L);
	
		// Generating G function
		unsigned char seed_G[VEC_L_BYTES];
		ffi_vec_to_string(seed_G, m, PARAM_L);
		AES_XOF_struct* G_seedexpander = (AES_XOF_struct*)malloc(sizeof(AES_XOF_struct));
		seedexpander_init(G_seedexpander, seed_G, seed_G + 28, SEEDEXPANDER_MAX_LENGTH);

		// Computing theta
		unsigned char theta[SK_SEED_LEN];
		seedexpander(G_seedexpander, theta, SK_SEED_LEN);

		// 用pk加密明文m，同时形成密文u、v
		start[1] = clock();
		piglet1_pke_encryption(m, pk, theta, u, v);
		end[1] = clock();
		total_time[1] = total_time[1] + end[1] - start[1];

		// 用 sk 解密密文u、v，生成明文dec_m
		ffi_vec dec_m;
		dec_m.SetLength(PARAM_L);
		
		start[2] = clock();
		piglet1_pke_decryption(sk, dec_m, u, v);
		end[2] = clock();
		total_time[2] = total_time[2] + end[2] - start[2];

		if (m==dec_m)	sum++;
		printf(". ");

	}

	printf("\n\n加解密测试次数: %d\n", sum);
	printf("密钥生成平均时间： %ums\n", total_time[0]/experiment);
	printf("公钥加密平均时间： %ums\n", total_time[1] / experiment);
	printf("解密算法平均时间： %ums\n", total_time[2]/ experiment);

	//封装与解封装测试
	experiment = 1;
	total_time[0] = 0;
	total_time[1] = 0;
	total_time[2] = 0;

	setbuf(stdin,NULL);
	printf("\n\n请输入封装与解封装测试次数: ");
	scanf_s("%d",&experiment);

	printf("\n正在进行封装与解封测试....(%d次)\n",experiment);
	for (unsigned int i = 0; i < experiment; i++) {
		// 密钥生成，保存在 pk and sk.
		start[0] = clock();
		piglet1_pke_keygen(pk, sk);
		end[0] = clock();
		total_time[0] = total_time[0] + end[0] - start[0];
		// 保存测试向量,只保存第一次的，内容依次是：公钥种子、私钥、公钥、密文、共享密钥
		//每个部分的格式是 Big endian的长度、信息内容
		if (i == 0) {
			write_buff(fpVector, pk + VEC_R_BYTES * PARAM_K + GD_SEED_LEN, PK_SEED_LEN);
			write_buff(fpVector, sk, KEM_SECRETKEYBYTES);
			write_buff(fpVector, pk, KEM_PUBLICKEYBYTES);
		}

		//封装
		start[1] = clock();
		kem_enc(pk, ss1, ct);
		end[1] = clock();
		total_time[1] = total_time[1] + end[1] - start[1];

		// 保存测试向量,只保存第一次的，内容依次是：公钥种子、私钥、公钥、密文、共享密钥
		//每个部分的格式是 Big endian的长度、信息内容
		if (i == 0) {
			write_buff(fpVector, ct, CIPHERTEXT_BYTES);
			write_buff(fpVector, ss1, KEM_BYTES);
		}
		// 解封装
		start[2] = clock();
		kem_dec(sk, ct, ss2);
		end[2] = clock();
		total_time[2] = total_time[2] + end[2] - start[2];
		printf(". ");
		//比较结果，算法确保了成功，结果不做比较

	}

	printf("\n封装与解封装次数: %d\n\n", experiment);
	printf("密钥生成平均时间: %ums\n", total_time[0] / experiment);
	printf("    封装平均时间: %ums\n", total_time[1] / experiment);
	printf("  解封装平均时间:  %ums\n", total_time[2] / experiment);
	
	//关闭测试向量文件
	if (fpVector) {
		fclose(fpVector);
	}

	setbuf(stdin,NULL);
	getchar();
    return 0;
}
