#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 "prg.h"
#include "prg_aes_impl.h"

#include <iostream>
 
#include "block.h"
#include <chrono>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <fstream> 
#include <../boost/asio.hpp>
#include <mutex>
#include <../boost/lexical_cast.hpp>
#include <memory>
#include <../boost/asio.hpp>
 
using boost::asio::ip::tcp;
//#include "network.h"

#include "duoram-utils.h"
 
using boost::asio::ip::tcp;

 
 
 

DB_t * X0;
DB_t * X1;
 
DB_t alpha;

 
typedef __m128i leaf_t;
typedef __m128i node_t;
__m128i seed0_, seed1_, seed2_;
  
const size_t target = 22;
const leaf_t val =   _mm_set1_epi8(0x01);

 using socket_t = boost::asio::ip::tcp::socket;
 


 using namespace std;
 

 
void quick_fix_send_ones_share(tcp::socket& out0, tcp::socket& out1)
{

  DB_t one = 47920;//0xffffffffffffffff;
  
  DB_t S0, S1;
  arc4random_buf(&S0, sizeof(DB_t));
  S1 = - S0 + one;
  // out0 << S0; 
  // out1 << S1;
  
  write(out0, boost::asio::buffer(&S0, sizeof(S0)));
  write(out1, boost::asio::buffer(&S1, sizeof(S1)));
    


  //arc4random_buf(&S0, sizeof(DB_t));
  S1 = - S0 + one;  
  // out0 << S0; 
  // out1 << S1;
 
  write(out0, boost::asio::buffer(&S0, sizeof(S0)));
  write(out1, boost::asio::buffer(&S1, sizeof(S1)));
 
 


  //arc4random_buf(&S0, sizeof(DB_t));
  S1 = - S0 + one;
  // out0 << S0; 
  // out1 << S1;
  write(out0, boost::asio::buffer(&S0, sizeof(S0)));
  write(out1, boost::asio::buffer(&S1, sizeof(S1)));
 
}




 
 

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);
  // sockets_0.emplace_back(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]);;

  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;
    
   arc4random_buf(&seed0_, sizeof(__m128i));
   arc4random_buf(&seed1_, sizeof(__m128i));
   arc4random_buf(&seed2_, sizeof(__m128i));
      
   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();

    //  int const in1_w_ { open( "party0_write_c", O_RDONLY ) };
    //  size_t r = read(in1_w_, c,  db_nitems * sizeof(c[0]));
    //  if(r < 0) {
    //  perror("Read error");
    //  close(in1_w_);
    //  return 1;
    //  }
    //  int const in2_w_ { open( "party1_write_d", O_RDONLY ) };
    //  r = read(in2_w_, d,  db_nitems * sizeof(d[0]));  
    //  if(r < 0) {
    //  perror("Read error");
    //  close(in2_w_);
    //  return 1;
    //  }

    //  int const in1_w { open( "party0_write_flags_c", O_RDONLY ) };
    //  r = read(in1_w, writing_c,  db_nitems * sizeof(writing_c[0]));
    //  if(r < 0) {
    //  perror("Read error");
    //  close(in1_w);
    //  return 1;
    //  }
    //  int const in2_w { open( "party1_write_flags_d", O_RDONLY ) };
    //  r = read(in2_w, writing_d,  db_nitems * sizeof(writing_d[0])); 
    //  if(r < 0) {
    //  perror("Read error");
    //  close(in2_w);
    //  return 1;
    //  }

    int const in2 { open( "party1_read_flags_d", O_RDONLY ) };
    size_t r = read(in2, reading_d,  db_nitems * sizeof(reading_d[0])); 
     if(r < 0) {
     perror("Read error");
     close(in2);
     return 1;
     }

     int const in2_ { open( "party0_read_flags_c", O_RDONLY ) };
     r = read(in2_, reading_c,  db_nitems * sizeof(reading_c[0])); 
     if(r < 0) {
     perror("Read error");
     close(in2_);
     return 1;
     }

  //quick_fix_send_ones_share(sockets_0[0], sockets_1[0]);

  block<__m128i> seed0, seed1;

  arc4random_buf(&seed0, sizeof(block<__m128i>));
  arc4random_buf(&seed1, sizeof(block<__m128i>));
  
  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));
    
 

  boost::asio::write(sockets_0[0], boost::asio::buffer(&seed0, sizeof(seed0))); 
  boost::asio::write(sockets_1[0], boost::asio::buffer(&seed1, sizeof(seed1)));
 
  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 w = 0; w < number_of_ind_writes; ++w)
    {
 

     int const in1_w { open( "party0_write_flags_c", O_RDONLY ) };
     size_t r = read(in1_w, writing_c,  db_nitems * sizeof(writing_c[0]));
     if(r < 0) {
     perror("Read error");
     close(in1_w);
     return 1;
     }
     
     int const in2_w { open( "party1_write_flags_d", O_RDONLY ) };
     r = read(in2_w, writing_d,  db_nitems * sizeof(writing_d[0])); 
     if(r < 0) {
     perror("Read error");
     close(in2_w);
     return 1;
     }
               int const in1_w_ { open( "party0_write_c", O_RDONLY ) };
           r = read(in1_w_, c,  db_nitems * sizeof(c[0]));
          if(r < 0) {
          perror("Read error");
          close(in1_w_);
          return 1;
          }
          int const in2_w_ { open( "party1_write_d", O_RDONLY ) };
          r = read(in2_w_, d,  db_nitems * sizeof(d[0]));  
          if(r < 0) {
          perror("Read error");
          close(in2_w_);
          return 1;
          }
        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)));

       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
    }



    
    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) rotate[r] = reads_shift_from_P0[r] + reads_shift_from_P1[r];

    for(size_t r = 0; r < number_of_ind_reads; ++r)
    {
    //  // rotate_all(reading_b, reading_b_rotated, reading_c, reading_c_rotated, reading_d, reading_d_rotated, rotate[r], db_nitems);
      
      #ifdef DEBUG
        std::cout << "rotate[r] " << rotate[r] << std::endl;
      #endif

      auto [Gamma0, Gamma1] =   generate_cancelation_terms(reading_c, reading_d, db_nitems, rotate[r]);
     
      Gamma0_reads[r] = Gamma0;
      Gamma1_reads[r] = Gamma1;
     }

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


     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;
        //rotate_all(reading_b, reading_b_rotated, reading_c, reading_c_rotated, reading_d, reading_d_rotated, rotate, 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)));
     }


  return 0;
}