#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()
, ctx(std::make_shared<boost::asio::ssl::context>(boost::asio::ssl::context::sslv23))
, server(io_service, listen_settings, 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()
, 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, *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, *ctx))))
, api(get_implementation())
, keypair(api.create_key_pair())
, network_key()
, precomputation_data()
, shutting_down(false)
{
	GOOGLE_PROTOBUF_VERIFY_VERSION;

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

Node::~Node() {
	api.free_key_pair(&keypair);
}

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()) {
		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;
	
	unsigned char* pub_key;
	size_t len;
	api.element_to_array(&pub_key, &len, keypair.pub);
	init.set_public_share(pub_key, len);
	free(pub_key);
	
	BOOST_LOG_TRIVIAL(trace) << "Sending intialization as first node";
	next_node.async_send(init);
}


void Node::handle_node_initialization(const cmix_proto::Initialization& init)
{
	if(network_settings.is_first) {
		network_key = api.array_to_element(init.public_share().c_str(), init.public_share().size(), true);
	
		cmix_proto::SecretKey sec;
		
		unsigned char* data;
		size_t len;
		
		api.element_to_array(&data, &len, network_key);
		sec.set_secret_key(data, len);
		
		api.free_buffer(data);
		
		next_node.async_send(sec);
	} else {
	
		char* buffer;
		size_t len;
		api.add_public_share(&buffer, &len, init.public_share().c_str(), init.public_share().size(), keypair.pub);
		
		cmix_proto::Initialization init;
		init.set_public_share(buffer, len);
		api.free_buffer(buffer);
		
		BOOST_LOG_TRIVIAL(trace) << "Sending intialization";
		next_node.async_send(init);
	}
}

void Node::handle_node_secretkey(cmix_proto::SecretKey const& secret)
{
	std::string share = secret.secret_key();
	
	network_key = api.array_to_element(secret.secret_key().c_str(), secret.secret_key().size(), true);
	
	if(network_settings.is_first) {
		start_precomputation();
	} else {
		next_node.async_send(secret);
	}
}

void Node::handle_node_prepre(cmix_proto::PrePre const& pre) {
	if(network_settings.is_first) {
		
		cmix_proto::PreMix premix;
		for(int i = 0; i < pre.m_er_size(); ++i) {
			GroupElement random_r = api.array_to_element(pre.r_er(i).data(), pre.r_er(i).size(), true);
			GroupElement message_r = api.array_to_element(pre.m_er(i).data(), pre.m_er(i).size(), true);
			
			GroupElement random_s;
			GroupElement message_s;
			
			api.encrypt(&random_s, &message_s, precomputation_data[precomputation_data[i].new_location].s, network_key);
			
			GroupElement random_pirs = api.multiply(random_r, random_s, true);
			GroupElement message_pirs = api.multiply(message_r, message_s, true);
			
			unsigned char* data;
			size_t len;
			
			api.element_to_array(&data, &len, random_pirs);
			premix.set_r_epirs(precomputation_data[i].new_location, data, len);
			api.free_buffer(data);
			
			api.element_to_array(&data, &len, message_pirs);
			premix.set_m_epirs(precomputation_data[i].new_location, data, len);
			api.free_buffer(data);
			
			api.free_group_element(random_r);
			api.free_group_element(message_r);
			api.free_group_element(random_s);
			api.free_group_element(message_s);
		}
		
		next_node.async_send(premix);
		
	} else {
		std::vector<size_t> permutation(clients.size());
		std::iota(permutation.begin(), permutation.end(), 0);
		//ToDo: generate something different than the ID permutation.
		
		cmix_proto::PrePre prepre;
		
		for(int i = 0; i < pre.m_er_size(); ++i) {
			GroupElement r = api.get_group_element(true);
			GroupElement s = api.get_group_element(true);
			
			GroupElement random_element;
			GroupElement message_element;
			
			api.encrypt(&random_element, &message_element, r, network_key);
			
			GroupElement other_random_element = api.array_to_element(pre.r_er(i).data(), pre.r_er(i).size(), true);
			GroupElement other_message_element = api.array_to_element(pre.m_er(i).data(), pre.m_er(i).size(), true);
			
			GroupElement new_random_element = api.multiply(random_element, other_random_element, true);
			GroupElement new_message_element = api.multiply(message_element, other_message_element, true);
			
			
			unsigned char* buffer;
			size_t len;
			
			api.element_to_array(&buffer, &len, new_random_element);
			prepre.add_r_er(buffer, len);
			api.free_buffer(buffer);
			
			api.element_to_array(&buffer, &len, new_message_element);
			prepre.add_m_er(buffer, len);
			api.free_buffer(buffer);
			
			api.free_group_element(random_element);
			api.free_group_element(message_element);
			api.free_group_element(other_random_element);
			api.free_group_element(other_message_element);
			api.free_group_element(new_random_element);
			api.free_group_element(new_message_element);
			
			precomputation_data.emplace_back(MixData{r, s, "", permutation[i]});
		}
		
		next_node.async_send(prepre);
	}
}

void Node::handle_node_message(cmix_proto::CMixMessage message)
{
	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kInitialization: {
			BOOST_LOG_TRIVIAL(trace) << "Handling initialization";
			handle_node_initialization(message.initialization());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kBye: {
			BOOST_LOG_TRIVIAL(trace) << "Handling bye";
			//Todo: find a nice way to handle network shutdown.
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kSecretkey: {
			BOOST_LOG_TRIVIAL(trace) << "Handling SecretKey";
			handle_node_secretkey(message.secretkey());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kPrepre: {
			BOOST_LOG_TRIVIAL(trace) << "Handling PrePre";
			handle_node_prepre(message.prepre());
			break;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_node_message: CMixMessage contains unknown contents.";
		}
	}
	prev_node.async_receive([this](cmix_proto::CMixMessage message) {
		handle_node_message(message);
	});
}

void Node::handle_client_keyexchange(ClientConnections::key_type handle, cmix_proto::KeyExchange ke) {
	void* priv_el = api.get_group_element(true);
	
	data[handle].shared_value = api.derive_shared_key(keypair, reinterpret_cast<uint8_t const*>(ke.public_key().c_str()), ke.public_key().size(), reinterpret_cast<uint8_t const*>(ke.value().c_str()), ke.value().size(), priv_el, true);
	
	void* ex_val = api.get_key_exchange_value(priv_el);
	
	api.free_group_element(priv_el);
	
	cmix_proto::KeyExchange exchange;
	
	unsigned char* buffer;
	size_t len;
	
	api.element_to_array(&buffer, &len, keypair.pub);
	exchange.set_public_key(buffer, len);
	api.free_buffer(buffer);
	
	api.element_to_array(&buffer, &len, ex_val);
	exchange.set_value(buffer, len);
	api.free_buffer(buffer);
	
	api.free_group_element(ex_val);
	
	clients.at(handle).async_send(exchange);
}

void Node::handle_client_bye(ClientConnections::key_type handle, cmix_proto::Bye)
{
	clients.at(handle).close();
	clients.erase(handle);
}

void Node::handle_client_message(ClientConnections::key_type handle, cmix_proto::CMixMessage message)
{
	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kKeyexchange: {
			BOOST_LOG_TRIVIAL(trace) << "Handling keyexchange";
			handle_client_keyexchange(handle, message.keyexchange());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kBye: {
			BOOST_LOG_TRIVIAL(trace) << "Handling bye";
			handle_client_bye(handle, message.bye());
			return;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_client_message: CMixMessage contains unknown contents.";
		}
	}
	clients.at(handle).async_receive([this, handle](cmix_proto::CMixMessage message){
		handle_client_message(handle, message);
	});
}

void Node::handle_imanode(Purgatory::iterator handle) {
	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 c) {
	std::string client_id = c.id();
	clients.emplace(c.id(), 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);
			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() {
	precomputation_data.clear();
	precomputation_data.reserve(clients.size());
	
	std::vector<size_t> permutation(clients.size());
	std::iota(permutation.begin(), permutation.end(), 0);
	//ToDo: generate something different than the ID permutation.
	
	cmix_proto::PrePre prepre;
	
	auto perm_it = permutation.begin();
	for(auto const& pair : clients) {
		GroupElement r = api.get_group_element(true);
		GroupElement s = api.get_group_element(true);
	
		GroupElement random_element;
		GroupElement message_element;
		
		api.encrypt(&random_element, &message_element, r, network_key);
		
		unsigned char* buffer;
		size_t len;
		
		api.element_to_array(&buffer, &len, random_element);
		prepre.add_m_er(buffer, len);
		api.free_buffer(buffer);
		
		api.element_to_array(&buffer, &len, message_element);
		prepre.add_m_er(buffer, len);
		api.free_buffer(buffer);
		
		next_node.async_send(prepre);
		
		precomputation_data.emplace_back(MixData{r, s, pair.first, *perm_it++});
		
		api.free_group_element(random_element);
		api.free_group_element(message_element);
	}
	
	next_node.async_send(prepre);
}