#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 <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <fstream> 
#include <mutex>
#include <../boost/asio.hpp>
#include <../boost/lexical_cast.hpp>
 
using boost::asio::ip::tcp;

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


DB_t * X0;
DB_t * X1; 
DB_t alpha;

typedef __m128i leaf_t;
typedef __m128i node_t;
  
using socket_t = boost::asio::ip::tcp::socket;

using namespace std;
 
size_t communication_cost_dep_read = 0;
size_t communication_cost_write    = 0;
size_t communication_cost_ind_read = 0;

double dependent_read_time = 0.0;
double write_time = 0.0;

auto generate_cancelation_terms(int8_t c0[], int8_t d1[], size_t db_nitems, size_t rotate_by = 0)
{
   DB_t Z0 =  dot_product_with_bool(X0, d1, db_nitems, rotate_by);
   DB_t Z1 =  dot_product_with_bool(X1, c0, db_nitems, rotate_by);

   DB_t cancelation_term0 = -Z0 ;
   DB_t cancelation_term1 = -Z1 ; 

   return std::make_pair(cancelation_term0, cancelation_term1);
}
 
 
 
void refresh_blinds(int8_t writing_c[], int8_t writing_d[], 
                    DB_t c[], DB_t d[],
                    DB_t finalcw0, DB_t finalcw1, size_t db_nitems, size_t rotate_by = 0)
{
  for(size_t j = 0; j < db_nitems; ++j)
  {
    X0[j] = X0[j] - c[(j + rotate_by) % db_nitems] - ((writing_c[(j + rotate_by) % db_nitems]) * finalcw0);//   tmp0; //c0
    X1[j] = X1[j] - d[(j + rotate_by) % db_nitems] - ((writing_d[(j + rotate_by) % db_nitems]) * finalcw1); //c2
  }
}
 


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

int main(int argc, char* argv[])
{ 
  size_t expo =  atoi(argv[3]);
  size_t db_nitems            = 1ULL << expo;
  size_t number_of_ind_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]);
  
  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));
   
  boost::asio::io_context io_context;
    
  AES_KEY aeskey;
       
  const size_t number_of_sockets = 40;

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

   std::vector<socket_t> sockets_0;
   std::vector<socket_t> sockets_1;
   sockets_0.reserve(number_of_sockets + 1);
   sockets_1.reserve(number_of_sockets + 1);

   boost::asio::thread_pool pool2(number_of_sockets * 2); 

   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_1), ports2_1[j],  j));
   }


   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_0), ports2_0[j],  j));
   }

   pool2.join();

   X0 = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
   X1 = (DB_t *) std::aligned_alloc(sizeof(__m256i), db_nitems * sizeof(DB_t));
    
   for(size_t j = 0; j < db_nitems; ++j)
   {
    X0[j] = 0;
    X1[j] = 0; 
   }

  
   DuORAM_Write * WriteP0_recv = new DuORAM_Write[number_of_ind_writes];
   DuORAM_Write * WriteP1_recv = new DuORAM_Write[number_of_ind_writes];

   DB_t * Gamma0_ = new DB_t[number_of_ind_writes];
   DB_t * Gamma1_ = new DB_t[number_of_ind_writes];
   DB_t * update_message0 = new DB_t[number_of_ind_writes];
   DB_t * update_message1 = new DB_t[number_of_ind_writes];
   DB_t * FCW = new DB_t[number_of_ind_writes];
 
   for(size_t i = 0; i < number_of_accesses; ++i)
   {
    auto start_writes = std::chrono::steady_clock::now();
     for(size_t w = 0; w < number_of_ind_writes; ++w)
     {
       read_flags_for_refreshing(db_nitems);

       size_t rotate_by_ = WriteP0_recv[w].shift + WriteP1_recv[w].shift;
     
       auto [Gamma0, Gamma1] = generate_cancelation_terms(writing_c, writing_d, db_nitems, rotate_by_);
       Gamma0_[w] = Gamma0;
       Gamma1_[w] = Gamma1; 
 
       boost::asio::read(sockets_0[0], boost::asio::buffer(&WriteP0_recv[w],  sizeof(DuORAM_Write)));
       boost::asio::read(sockets_1[0], boost::asio::buffer(&WriteP1_recv[w],  sizeof(DuORAM_Write)));
       boost::asio::write(sockets_0[1], boost::asio::buffer(&Gamma0_[w], sizeof(DB_t))); 
       boost::asio::write(sockets_1[1], boost::asio::buffer(&Gamma1_[w], sizeof(DB_t)));
       communication_cost_write += sizeof(DB_t);
       communication_cost_write += sizeof(DB_t);
       boost::asio::read(sockets_0[2], boost::asio::buffer(&update_message0[w],  sizeof(DB_t))); 
       boost::asio::read(sockets_1[2], boost::asio::buffer(&update_message1[w],  sizeof(DB_t)));
       
       FCW[w] = update_message0[w] + update_message1[w];      
     }

     for(size_t w = 0; w < number_of_ind_writes; ++w)
     {
       size_t rotate_by_ = WriteP0_recv[w].shift + WriteP1_recv[w].shift;
       refresh_blinds(writing_c, writing_d, c, d, FCW[w], FCW[w], db_nitems, -rotate_by_);
   
       #ifdef DEBUG
       for(size_t j = 0; j < db_nitems; ++j)
       {
         boost::asio::write(sockets_0[0], boost::asio::buffer(&X0[j], sizeof(X0[j]))); //DEBUG
         boost::asio::write(sockets_1[0], boost::asio::buffer(&X1[j], sizeof(X1[j]))); //DUBUG
       }
       #endif
     }

   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();

    auto start_ind_reads = std::chrono::steady_clock::now();
     DB_t * Gamma0_reads = new DB_t[number_of_ind_reads];
     DB_t * Gamma1_reads = new DB_t[number_of_ind_reads];

     
     size_t * reads_shift_from_P0 = new size_t[number_of_ind_reads];
     size_t * reads_shift_from_P1 = new size_t[number_of_ind_reads];
     size_t *  rotate              = new size_t[number_of_ind_reads];
     
     boost::asio::read(sockets_0[3], boost::asio::buffer(reads_shift_from_P0, number_of_ind_reads * sizeof(size_t))); 
     boost::asio::read(sockets_1[3], boost::asio::buffer(reads_shift_from_P1, number_of_ind_reads * sizeof(size_t)));

     for(size_t r = 0; r < number_of_ind_reads; ++r)
     {  
       #ifdef VERBOSE
         std::cout << "rotate[r] " << rotate[r] << std::endl;
       #endif
       rotate[r] = reads_shift_from_P0[r] + reads_shift_from_P1[r];
       read_flags_for_generating_cancellation_terms(db_nitems); 
       auto [Gamma0, Gamma1] =   generate_cancelation_terms(reading_c, reading_d, db_nitems, rotate[r]);     
       Gamma0_reads[r] = Gamma0;
       Gamma1_reads[r] = Gamma1;
     }

     delete[] reads_shift_from_P0;
     delete[] reads_shift_from_P1;
     delete[] rotate;

     boost::asio::write(sockets_0[4], boost::asio::buffer(Gamma0_reads, number_of_ind_reads * sizeof(DB_t)));  
     boost::asio::write(sockets_1[4], boost::asio::buffer(Gamma1_reads, number_of_ind_reads * sizeof(DB_t)));
     delete[] Gamma0_reads;
     delete[] Gamma1_reads;
     communication_cost_ind_read += number_of_ind_reads * sizeof(DB_t);
     communication_cost_ind_read += number_of_ind_reads * sizeof(DB_t);
    auto end_ind_reads = std::chrono::steady_clock::now();
    std::chrono::duration<double> elapsed_seconds_ind_reads = end_ind_reads - start_ind_reads;

    auto start_dep_reads = std::chrono::steady_clock::now();
     for(size_t d = 0; d < number_of_dep_reads; ++d)
     {  
       size_t shifts0, shifts1, rotate;
       boost::asio::read(sockets_0[5], boost::asio::buffer(&shifts0, sizeof(shifts0)));
       boost::asio::read(sockets_1[5], boost::asio::buffer(&shifts1, sizeof(shifts1)));
       rotate = shifts0 + shifts1;
         
       read_flags_for_generating_cancellation_terms(db_nitems);
         
       auto [Gamma0, Gamma1] =   generate_cancelation_terms(reading_c, reading_d, db_nitems, rotate);
       boost::asio::write(sockets_0[6], boost::asio::buffer(&Gamma0, sizeof(Gamma0)));  
       boost::asio::write(sockets_1[6], boost::asio::buffer(&Gamma1, sizeof(Gamma1)));
       communication_cost_dep_read += sizeof(Gamma0);
       communication_cost_dep_read += sizeof(Gamma1);
     }



    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();
  }

delete[] WriteP0_recv;
delete[] WriteP1_recv;

delete[] Gamma0_;
delete[] Gamma1_;
delete[] update_message0;
delete[] update_message1;
delete[] FCW;
free(X0);
free(X1);
free(b);
free(c);
free(d);
free(reading_b);
free(reading_c);
free(reading_d);
free(writing_b);
free(writing_c);
free(writing_d);

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

 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; 

  return 0;
}