
#include <type_traits>  // std::is_same<>
#include <limits>       // std::numeric_limits<>
#include <climits>      // CHAR_BIT
#include <cmath>        // std::log2, std::ceil, std::floor
#include <stdexcept>    // std::runtime_error
#include <array>        // std::array<>
#include <iostream>     // std::istream and std::ostream
#include <vector>       // std::vector<>
#include <memory>       // std::shared_ptr<>
#include <utility>      // std::move
#include <algorithm>    // std::copy
#include <cstring>      // std::memcpy
#include <bsd/stdlib.h> // arc4random_buf
#include <x86intrin.h>  // SSE and AVX intrinsics
#include <chrono>
#include <thread>
#include <deque>
#include <../boost/asio.hpp>

using boost::asio::ip::tcp;
using namespace std::chrono;
using namespace std;
using socket_t = boost::asio::ip::tcp::socket;

#include "prg.h"
#include "prg_aes_impl.h"
#include "block.h"  
#include "duoram-utils.h"
#include "readvectors.h"
#include "duoram-read.h"
#include "duoram-refresh.h"
#include "duoram-write.h"



 
void accept_conncections_from_Pb(boost::asio::io_context&io_context, std::vector<socket_t>& sockets_, int port, size_t j)
{
  tcp::acceptor acceptor_a(io_context, tcp::endpoint(tcp::v4(), port));
	 tcp::socket sb_a(acceptor_a.accept());
	 sockets_[j] = std::move(sb_a);    
}

 

int main(const int argc, const char * argv[])
{
	
  size_t expo =  atoi(argv[3]);
  size_t db_nitems = 1ULL << expo;
  size_t number_of_writes	    =  atoi(argv[4]);
  size_t number_of_ind_reads  =  atoi(argv[5]); 
  size_t number_of_dep_reads  =  atoi(argv[6]);
  size_t number_of_accesses   =  atoi(argv[7]);
   
  reading_temp = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  DB = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  updated_DB = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  blinded_DB = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  	
  blinded_DB_recv = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  updated_blinded_DB_recv = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
	 
  b = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  c = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
  d = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
		
	 reading_b = (int8_t *) malloc(db_nitems * sizeof(int8_t));
	 reading_c = (int8_t *) malloc(db_nitems * sizeof(int8_t));
	 reading_d = (int8_t *) malloc(db_nitems * sizeof(int8_t));
	 writing_b = (int8_t *) malloc(db_nitems * sizeof(int8_t));
	 writing_c = (int8_t *) malloc(db_nitems * sizeof(int8_t));
	 writing_d = (int8_t *) malloc(db_nitems * sizeof(int8_t));

 
  size_t * rotate_by_ = new size_t[number_of_writes];


  boost::asio::io_context io_context;
  tcp::resolver resolver(io_context);
  std::string addr = "127.0.0.1";
  
  const std::string host1 = (argc < 2) ? "127.0.0.1" : argv[1];
  const std::string host2 = (argc < 3) ? "127.0.0.1" : argv[2];

  
  bool party;
  const size_t number_of_sockets = 40;
  std::vector<socket_t> sockets_;
  for(size_t j = 0; j < number_of_sockets + 1; ++j)
  {
 	 	tcp::socket emptysocket(io_context);
 		 sockets_.emplace_back(std::move(emptysocket));
  }
	 sockets_.reserve(number_of_sockets + 1);
 	printf("number_of_sockets = %zu\n", number_of_sockets);

	 std::vector<socket_t> sockets_2;

	 std::vector<int> ports;
	 for(size_t j = 0; j < number_of_sockets; ++j) 
	 {
	 	int port = 6000;
		 ports.push_back(port + j);
 	}
	
	 std::vector<int> ports2_0;
	 for(size_t j = 0; j < number_of_sockets; ++j) 
	 {
		 int port = 8000;
		 ports2_0.push_back(port + j);
	 }

	std::vector<int> ports2_1;
	for(size_t j = 0; j < number_of_sockets; ++j) 
	{
		int port = 9000;
		ports2_1.push_back(port + j);
	}



#if (PARTY == 0)    
  	 
 	party = false;

 	#ifdef ThreeParty	
		for(size_t j = 0; j < number_of_sockets; ++j)
		{
			tcp::socket sb_a(io_context);
			boost::asio::connect(sb_a, resolver.resolve({host2, std::to_string(ports2_0[j])}));
			sockets_2.emplace_back(std::move(sb_a)); 
		}
	#endif
  for(size_t j = 0; j < number_of_sockets; ++j)
  {
   tcp::socket sb_a(io_context);
    boost::asio::connect(sb_a, resolver.resolve({host1, std::to_string(ports[j])}));
    sockets_[j] = std::move(sb_a); 
  }
 #else	

	party = true;	
 	
	#ifdef ThreeParty	
	 for(size_t j = 0; j < number_of_sockets; ++j)
  {
 	 tcp::socket sb_a(io_context);
  	boost::asio::connect(sb_a, resolver.resolve({host2, std::to_string(ports2_1[j])}));
  	sockets_2.emplace_back(std::move(sb_a)); 
	 }
	#endif
   
  boost::asio::thread_pool pool2(number_of_sockets); 
  for(size_t j = 0; j < number_of_sockets; ++j)
  {
	   boost::asio::post(pool2, std::bind(accept_conncections_from_Pb,  std::ref(io_context), std::ref(sockets_), ports[j],  j));
  } 
  pool2.join();

#endif
 
 
	generate_random_distinguished_points(party);
  
 AES_KEY aeskey;	  
 blinds =	(DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));	
 updated_blinds =	(DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));

	size_t * where_to_write = new size_t[number_of_writes];
	size_t * where_to_read_dependent = new size_t[number_of_dep_reads];
	size_t * where_to_read_independent = new size_t[number_of_ind_reads];
   
   	
	auto start_total = std::chrono::steady_clock::now();	

	setup(DB, updated_DB, blinded_DB_recv, blinds, updated_blinds, updated_blinded_DB_recv, db_nitems, party);

	for(size_t i = 0; i < number_of_accesses; ++i)
	{
		for(size_t j = 0; j < number_of_writes; ++j)
		{
			where_to_write[j] = j + 4;
		}	   
		for(size_t j = 0; j < number_of_dep_reads; ++j)
		{
			where_to_read_dependent[j] = j + 4;
		}
		for(size_t j = 0; j < number_of_ind_reads; ++j)
		{
			where_to_read_independent[j] = j + 4;
		}			
		for(size_t j = 0; j < db_nitems; ++j) 
		{
			blinds[j] = 0;
			updated_blinds[j] = blinds[j];
			updated_blinded_DB_recv[j] = blinded_DB_recv[j];
		}

		int64_t ri, ri_recv;
		arc4random_buf(&ri, sizeof(ri));
		ri = ri % db_nitems;
		boost::asio::write(sockets_[0], boost::asio::buffer(&ri, sizeof(ri)));
		boost::asio::read(sockets_[0], boost::asio::buffer(&ri_recv, sizeof(ri_recv)));
		
	std::cout << "ri = " << ri << std::endl;
 if(party) ri = 2 - ri_recv;

		#ifdef VERBOSE
			boost::asio::write(sockets_[0], boost::asio::buffer(&ri, sizeof(ri)));
			boost::asio::read(sockets_[0], boost::asio::buffer(&ri_recv, sizeof(ri_recv)));		
			int64_t ri_reconstruction = ri + ri_recv;
			std::cout << "ri_reconstruction = " << ri_reconstruction << std::endl;
		#endif

		#ifdef ThreeParty
			DuORAM_Write * WritePb_     = new DuORAM_Write[number_of_writes];
			DuORAM_Write * WritePb_recv = new DuORAM_Write[number_of_writes];
			DB_t * read_out = new DB_t[number_of_writes];
			DB_t * Gamma =  new DB_t[number_of_writes];
		#endif

		DB_t * CW = new DB_t[number_of_writes];  
		DB_t * update_message = new DB_t[number_of_writes];

		auto start_writes = std::chrono::steady_clock::now();

		#ifdef ThreeParty
		for(size_t w = 0; w < number_of_writes; ++w)
		{
				DB_t FCW_read = 0;

				/*The definition of read_final_correction_word appears in duoram-utils.h*/ 	
				read_final_correction_word(party, FCW_read);
				
				#ifdef VERBOSE
					std::cout << "FCW_read (from) = " << FCW_read << std::endl;	
				#endif

				DB_t alpha0 = -FCW_read;
				WritePb_[w].shift = where_to_write[w] -ri;
				WritePb_[w].CW = distinguised_value[0];
				boost::asio::write(sockets_2[0], boost::asio::buffer(&WritePb_[w],   sizeof(DuORAM_Write)));
				read(sockets_2[1], boost::asio::buffer(&Gamma[w],  sizeof(DB_t))); 
				boost::asio::write(sockets_[0],  boost::asio::buffer(&WritePb_[w],    sizeof(DuORAM_Write)));
				boost::asio::read(sockets_[0],   boost::asio::buffer(&WritePb_recv[w], sizeof(DuORAM_Write)));

				/*The definition of read_flags_for_writing appears in duoram-utils.h*/ 	
				read_flags_for_writing(party, db_nitems);
	
				rotate_by_[w] = WritePb_[w].shift + WritePb_recv[w].shift;		 			

				#ifdef VERBOSE 		
					std::cout << "print database (prints the non-zero database entries): " << std::endl;
					/* The definition of reconstruct_database appears in duoram-utils.h */							
					reconstruct_database(sockets_[0],  DB, db_nitems);
				#endif

				for(size_t j = 0; j < db_nitems; ++j) reading_temp[j] = DB[j] + updated_blinded_DB_recv[j];
				
				/*The definition of dot_product_with_bool appears in duoram-write.h*/ 
				if(!party) read_out[w] 		=   dot_product_with_bool(reading_temp,   writing_b, db_nitems, rotate_by_[w]) + 
												dot_product_with_bool(updated_blinds, writing_b, db_nitems, rotate_by_[w]) - 
												dot_product_with_bool(updated_blinds, writing_c, db_nitems, rotate_by_[w]) + Gamma[w];
							
				if(party)  read_out[w] 		=   dot_product_with_bool(reading_temp,   writing_c, db_nitems, rotate_by_[w]) + 
												dot_product_with_bool(updated_blinds, writing_c, db_nitems, rotate_by_[w])	- 
												dot_product_with_bool(updated_blinds, writing_d, db_nitems, rotate_by_[w]) + Gamma[w];
							
				#ifdef VERBOSE
					std::cout << "read_out[" << w << "] = " << read_out[w] << std::endl;
				#endif

				#ifdef VERBOSE
				std::cout << "reconstructing the output: " << print_reconstruction(sockets_[0], read_out[w]) << "\n";	
				#endif
			
				distinguised_value[0] = 80 * (1 + w);
				update_message[w] = distinguised_value[0] - read_out[w]   + alpha0;

					#ifdef VERBOSE
						std::cout << "The updated message shares is = " << update_message[w] << std::endl;
					#endif 
				boost::asio::write(sockets_2[2], boost::asio::buffer(&update_message[w],  sizeof(DB_t)));
				boost::asio::write(sockets_[2], boost::asio::buffer(&update_message[w], sizeof(DB_t)));
				boost::asio::read(sockets_[2], boost::asio::buffer(&CW[w], sizeof(DB_t)));
					
				CW[w] = CW[w] + update_message[w];

				#ifdef VERBOSE
					std::cout << "cw = " << CW[w] << std::endl;
				#endif 
			}
		#endif

		for(size_t w = 0; w < number_of_writes; ++w)
 	{	
			
			/*The definition of DuoramUpdate appears in duoram-write.h*/  	 
		  	DuoramUpdate(party,  db_nitems,  rotate_by_[w],  DB, updated_DB, writing_b, b,  CW[w],  update_message[w],  writing_c, writing_d, c, d);
			
		 	#ifdef DEBUG		
				 #ifdef ThreeParty
					/*The definition debug_ appears in duoram-utils.h*/  
					debug_(sockets_2[0], sockets_[0], db_nitems);
			 	#endif
		 	#endif
 	}  

   	auto end_writes = std::chrono::steady_clock::now();
   	std::chrono::duration<double> elapsed_seconds_writes = end_writes - start_writes;
   	printf("elapsed_seconds_writes = %f\n",elapsed_seconds_writes.count());
	
	#ifdef VERBOSE
		std::cout << "Reconstructing the database after doing " << number_of_writes << " writes " << std::endl;
		reconstruct_database(sockets_[0],  DB, db_nitems);
	#endif
	// WRITES END.

  #ifdef ThreeParty
		std::cout << std::endl << std::endl << "============== WRITES END ==============" << std::endl << std::endl;
 	auto start_ind_reads = std::chrono::steady_clock::now();
		size_t *  WritePb_ind_reads = new size_t[number_of_ind_reads];
		size_t *  WritePb_ind_reads_recv = new size_t[number_of_ind_reads];
		size_t *  rotate = new size_t[number_of_ind_reads];
	 	 
	 	
	 for(size_t r = 0; r < number_of_ind_reads; ++r) WritePb_ind_reads[r] = where_to_read_independent[r] -ri;

	 boost::asio::write(sockets_2[3], boost::asio::buffer(WritePb_ind_reads, number_of_ind_reads * sizeof(size_t)));
	 boost::asio::write(sockets_[3], boost::asio::buffer(WritePb_ind_reads, number_of_ind_reads * sizeof(size_t)));
	 boost::asio::read(sockets_[3], boost::asio::buffer(WritePb_ind_reads_recv, number_of_ind_reads * sizeof(size_t)));
	 	
		DB_t * Gamma_reads = new DB_t[number_of_ind_reads];
 	boost::asio::read(sockets_2[4], boost::asio::buffer(Gamma_reads, number_of_ind_reads * sizeof(DB_t)));

		for(size_t j = 0; j < number_of_ind_reads; ++j)
		{
			rotate[j] = WritePb_ind_reads[j] + WritePb_ind_reads_recv[j];
		}

		DB_t * read_out_independent_reads = new DB_t[number_of_ind_reads];
		for(size_t r = 0; r < number_of_ind_reads; ++r)
		{
	 		#ifdef VERBOSE
		    	std::cout << "rotate[r]" << rotate[r] << std::endl; 
		    	std::cout << "Gamma_reads[r] = " << Gamma_reads[r] << std::endl;
		    #endif
			
			read_out_independent_reads[r] = DuoramIndependentRead(party,  db_nitems,  ri, Gamma_reads,  rotate,  r);

		    #ifdef VERBOSE
		  		std::cout << "---> [duoram independent reads] " <<  print_reconstruction(sockets_[0], read_out_independent_reads[r]) << std::endl;
		    #endif
	    }

 auto end_ind_reads = std::chrono::steady_clock::now();
 std::chrono::duration<double> elapsed_seconds_ind_reads = end_ind_reads - start_ind_reads;
 printf("elapsed_seconds_ind_reads = %f\n",elapsed_seconds_ind_reads.count());
	std::cout << std::endl << std::endl << "============== INDEPENDENT READS END  ==============" << std::endl << std::endl;
	auto start_dep_reads = std::chrono::steady_clock::now();
	
 
	std::cout << std::endl << std::endl << "============== DEPENDENT READS START  ==============" << std::endl << std::endl;
	DB_t * read_out_dependent_reads = new DB_t[number_of_dep_reads];
	 
 for(size_t r = 0; r < number_of_dep_reads; ++r)
	{
	  read_out_dependent_reads[r] = DuoramRead(party, db_nitems, ri, where_to_read_dependent[r], sockets_2[5], sockets_2[6], sockets_[5]);
	  #ifdef VERBOSE
    std::cout << "dependent read (share) " << r << " -> " << read_out_dependent_reads[r] << std::endl;	
		 	std::cout << print_reconstruction(sockets_[0], read_out_dependent_reads[r]) << std::endl;
		  #endif
  }
		
	auto end_dep_reads = std::chrono::steady_clock::now();
	std::chrono::duration<double> elapsed_seconds_dep_reads = end_dep_reads - start_dep_reads;
	printf("elapsed_seconds_dep_reads = %f\n",elapsed_seconds_dep_reads.count()); 
	std::cout << std::endl << std::endl << "============== DEPENDENT READS END  ==============" << std::endl << std::endl;

	 
 #endif
}  

auto end_total = std::chrono::steady_clock::now();
std::chrono::duration<double> elapsed_seconds_total = end_total - start_total;
printf("elapsed_seconds_total = %f\n",elapsed_seconds_total.count());

free(reading_temp);

	return 0;
}
