#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;

size_t communication_cost_dep_read = 0;
size_t communication_cost_write    = 0;
#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"

double dependent_read_time = 0.0;
double write_time = 0.0;

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];
    }

   
    // 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)));
		
	 
    // if(party) ri = 2 - ri_recv;
     int64_t ri;
    read_rand_indx(party, ri);

   
    #ifdef DEBUG
    int64_t  ri_recv;
    boost::asio::write(sockets_[0], boost::asio::buffer(&ri, sizeof(ri)));
    boost::asio::read(sockets_[0], boost::asio::buffer(&ri_recv, sizeof(ri_recv)));
    ri_recv = ri_recv + ri;
    std::cout << "ri_recv = "  << ri_recv % db_nitems << std::endl;
    #endif
    #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

     DuORAM_Write * WritePb_     = new DuORAM_Write[number_of_writes];
     DuORAM_Write * WritePb_recv = new DuORAM_Write[number_of_writes];
    #ifdef ThreeParty
     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)));
      communication_cost_write += 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)));
      communication_cost_write += 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)));
      communication_cost_write += 2 * 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 
     }
    #else // 2-party case
    for(size_t w = 0; w < number_of_writes; ++w)
    {
      WritePb_[w].shift = where_to_write[w] -ri;
      WritePb_[w].CW = distinguised_value[0];
      boost::asio::write(sockets_[0],  boost::asio::buffer(&WritePb_[w],    sizeof(DuORAM_Write)));
      communication_cost_write += sizeof(DuORAM_Write);
      boost::asio::read(sockets_[0],   boost::asio::buffer(&WritePb_recv[w], sizeof(DuORAM_Write)));
      rotate_by_[w] = WritePb_[w].shift + WritePb_recv[w].shift;
    }
    #endif
    delete[] WritePb_;
    delete[] WritePb_recv;
    delete[] where_to_write;

    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
    }  

    delete[] rotate_by_;
    delete[] CW;
    delete[] update_message;

   auto end_writes = std::chrono::steady_clock::now();
   std::chrono::duration<double> elapsed_seconds_writes = end_writes - start_writes;
   write_time = elapsed_seconds_writes.count();
   // printf("elapsed_seconds_writes = %f\n",elapsed_seconds_writes.count());
   // std::cout << "communication_cost_writes = " << communication_cost_write << std::endl; 
   #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
    #ifdef VERBOSE
     std::cout << std::endl << std::endl << "============== WRITES END ==============" << std::endl << std::endl;
    #endif
    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;
    delete[] where_to_read_independent;
    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
    }
    delete[] rotate;
    delete[] Gamma_reads;
    delete[] WritePb_ind_reads_recv;
    delete[] WritePb_ind_reads;
    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());

    #ifdef VERBOSE
    std::cout << std::endl << std::endl << "============== INDEPENDENT READS END  ==============" << std::endl << std::endl;
    #endif

    auto start_dep_reads = std::chrono::steady_clock::now();
    #ifdef VERBOSE
    std::cout << std::endl << std::endl << "============== DEPENDENT READS START  ==============" << std::endl << std::endl;
    #endif
    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]);
      communication_cost_dep_read += sizeof(DB_t);
      communication_cost_dep_read += sizeof(DB_t);
      #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
    }
    delete[] where_to_read_dependent;
    auto end_dep_reads = std::chrono::steady_clock::now();
    std::chrono::duration<double> elapsed_seconds_dep_reads = end_dep_reads - start_dep_reads;
    dependent_read_time = elapsed_seconds_dep_reads.count();
    // printf("elapsed_seconds_dep_reads = %f\n",elapsed_seconds_dep_reads.count()); 
    // std::cout << "communication_cost_dep_read = " << communication_cost_dep_read << std::endl;
    #ifdef VERBOSE
     std::cout << std::endl << std::endl << "============== DEPENDENT READS END  ==============" << std::endl << std::endl;
    #endif
   #endif
    
    free(reading_b);
    free(reading_c);
    free(reading_d);
    free(writing_b);
    free(writing_c);
    free(writing_d);
    free(reading_temp);
    free(DB);
    free(updated_DB);
    free(blinded_DB);
    free(blinded_DB_recv);
    free(updated_blinded_DB_recv);
    free(b);
    free(c);
    free(d);
    free(blinds);
    free(updated_blinds);
    
    #ifdef ThreeParty
     delete[] read_out;
     delete[] Gamma;
    #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());

std::cout << "write_time = " << write_time << std::endl;
std::cout << "communication_cost_writes = " << communication_cost_write << " bytes" << std::endl; 

#ifdef ThreeParty
 std::cout << "dependent_read_time =  " << dependent_read_time << std::endl;
 std::cout << "communication_cost_dep_read = " << communication_cost_dep_read << " bytes" << std::endl; 
 std::cout << "interleaved_time =  " << dependent_read_time + write_time  << std::endl;
 std::cout << "communication_cost_interleaved = " << (communication_cost_dep_read + communication_cost_write) << " bytes" << std::endl; 
#endif
free(reading_temp);

	return 0;
}