#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 <../boost/asio.hpp>
#include <../boost/lexical_cast.hpp>
#include <iostream> 
#include <chrono>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <fstream>
#include <mutex>
#include "bitutils.h"
#include "block.h"
#include "prg_aes_impl.h" 
#include "filesio.h"

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

#include "mpc.h"
#include "dpfgen.h"
#include "share-conversion.h"

void mpc_gen(const size_t depth, AES_KEY& prgkey, const size_t db_nitems, const size_t n_threads, std::vector<socket_t>& sockets0, std::vector<socket_t>& sockets1, 
             __m128i** output0, int8_t ** flags0, __m128i** output1, int8_t ** flags1, dpfP2 * dpf_instance0 , dpfP2 * dpf_instance1, size_t ind,
             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;

      gamma0 = (bit1 == 1) ? rand0 : _mm_setzero_si128();
      gamma1 = (bit0 == 1) ? rand1 : _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

  } 
  
  boost::asio::read(sockets0[socket_no+1], boost::asio::buffer(&dpf_instance0[ind], sizeof(dpfP2)));
  boost::asio::read(sockets1[socket_no+1], boost::asio::buffer(&dpf_instance1[ind], sizeof(dpfP2)));
 

  evaluate_dpfs(db_nitems,  dpf_instance0[ind],   prgkey,  0, db_nitems-1,	output0[ind],  flags0[ind],  false,  ind);
  evaluate_dpfs(db_nitems,  dpf_instance1[ind],   prgkey,  0, db_nitems-1, output1[ind],  flags1[ind],  true ,  ind);


  // P2_write_evalfull_outs_into_a_file(false, 0, db_nitems,  flags0[0], 	output0[0]);
  // P2_write_evalfull_outs_into_a_file(true,  0, db_nitems,  flags1[0], 	output1[0]);
 #ifdef DEBUG
    for(size_t j = 0; j < db_nitems; ++j)
    {
      std::cout << j << "-> "  << (int) flags0[0][j] << " <-> " << (int) flags1[0][j] << std::endl;
      std::cout << j << "-> " << output0[0][j][0] << " <-> " << output1[0][j][0] << std::endl << std::endl;
    }
  #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 db_nitems = 1ULL << atoi(argv[4]);
  const size_t depth = std::ceil(std::log2(db_nitems));
    
  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);  

	__m128i ** output0 = (__m128i ** ) malloc(sizeof(__m128i *) * n_threads);
	int8_t  ** flags0  = (int8_t ** ) malloc(sizeof(uint8_t *) * n_threads);
	 
	for(size_t j = 0; j < n_threads; ++j)
	{
		output0[j] = (__m128i *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(__m128i));
		flags0[j]  = (int8_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(uint8_t));
	}


  
	__m128i ** output1 = (__m128i ** ) malloc(sizeof(__m128i *) * n_threads);
	int8_t  ** flags1  = (int8_t ** ) malloc(sizeof(uint8_t *) * n_threads);
	 
	for(size_t j = 0; j < n_threads; ++j)
	{
		output1[j] = (__m128i *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(__m128i));
		flags1[j]  = (int8_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(uint8_t));
	}
 
 dpfP2 * dpf_instance0 = (dpfP2 * ) malloc (sizeof(dpfP2) * n_threads);
 dpfP2 * dpf_instance1 = (dpfP2 * ) malloc (sizeof(dpfP2) * n_threads);
 
 for(size_t j = 0; j < n_threads; ++j)
 {
   boost::asio::post(pool, std::bind(mpc_gen,  std::ref(depth), std::ref(aeskey), db_nitems, n_threads,  std::ref(sockets0), std::ref(sockets1), 
                                       output0, flags0,  output1, flags1,  std::ref(dpf_instance0), std::ref(dpf_instance1), j, 5 * j));    
 }

  pool.join();


  boost::asio::thread_pool pool3(n_threads); 
  
  int64_t ** leaves0    = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads);
	 int64_t ** leafbits0  = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads); 
  int64_t ** leaves1    = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads);
	 int64_t ** leafbits1  = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads); 
	
  for(size_t j = 0; j < n_threads; ++j)
	 {
		 leaves0[j]    = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
		 leafbits0[j]  = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
   leaves1[j]    = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
		 leafbits1[j]  = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
	 }

 /* The function convert_sharesP2 appears in share-conversion.h */
  for(size_t j = 0; j < n_threads; ++j)
  {
   boost::asio::post(pool3, std::bind(convert_sharesP2, db_nitems,  output0, flags0,  output1, flags1, leaves0, leafbits0, leaves1, leafbits1,  std::ref(sockets0), std::ref(sockets1), j, j));    
  }

  pool3.join(); 
  
  /* The function P2_xor_to_additive appears in share-conversion.h */
  boost::asio::thread_pool pool4(n_threads); 
  for(size_t j = 0; j < n_threads; ++j)
  {
     boost::asio::post(pool4,  std::bind(P2_xor_to_additive, std::ref(sockets0[j]), std::ref(sockets1[j]), j));
  }
  pool4.join();



  for(size_t i = 0; i < n_threads; ++i)
  {
      P2_write_evalfull_outs_into_a_file(false, i, db_nitems,  flags0[i], 	leaves0[i]);
      P2_write_evalfull_outs_into_a_file(true,  i, db_nitems,  flags1[i], 	leaves1[i]);
  }

  #ifdef DEBUG
    for(size_t ind = 0; ind < n_threads; ++ind)
    {
     for(size_t j = 0; j < db_nitems; ++j)
     {
      if(flags0[ind][j] + flags1[ind][j] != 0)
      {
       std::cout << j << "-> "  << (int) (flags0[ind][j] + flags1[ind][j]) << " = " << (int) (flags0[ind][j])  << " + " << (int) (flags1[ind][j]) << std::endl;
       std::cout << j << "-> " << output0[ind][j][0] << " <-> " << output1[ind][j][0] << std::endl << std::endl;
      }
     }
   }
  #endif
  
  return 0;
}