#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 <boost/asio/thread_pool.hpp>
#include "bitutils.h"
#include "block.h"
#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>
 

using boost::asio::ip::tcp;

using namespace dpf;
 

typedef __m128i leaf_type;
 
 
typedef __m128i leaf_t;
typedef __m128i node_t;
 
 

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

size_t bits_per_leaf = std::is_same<leaf_t, bool>::value ? 1 : sizeof(leaf_t) * CHAR_BIT;
bool is_packed = (sizeof(leaf_t) < sizeof(node_t));
size_t leaves_per_node = is_packed ? sizeof(node_t) * CHAR_BIT / bits_per_leaf : 1;
size_t __depth(const size_t nitems) { return std::ceil(std::log2(std::ceil(static_cast<double>(nitems) / leaves_per_node))); }
 
 
void mpc_gen(const size_t depth, std::vector<socket_t>& sockets0, std::vector<socket_t>& sockets1, size_t socket_no = 0)
{

  for(size_t j = 0; j < depth; ++j)
  {
   
     __m128i rand0, rand1, gamma0, gamma1;

     arc4random_buf(&rand0, sizeof(__m128i));
     arc4random_buf(&rand1, sizeof(__m128i));
     uint8_t bit0, bit1; 
     bit0 = rand();
     bit0 = bit0 % 2;
     bit1 = rand();
     bit1 = bit1 %2;

     if(bit1 == 1)
     {
      gamma0 = rand0;
     }
     else
     {
      gamma0 = _mm_setzero_si128();
     }
     
     if(bit0 == 1)
     {
      gamma1 = rand1;
     }
     else
     {
      gamma1 = _mm_setzero_si128();
     }
    
    struct cw_construction
    {
      __m128i rand_b, gamma_b;
      uint8_t bit_b;
    };
    
    cw_construction computecw0, computecw1;
    
    computecw0.rand_b   = rand0;
    computecw0.gamma_b  = gamma0;
    computecw0.bit_b    = bit0;

    computecw1.rand_b   = rand1;
    computecw1.gamma_b  = gamma1;
    computecw1.bit_b    = bit1;
    
    boost::asio::write(sockets0[socket_no], boost::asio::buffer(&computecw0, sizeof(computecw0)));
    boost::asio::write(sockets1[socket_no], boost::asio::buffer(&computecw1, sizeof(computecw1)));

    #ifdef DEBUG
      boost::asio::write(sockets0[socket_no], boost::asio::buffer(&rand0, sizeof(rand0)));
      boost::asio::write(sockets0[socket_no], boost::asio::buffer(&gamma0, sizeof(gamma0)));
      boost::asio::write(sockets0[socket_no], boost::asio::buffer(&bit0, sizeof(bit0)));

      boost::asio::write(sockets1[socket_no], boost::asio::buffer(&rand1, sizeof(rand1)));
      boost::asio::write(sockets1[socket_no], boost::asio::buffer(&gamma1, sizeof(gamma1)));
      boost::asio::write(sockets1[socket_no], boost::asio::buffer(&bit1, sizeof(bit1)));
    #endif

  } 
  
}
 
void accept_conncections_from_Pb(boost::asio::io_context&io_context, std::vector<socket_t>& sockets0, int port, size_t j)
{
  tcp::acceptor acceptor2_(io_context, tcp::endpoint(tcp::v4(), port));
  tcp::socket s2(acceptor2_.accept());
  sockets0[j] = std::move(s2);
}
int main(int argc, char* argv[])
{ 
  
  AES_KEY aeskey;
  boost::asio::io_context io_context;
  tcp::resolver resolver(io_context);
   
  const std::string host0 = (argc < 2) ? "127.0.0.1" : argv[1];
  const std::string host1 = (argc < 3) ? "127.0.0.1" : argv[2];
   
  const size_t n_threads = atoi(argv[3]);
  const size_t number_of_sockets = 5 * n_threads;
       
   
   const size_t nitems2 = 1ULL << atoi(argv[4]);
   const size_t depth = std::ceil(std::log2(nitems2));

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

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


    std::vector<socket_t> sockets0;
    std::vector<socket_t> sockets1;
    sockets0.reserve(number_of_sockets + 1);
    sockets1.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(sockets1), 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(sockets0), ports2_0[j],  j));
    }

pool2.join();

boost::asio::thread_pool pool(n_threads); 

for(size_t j = 0; j < n_threads; ++j)
{

  boost::asio::post(pool, std::bind(mpc_gen,  std::ref(depth), std::ref(sockets0), std::ref(sockets1), 5 * j));
    
}

pool.join();

 
    return 0;
}