#include "ot_blinds.h"
#include <fcntl.h>
#include <cstdlib>
#include <iostream>
#include <ENCRYPTO_utils/connection.h> 
#include <fstream>
#include <x86intrin.h>  // SSE and AVX intrinsics
//const uint64_t n_OTS = 128;
 
ot_ext_prot test_prots[] = {IKNP};
snd_ot_flavor test_sflavor[] = {Snd_OT, Snd_C_OT, Snd_GC_OT, Snd_R_OT};
rec_ot_flavor test_rflavor[] = {Rec_OT, Rec_R_OT};
//uint64_t test_numots[] = {n_OTS, 3215, 100000};
uint64_t test_bitlen[] = {1, 3, 8, 191};
uint32_t test_nthreads[] = {1, 4};
field_type test_ftype[] = {P_FIELD, ECC_FIELD};
bool test_usemecr[] = {false};

BOOL Cleanup()
{
	delete sndthread;
	delete rcvthread;

	return true;
}

void InitSender(const std::string &address, const int port, CLock *glock)
{
	m_nPort = (uint16_t)port;
	m_nAddr = &address;

	//Server listen
	m_Socket = Listen(address, port);
	if (!m_Socket)
	{
		std::cerr << "Listen failed on " << address << ":" << port << "\n";
		std::exit(1);
	}

	sndthread = new SndThread(m_Socket.get(), glock);
	rcvthread = new RcvThread(m_Socket.get(), glock);

	sndthread->Start();
	rcvthread->Start();
}

void InitReceiver(const std::string &address, const int port, CLock *glock)
{
	m_nPort = (uint16_t)port;
	m_nAddr = &address;

	//Client connect
	m_Socket = Connect(address, port);
	if (!m_Socket)
	{
		std::cerr << "Connect failed on " << address << ":" << port << "\n";
		std::exit(1);
	}

	sndthread = new SndThread(m_Socket.get(), glock);
	rcvthread = new RcvThread(m_Socket.get(), glock);

	sndthread->Start();
	rcvthread->Start();
}

OTExtSnd *InitOTExtSnd(ot_ext_prot m_eProt, uint32_t nbaseots, uint32_t nchecks, bool enablemecr, field_type ftype, crypto *crypt)
{
	OTExtSnd *sender;
	switch (m_eProt)
	{
	case ALSZ:
		sender = new ALSZOTExtSnd(crypt, rcvthread, sndthread, nbaseots, nchecks);
		break;
	case IKNP:
		sender = new IKNPOTExtSnd(crypt, rcvthread, sndthread);
		break;
	case NNOB:
		sender = new NNOBOTExtSnd(crypt, rcvthread, sndthread);
		break;
	default:
		sender = new ALSZOTExtSnd(crypt, rcvthread, sndthread, nbaseots, nchecks);
		break;
	}

	if (enablemecr)
		sender->EnableMinEntCorrRobustness();
	sender->ComputeBaseOTs(ftype);
	return sender;
}

OTExtRec *InitOTExtRec(ot_ext_prot m_eProt, uint32_t nbaseots, uint32_t nchecks, bool enablemecr, field_type ftype, crypto *crypt)
{
	OTExtRec *receiver;
	switch (m_eProt)
	{
	case ALSZ:
		receiver = new ALSZOTExtRec(crypt, rcvthread, sndthread, nbaseots, nchecks);
		break;
	case IKNP:
		receiver = new IKNPOTExtRec(crypt, rcvthread, sndthread);
		break;
	case NNOB:
		receiver = new NNOBOTExtRec(crypt, rcvthread, sndthread);
		break;
	default:
		receiver = new ALSZOTExtRec(crypt, rcvthread, sndthread, nbaseots, nchecks);
		break;
	}

	if (enablemecr)
		receiver->EnableMinEntCorrRobustness();
	receiver->ComputeBaseOTs(ftype);
	return receiver;
}

// This is P0
void run_test_sender(uint32_t numots, uint32_t bitlength, snd_ot_flavor stype, rec_ot_flavor rtype, uint32_t numthreads,
					 crypto *crypt, OTExtSnd *sender, OTExtRec *receiver)
{
	CBitVector delta;
	CBitVector X0;
	uint32_t nsndvals = 2;
	CBitVector **X = (CBitVector **)malloc(sizeof(CBitVector *) * nsndvals);
	CBitVector **Y = (CBitVector **)malloc(sizeof(CBitVector *) * nsndvals);
	//The masking function with which the values that are sent in the last communication step are processed
	XORMasking *m_fMaskFct = new XORMasking(bitlength, delta);

	//creates delta as an array with "numOTs" entries of "bitlength" bit-values and fills delta with random values
	delta.Create(numots, bitlength, crypt);

	//Create X1 and X2 as two arrays with "numOTs" entries of "bitlength" bit-values and resets them to 0
	//X1.Create(numots, bitlength, crypt);
	//X2.Create(numots, bitlength, crypt);

	//X[0] --> T0 
	//X[1] --> X0
	//Y[0] --> Y0
	printf("run_test_sender\n");
	//(T, X0 \oplus T)
	for (uint32_t i = 0; i < nsndvals; i++)
	{
		X[i] = new CBitVector();
		// X[i]->Create(numots, bitlength, crypt);
		X[i]->Create(numots, bitlength);
        X[i]->FillRand(numots, crypt);

		Y[i] = new CBitVector();
		// X[i]->Create(numots, bitlength, crypt);
		Y[i]->Create(numots, bitlength);
       
	   // Y[i]->FillRand(n_OTS, crypt);
		
		 Y[i]->SetToOne();  // I added this line 
        //X[i]->PrintBinary();  // I added this line 

	}

	 
 	X0.Create(numots, bitlength);
	X0.Reset();
	X0.XOR(X[1]);
	//Y[0]->Copy(X[1]);
	
 #ifdef VERBOSE
  printf("X0:\n");
	 //X[1]->PrintBinary();  // I added this line 
	 X0.PrintHex();  // I added this line 
 	printf("Y0:\n");
 	Y[0]->PrintHex();  // I added this line
	 printf("T0: \n");
	 X[0]->PrintHex();
	#endif

	X[1]->XOR(X[0]);

	//X[1] --> X0 \oplus T0
	//for (uint32_t i = 0; i < nsndvals; i++)
	{
	//   printf("T0:\n");
    //   X[0]->PrintBinary();  // I added this line
    //   printf("X0 ^ T0:\n");
    //   X[1]->PrintBinary();  // I added this line 
	}

	sender->send(numots, bitlength, nsndvals, X, stype, rtype, numthreads, m_fMaskFct);
	CBitVector response;


	//Pre-generate the respose vector for the results
	response.Create(numots, bitlength);
	response.Reset();

	//Y[0] is the choice vector.
	receiver->receive(numots, bitlength, nsndvals, Y[0], &response, stype, rtype, numthreads, m_fMaskFct);
	
	
	
	response.XOR(X[0]);
	
 #ifdef VERBOSE
  printf("\n\nLearnt Gamma0 = (X1 /cdot Y0) /oplus T------>>>>>> \n");
  response.PrintHex();
 #endif

	  int fd = open("./y0", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 	  ssize_t bytes = write(fd, Y[0]->GetArr(), Y[0]->GetSize());
	  close(fd);

	  int fd2 = open("./x0", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 	  ssize_t bytes2 = write(fd2, X0.GetArr(), X0.GetSize());
	  close(fd2);

	   int fd3 = open("./gamma0", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 	  ssize_t bytes3 = write(fd3, response.GetArr(), response.GetSize());
	  close(fd3);

	// for(size_t j = 0; j < numots; ++j)
	// {
	// //	std::cout << (int) response.GetBitNoMask(j);
	// 	gammafile0 << (int) response.GetBitNoMask(j);
	// 	X0file 	  << (int) X0.GetBitNoMask(j);
	// 	Y0file 	  << (int) Y[0]->GetBitNoMask(j);
	// }

	for (uint32_t i = 0; i < nsndvals; i++)
	{
		delete (X[i]);
		delete (Y[i]);
	}
	free(X);
	free(Y);
 
	delta.delCBitVector();
	delete m_fMaskFct;
}


// This is P1
void run_test_receiver(uint32_t numots, uint32_t bitlength, snd_ot_flavor stype, rec_ot_flavor rtype, uint32_t numthreads,
					   crypto *crypt, OTExtSnd *sender, OTExtRec *receiver, int m_nPID)
{
	CBitVector X1, response;
	uint32_t nsndvals = 2;

	CBitVector **X = (CBitVector **)malloc(sizeof(CBitVector *) * nsndvals);
	CBitVector **Y = (CBitVector **)malloc(sizeof(CBitVector *) * nsndvals);
	

	//X[0] --> T1 
	//X[1] --> X1
	//Y[0] --> Y1
	for (uint32_t i = 0; i < nsndvals; i++)
	{
		X[i] = new CBitVector();
		X[i]->Create(numots, bitlength);
        X[i]->FillRand(numots, crypt);

		Y[i] = new CBitVector();
		Y[i]->Create(numots, bitlength);
        //Y[i]->FillRand(n_OTS, crypt);
		 Y[i]->SetToOne(); 
	}


	X1.Create(numots, bitlength);
	X1.Reset();
	X1.XOR(X[1]);
	
 #ifdef VERBOSE
  printf("X1:\n");
  //X[1]->PrintBinary();  // I added this line 
  X1.PrintHex();  // I added this line 
  printf("Y1:\n");
  Y[0]->PrintHex();  // I added this line
  printf("T1: \n");
  X[0]->PrintHex();
	#endif
	// X[1] -- > X1 \oplus T1
	X[1]->XOR(X[0]);


	//The masking function with which the values that are sent in the last communication step are processed
	XORMasking *m_fMaskFct = new XORMasking(bitlength);

 

 

	//Pre-generate the respose vector for the results
	response.Create(numots, bitlength);
	response.Reset();

	/*
	 * The inputs of the receiver in G_OT, C_OT and R_OT are the same. The only difference is the version
	 * variable that has to match the version of the sender.
	*/

	//Learns: (X0 \cdot Y1) \oplus T0
	//Y0 is the choice bits.
	receiver->receive(numots, bitlength, nsndvals, Y[0], &response, stype, rtype, numthreads, m_fMaskFct);

	sender->send(numots, bitlength, nsndvals, X, stype, rtype, numthreads, m_fMaskFct);

	// printf("Learnt: (X0 /cdot Y1) /oplus T0------>>>>>> \n");

	// response.PrintBinary();
	response.XOR(X[0]);
	
 #ifdef VERBOSE
  printf("\n\nLearnt: Gamma1 =  (X0 /cdot Y1) /oplus T------>>>>>> \n");
 	response.PrintHex();
	#endif
 //std::ofstream gammafile1, X1file, Y1file;
	
	

	int fd = open("./y1", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 ssize_t bytes = write(fd, Y[0]->GetArr(), Y[0]->GetSize());
 
 #ifdef VERBOSE
  std::cout << "bytes = " << bytes << std::endl;
	#endif

 close(fd);

	int fd2 = open("./x1", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 ssize_t bytes2 = write(fd2, X1.GetArr(), X1.GetSize());
 
	close(fd2);

	int fd3 = open("./gamma1", O_CREAT | O_WRONLY | O_TRUNC, 0600);;
 ssize_t bytes3 = write(fd3, response.GetArr(), response.GetSize());
 
	close(fd3);

  #ifdef VERBOSE
	   std::cout << "X1: \n\n";
	   X1.PrintHex();
   #endif

	for (uint32_t i = 0; i < nsndvals; i++)
	{
		delete (X[i]);
		delete (Y[i]);
	}
	free(X);
	free(Y);
	delete m_fMaskFct;
//	choices.delCBitVector();
	response.delCBitVector();
}


int main(int argc, char **argv)
{
	std::string addr0 = argv[1]; // "127.0.0.1";
	std::string addr1 = argv[2]; // "127.0.0.1";
	int port = 7766;

	__m128i blinds;


	// if (argc != 2)
	// {
	// 	std::cout << "Please call with 0 if acting as server or 1 if acting as client" << std::endl;
	// 	return EXIT_FAILURE;
	// }

	//Determines whether the program is executed in the sender or receiver role
	m_nPID 				= atoi(argv[3]);
	int64_t nblindbits  = atoi(argv[4]);


	std::cout << "Playing as role: " << m_nPID << std::endl;
	assert(m_nPID >= 0 && m_nPID <= 1);

	//The symmetric security parameter (80, 112, 128)
	uint32_t m_nSecParam = 128;

	crypto *crypt = new crypto(m_nSecParam, (uint8_t *)m_cConstSeed[m_nPID]);
	CLock *glock = new CLock(); // pass this to sender and receiver constructors

	uint32_t m_nBaseOTs = 190;
	uint32_t m_nChecks = 380;

	// NOTE: This vector controls the settings used by the oblivious transfer.
	test_options selected_options; // = {IKNP, 128, 1, Snd_C_OT, Rec_OT, 1, ECC_FIELD, false};
	selected_options.prot     = IKNP;
	selected_options.numots   = nblindbits; // Number of OTs performed using the extended COT.
	selected_options.bitlen   = 1;
	selected_options.sflavor  = Snd_OT; //Snd_C_OT;
	selected_options.rflavor  = Rec_OT;
	selected_options.nthreads = 1; // Number of threads
	selected_options.ftype    = P_FIELD; // Type of field to use for the base OT.
	selected_options.usemecr  = false;

	// test_options selected_options = {IKNP, 128, 1, Snd_C_OT, Rec_OT, 1, P_FIELD, false}; // Alternative using P_FIELD
	
	if (m_nPID == SERVER_ID) //Play as OT sender
	{
		InitSender(addr0, port, glock);
 
		OTExtSnd *sender = NULL;
		
		InitReceiver(addr1, port, glock);

		OTExtRec *receiver = NULL;

 
	 auto start = std::chrono::steady_clock::now();
		sender = InitOTExtSnd(selected_options.prot, m_nBaseOTs, m_nChecks, selected_options.usemecr, selected_options.ftype, crypt);
		receiver = InitOTExtRec(selected_options.prot, m_nBaseOTs, m_nChecks, selected_options.usemecr, selected_options.ftype, crypt);
		std::cout << "--> : " << getProt(selected_options.prot) << " Sender " << selected_options.numots << " " << getSndFlavor(selected_options.sflavor) << " / " << getRecFlavor(selected_options.rflavor) << " on " << selected_options.bitlen << " bits with " << selected_options.nthreads << " threads, " << getFieldType(selected_options.ftype) << " and" << (selected_options.usemecr ? "" : " no") << " MECR" << std::endl;

		run_test_sender(selected_options.numots, selected_options.bitlen, selected_options.sflavor, selected_options.rflavor, selected_options.nthreads, crypt, sender, receiver);

  auto end = std::chrono::steady_clock::now();
	 
  std::chrono::duration<double> elapsed_seconds = end - start;
	 std::cout << "time to generate and evaluate " << nblindbits << " OTs is: " << elapsed_seconds.count() << "s\n";
         std::cout << "bytes transmitted for OT = " << (54292 + 33 * nblindbits / 2) << " bytes\n";
		
  delete sender;
	}
	else //Play as OT receiver
	{
		InitReceiver(addr0, port, glock);

		OTExtRec *receiver = NULL;

			InitSender(addr1, port, glock);

	 auto start = std::chrono::steady_clock::now();
		OTExtSnd *sender = NULL;
		receiver = InitOTExtRec(selected_options.prot, m_nBaseOTs, m_nChecks, selected_options.usemecr, selected_options.ftype, crypt);
		sender = InitOTExtSnd(selected_options.prot, m_nBaseOTs, m_nChecks, selected_options.usemecr, selected_options.ftype, crypt);
		std::cout << "--> : " << getProt(selected_options.prot) << " Receiver " << selected_options.numots << " " << getSndFlavor(selected_options.sflavor) << " / " << getRecFlavor(selected_options.rflavor) << " on " << selected_options.bitlen << " bits with " << selected_options.nthreads << " threads, " << getFieldType(selected_options.ftype) << " and" << (selected_options.usemecr ? "" : " no") << " MECR" << std::endl;

		run_test_receiver(selected_options.numots, selected_options.bitlen, selected_options.sflavor, selected_options.rflavor, selected_options.nthreads, crypt, sender, receiver, m_nPID);

		delete receiver;
  auto end = std::chrono::steady_clock::now();
  std::chrono::duration<double> elapsed_seconds = end - start;
	 std::cout << "time to generate and evaluate " << nblindbits << " OTs is: " << elapsed_seconds.count() << "s\n";
         std::cout << "bytes transmitted for OT = " << (54274 + 33 * nblindbits / 2) << " bytes\n";
	}

	Cleanup();
	delete crypt;
	delete glock;

	return EXIT_SUCCESS;
}