#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/lexical_cast.hpp>
#include <../boost/asio.hpp>
#include <fcntl.h>
#include <cstdlib>
#include <chrono>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fstream>
#include <future>
#include <mutex>


#include "bitutils.h"
#include "block.h" 
#include "prg_aes_impl.h"
 

using boost::asio::ip::tcp;
using socket_t = boost::asio::ip::tcp::socket;
 
typedef unsigned char byte_t;

typedef __m128i node_t;
constexpr size_t leaf_size = 1;
typedef __m128i leaf_type;

typedef std::array<leaf_type, leaf_size> leaf_t;

using namespace dpf;  // The namespace is found in bitutils.h

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

int main(int argc, char * argv[])
{ 	
	
	boost::asio::io_context 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];
	const size_t n_threads = atoi(argv[3]);
	const size_t expo = atoi(argv[4]);
	const size_t number_of_sockets = 5 * n_threads;

	std::vector<socket_t> socketsPb, socketsP2;
	std::vector<int> ports, ports2_1, ports2_0;
 
  	bool party;

	make_connections(party, host1, host2,  io_context, socketsPb, socketsP2, ports,  ports2_1, ports2_0, number_of_sockets);
 
	const size_t db_nitems = 1ULL << atoi(argv[4]);

 	auto start = std::chrono::steady_clock::now(); 
	
	uint8_t ** target_share_read = new uint8_t*[n_threads];

	generate_random_targets(target_share_read,  n_threads, party, expo);
	
	AES_KEY aeskey;

	__m128i * final_correction_word = (__m128i *) std::aligned_alloc(sizeof(__m256i), n_threads * sizeof(__m128i));
	__m128i ** output = (__m128i ** ) malloc(sizeof(__m128i *) * n_threads);
	int8_t  ** flags  = (int8_t ** ) malloc(sizeof(uint8_t *) * n_threads);
	 
	for(size_t j = 0; j < n_threads; ++j)
	{
		output[j] = (__m128i *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(__m128i));
		flags[j]  = (int8_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(uint8_t));
	}
	 
	boost::asio::thread_pool pool(n_threads);
	boost::asio::thread_pool pool_share_conversion(n_threads);
	
	int64_t ** leaves = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads);
	int64_t ** leafbits  = (int64_t ** ) malloc(sizeof(int64_t *) * n_threads);
	 
	for(size_t j = 0; j < n_threads; ++j)
	{
		leaves[j] = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
		leafbits[j]  = (int64_t *)std::aligned_alloc(sizeof(node_t), db_nitems * sizeof(int64_t));
	}
	
	// The following function call creates and evaluates DPFs at target_share_read[j] for j \in \{0, \ldots, n_threads}
	// the flag vectors are stored in flags
	// the leaves are stored in output
	// the final correctionword is stored in final_correction_word
	/* The function create_dpfs is defined dpf-gen.h*/
	for(size_t j = 0; j < n_threads; ++j)
	{
		boost::asio::post(pool,  std::bind(create_dpfs,  db_nitems,	std::ref(aeskey),  target_share_read[j],  std::ref(socketsPb), std::ref(socketsP2), 0, db_nitems-1, 
						  output[j],  flags[j], std::ref(final_correction_word[j]),  party, 5 * j, j));	 	  
	}
	pool.join();  

	#ifdef DEBUG
		for(size_t j = 0; j < n_threads; ++j)
		{
			for(size_t i = 0; i < db_nitems; ++i)
			{
				int8_t flags_reconstruction;
				boost::asio::write(socketsPb[0], boost::asio::buffer(&flags[j][i], sizeof(flags[j][i])));
				boost::asio::read(socketsPb[0], boost::asio::buffer(&flags_reconstruction, sizeof(flags_reconstruction)));
				flags_reconstruction -= flags[j][i];
				if(flags_reconstruction != 0) std::cout << i << " ---> " << (int) flags_reconstruction << std::endl;

				int64_t output_reconstruction;
				boost::asio::write(socketsPb[0], boost::asio::buffer(&output[j][i][0], sizeof(output[j][i][0])));
				boost::asio::read(socketsPb[0], boost::asio::buffer(&output_reconstruction, sizeof(output_reconstruction)));
				output_reconstruction -= output[j][i][0];
				if(output_reconstruction != 0) std::cout << i << "---> " << output_reconstruction << std::endl;
			}
			int64_t final_correction_word_reconstruction = 0;
			boost::asio::write(socketsPb[0], boost::asio::buffer(&final_correction_word[j][0], sizeof(final_correction_word[j][0])));
			boost::asio::read(socketsPb[0], boost::asio::buffer(&final_correction_word_reconstruction, sizeof(final_correction_word_reconstruction)));
			final_correction_word_reconstruction = final_correction_word_reconstruction + final_correction_word[j][0];
			std::cout << "final_correction_word_reconstruction = " << final_correction_word_reconstruction << std::endl << std::endl;
		}
	#endif
 

	for(size_t j = 0; j < n_threads; ++j)
	{
		boost::asio::post(pool_share_conversion,  std::bind(convert_shares, j, output, flags, n_threads, db_nitems, final_correction_word, 	leaves, leafbits, 
												  std::ref(socketsPb[j]), std::ref(socketsP2[j]), party ));	 	
	}
	
	pool_share_conversion.join();

 	xor_to_additive( party, target_share_read, socketsPb, socketsP2);

	auto end = std::chrono::steady_clock::now();
	std::chrono::duration<double> elapsed_seconds = end-start;
	std::cout << "time to generate and evaluate " << n_threads << " dpfs of size 2^" << atoi(argv[4]) << " is: " << elapsed_seconds.count() << "s\n";

	return 0;
}