#include "node.hpp"

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

#include "logging.hpp"

#include <iostream>

using namespace boost::asio::ip;

Node::Node(ListenSettings const& listen_settings, NodeNetworkSettings network_settings)
: io_service()
, server(io_service, listen_settings, [this](boost::asio::ip::tcp::socket&& socket){accept_handler(std::move(socket));})
, clients()
, data()
, network_settings(network_settings)
, prev_node(Receiver(tcp::socket(io_service)))
, next_node(Sender(tcp::socket(io_service)))
, api(get_implementation())
, keypair(api.create_key_pair())
, network_key()
, shutting_down(false)
{
	GOOGLE_PROTOBUF_VERIFY_VERSION;

	auto on_connect = [this, network_settings](){
		next_node.async_send(cmix_proto::ImANode());
		if(network_settings.is_first) {
			start_initialisation();
		}
	};
	
	next_node.async_connect(network_settings.next_host, network_settings.next_port, on_connect);
}

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

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

void Node::accept_handler(boost::asio::ip::tcp::socket&& socket)
{
	std::list<Receiver>::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::start_initialisation() {
	cmix_proto::Initialization init;
	init.set_public_share(keypair.pub, keypair.pub_len);
	
	next_node.async_send(init);
}

void Node::handle_node_initialization(const cmix_proto::Initialization& init)
{
	if(init.public_share().size() != keypair.pub_len) {
		prev_node.close();
		cmix_proto::Bye bye;
		next_node.async_send(bye, [this]{stop();});
		BOOST_LOG_TRIVIAL(fatal) << "Previous node dit not send proper initialization message.";
		return;
	}

	if(network_settings.is_first) {
		cmix_proto::SecretKey sec;
		sec.set_secret_key(network_key.data(), network_key.size());
		next_node.async_send(sec);
	} else {
	
		Bignum shared = allocate_bignum(init.public_share().size());
		std::copy_n(init.public_share().data(), init.public_share().size(), shared.data);
		Bignum my_share = allocate_bignum(keypair.pub_len);
		std::copy_n(my_share.data, my_share.len, keypair.pub);
		Bignum mod = allocate_bignum(keypair.pub_len);
		get_curve25519_mod(&mod);
		Bignum new_shared = allocate_bignum(keypair.pub_len);
		
		calculate_shared_key_part(&new_shared, shared, my_share, mod);
		
		cmix_proto::Initialization init;
		init.set_public_share(new_shared.data, new_shared.len);
		
		next_node.async_send(init);
	
		free_bignum(&shared);
		free_bignum(&my_share);
		free_bignum(&mod);
		free_bignum(&new_shared);
	}
}

void Node::handle_node_bye(cmix_proto::Bye) {
	prev_node.close();
	if (!shutting_down) {
		send_bye();
		prev_node.async_receive([this](cmix_proto::CMixMessage message){
			handle_node_message(message);
		});
	} else {
		io_service.stop();
	}
}

void Node::handle_node_secretkey(cmix_proto::SecretKey const& secret)
{
	std::string share = secret.secret_key();
	network_key = std::vector<uint8_t>(share.begin(), share.end());
	
	if(network_settings.is_first) {
		start_precomputation();
	}
}

void Node::send_bye()
{
	cmix_proto::Bye bye;
	next_node.async_send(bye);
	shutting_down = true;
}

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";
			handle_node_bye(message.bye());
			return;
		}
		case cmix_proto::CMixMessage::ContentsCase::kSecretkey: {
			BOOST_LOG_TRIVIAL(trace) << "Handling SecretKey";
			handle_node_secretkey(message.secretkey());
			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) {
	data[handle].shared_value = api.derive_shared_key(keypair, reinterpret_cast<uint8_t const*>(ke.public_key().c_str()), true);
	
	cmix_proto::KeyExchange exchange;
	exchange.set_public_key(keypair.pub, keypair.pub_len);
	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);
	if(clients.size() == 0) {
		send_bye();
	}
}

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 = Receiver(std::move(*handle));
	purgatory.erase(handle);
	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(), SenderReceiver(std::move(*handle)));
	clients.at(c.id()).on_done([this, client_id]{
		clients.erase(client_id);
	});
	purgatory.erase(handle);
	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() {
}