#include "node.hpp"

#include "cmix.h"
#include "bignum.h"

#include "logging.hpp"

#include <numeric>

using namespace boost::asio::ip;

Node::Node(ListenSettings const& listen_settings, NodeNetworkSettings network_settings)
: io_service()
, timer(io_service)
, ssl_ctx(std::make_shared<boost::asio::ssl::context>(boost::asio::ssl::context::sslv23))
, server(io_service, listen_settings, ssl_ctx, [this](std::unique_ptr<boost::asio::ssl::stream<boost::asio::ip::tcp::socket>>&& socket, std::shared_ptr<boost::asio::ssl::context> ctx){accept_handler(std::move(socket), ctx);})
, clients()
, data()
, messages()
, network_settings(network_settings)
, prev_node(SSLReceiver(std::unique_ptr<boost::asio::ssl::stream<tcp::socket>>(new boost::asio::ssl::stream<boost::asio::ip::tcp::socket>(io_service, *ssl_ctx))))
, next_node(SSLSender(std::unique_ptr<boost::asio::ssl::stream<tcp::socket>>(new boost::asio::ssl::stream<boost::asio::ip::tcp::socket>(io_service, *ssl_ctx))))
, cmix_ctx(initialize_cmix_context(get_implementation()))
, shutting_down(false)
{
	initialize_keypair(&cmix_ctx);
	
	std::string x = to_string(cmix_ctx.keypair.sec, cmix_ctx);
	std::string y = to_string(cmix_ctx.keypair.pub, cmix_ctx);
	{
		std::stringstream ss;
		ss << "sec: ";
		for(auto&& c : x) {
			ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
		}
		BOOST_LOG_TRIVIAL(trace) << ss.str();
	}
	
	{
		std::stringstream ss;
		ss << "pub: ";
		for(auto&& c : y) {
			ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
		}
		BOOST_LOG_TRIVIAL(trace) << ss.str();
	} 
	
	GOOGLE_PROTOBUF_VERIFY_VERSION;

	if(network_settings.is_first) {
		connect_to_next_node();
	}
}

Node::~Node() {
	deinitialize(&cmix_ctx);
}

void Node::run() {
	io_service.run();
}

void Node::accept_handler(std::unique_ptr<boost::asio::ssl::stream<boost::asio::ip::tcp::socket>>&& socket, std::shared_ptr<boost::asio::ssl::context> ctx)
{
	Purgatory::iterator it = purgatory.emplace(purgatory.end(), std::move(socket));
	purgatory.back().on_done(
		[this, it]() {
			purgatory.erase(it);
		}
	);

	it->async_receive([this, it](cmix_proto::CMixMessage message) {
		handle_message(it, message);
	});
}

void Node::connect_to_next_node()
{
	if(!network_settings.certdir.empty()) {
		ssl_ctx->add_verify_path(network_settings.certdir);
	}
	
	auto on_connect = [this](){
		BOOST_LOG_TRIVIAL(trace) << "Connected to next_node";
		next_node.async_send(cmix_proto::ImANode());
	};
	
	next_node.async_connect(network_settings.next_host, network_settings.next_port, on_connect);
}

void Node::start_initialisation() {
	cmix_proto::Initialization init;
	
	size_t len = get_group_element_array_size(&cmix_ctx);
	init.mutable_public_share()->resize(len);
	
	get_public_key(&cmix_ctx, &(*init.mutable_public_share())[0]);
	
	BOOST_LOG_TRIVIAL(trace) << "Sending intialization as first node";
	next_node.async_send(init);
}

void Node::handle_imanode(Purgatory::iterator handle, cmix_proto::ImANode const&) {
	handle->on_done([]{});
	prev_node = std::move(*handle);
	purgatory.erase(handle);
	
	if(network_settings.is_first) {
		start_initialisation();
	} else {
		connect_to_next_node();
	}
	
	prev_node.async_receive([this](cmix_proto::CMixMessage message){
		handle_node_message(message);
	});
}

void Node::handle_imaclient(Purgatory::iterator handle, cmix_proto::ImAClient const& c) {
	std::string client_id = c.id();
	clients.emplace(c.id(), decltype(clients)::mapped_type(std::move(*handle)));
	clients.at(c.id()).on_done([this, client_id]{
		clients.erase(client_id);
	});
	purgatory.erase(handle);
	clients.at(c.id()).async_send(cmix_proto::NodeReady());
	clients.at(c.id()).async_receive([this, client_id](cmix_proto::CMixMessage message) {
		handle_client_message(client_id, message);
	});
}

void Node::handle_message(Purgatory::iterator handle, cmix_proto::CMixMessage message)
{
	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kImanode: {
			BOOST_LOG_TRIVIAL(trace) << "Handling imanode";
			handle_imanode(handle, message.imanode());
			return;
		}
		case cmix_proto::CMixMessage::ContentsCase::kImaclient: {
			BOOST_LOG_TRIVIAL(trace) << "Handling imaclient";
			handle_imaclient(handle, message.imaclient());
			return;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_message: CMixMessage contains unknown contents.";
		}
	}
	
	handle->close();
	purgatory.erase(handle);
}

void Node::start_precomputation() {
	BOOST_LOG_TRIVIAL(trace) << "Starting precomputation for " << clients.size() << " clients.";
	index_map.clear();
	if(start_mix(&cmix_ctx, clients.size()) != no_error) {
		exit(-1);
	}
	
	unsigned int i = 0;
	for(auto&& pair : clients) {
		index_map[pair.first] = i++;
	}
	
	if(initialize_mix_randomness(&cmix_ctx) != no_error) {
		exit(-1);
	}
	
	cmix_proto::PrePre prepre;
	
	for(size_t i = 0; i < cmix_ctx.nr_participants; ++i) {
		prepre.add_m_er();
		prepre.add_r_er();
	
		size_t len = get_group_element_array_size(&cmix_ctx);
		
		prepre.mutable_r_er(i)->resize(len);
		prepre.mutable_m_er(i)->resize(len);
		
		{
			std::stringstream ss;
			ss << "r: ";
			std::string r = to_string(cmix_ctx.r[i], cmix_ctx);
			for(auto&& c : r) {
				ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
			}
			BOOST_LOG_TRIVIAL(trace) << ss.str();
		} 
		
		if(encrypt_r(&cmix_ctx, &(*prepre.mutable_r_er(i))[0], &(*prepre.mutable_m_er(i))[0], i) != no_error) {
			exit(-1);
		}
	}
	
	BOOST_LOG_TRIVIAL(trace) << "Sending prepre message: " << prepre.ShortDebugString();
	next_node.async_send(prepre);
}

void Node::start_realtime_phase() {
	
	cmix_proto::RealPre realpre;
	
	size_t len = get_group_element_array_size(&cmix_ctx);
	for(auto&& pair : index_map) {
		realpre.add_h();
		realpre.add_m();
		realpre.mutable_m(pair.second)->resize(len);
		std::decay<decltype(pair)>::type::first_type const& handle = pair.first;
		std::decay<decltype(pair)>::type::second_type const& index = pair.second;
		
		auto& queue = messages[pair.first];
		if(queue.empty()) {
			std::vector<char> v(len);
			generate_random_message(&cmix_ctx, v.data());
			swap_k_for_r(&cmix_ctx, &(*realpre.mutable_m(index))[0], v.data(), data[handle].shared_value, index);
		} else {
			std::string x = to_string(data.at(handle).shared_value, cmix_ctx);
			{
				std::stringstream ss;
				ss << "shared_key: ";
				for(auto&& c : x) {
					ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
				}
				BOOST_LOG_TRIVIAL(trace) << ss.str();
			}
		
			BOOST_LOG_TRIVIAL(trace) << 
		
			swap_k_for_r(&cmix_ctx, &(*realpre.mutable_m(index))[0], queue.front().data(), data.at(handle).shared_value, index);
		}
		*realpre.mutable_h(index) = handle;
	}
	next_node.async_send(realpre);
}