#include "node.hpp"

#include "logging.hpp"

#include "bignum.h"

#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()
, network_settings(network_settings)
, prev_node(Client(tcp::socket(io_service)))
, next_node(tcp::socket(io_service))
, api(get_implementation())
, keypair(api.create_key_pair())
, network_pub_key()
{
	GOOGLE_PROTOBUF_VERIFY_VERSION;

	auto on_connect = [this, network_settings](){
		next_node.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)
{
	purgatory.emplace_back(std::move(socket));

	std::list<Client>::iterator it = --purgatory.end();
	purgatory.back().on_done(
		[this, it]() {
			purgatory.erase(it);
		}
	);

	it->receive([this, it](std::vector<uint8_t> const& message_buffer) {
		handle_message(it, message_buffer);
	});
}

void Node::start_initialisation() {
	cmix_proto::Initialization init;
	init.set_public_share(keypair.pub, keypair.pub_len);
	
	next_node.send(init);
}

cmix_proto::CMixMessage Node::parse_cmix_message(const std::vector<uint8_t>& buffer)
{
	cmix_proto::CMixMessage message;
	if(!message.ParseFromArray(buffer.data(), buffer.size())) {
		BOOST_LOG_TRIVIAL(error) << "Received something which was not a CMixMessage";
		io_service.stop();
		throw std::runtime_error("Network communication was disrupted in a unrecoverable way.");
	}
	return message;
}

void Node::handle_initialization(const cmix_proto::Initialization& init)
{
	if(network_settings.is_first) {
		std::string share = init.public_share();
		
		network_pub_key = std::vector<uint8_t>(share.begin(), share.end());
		start_precomputation();
	} 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{keypair.pub, keypair.pub_len};
		Bignum mod = allocate_bignum(keypair.pub_len);
		Bignum new_shared = allocate_bignum(keypair.pub_len);
		
		if(multiply_mod(&new_shared, shared, my_share, mod) != NoError) {
			BOOST_LOG_TRIVIAL(fatal) << "Group multiplication failed";
			throw std::runtime_error("Group multiplication failed");
		}
		
		cmix_proto::Initialization init;
		init.set_public_share(new_shared.data, new_shared.len);
		
		next_node.send(init);
	
		free_bignum(&shared);
		free_bignum(&mod);
		free_bignum(&new_shared);
	}
}

void Node::handle_node_message(const std::vector<uint8_t>& message_buffer)
{
	cmix_proto::CMixMessage message = parse_cmix_message(message_buffer);
	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kInitialization: {
			handle_initialization(message.initialization());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kBye: {
			prev_node.close();
			return;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_node_message: CMixMessage contains unknown contents.";
		}
	}
	prev_node.receive([this](std::vector<uint8_t> const& buffer) {
		handle_node_message(buffer);
	});
}

void Node::handle_client_message(std::list<Client>::iterator handle, const std::vector<uint8_t>& message_buffer)
{
	cmix_proto::CMixMessage message = parse_cmix_message(message_buffer);
	BOOST_LOG_TRIVIAL(trace) << "case: " << message.contents_case();
	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kKeyexchange: {
			BOOST_LOG_TRIVIAL(trace) << "Deriving shared key";
			api.derive_shared_key(keypair, reinterpret_cast<uint8_t const*>(message.keyexchange().public_key().c_str()), true);
			
			handle->receive([this, handle](std::vector<uint8_t> const& buffer){
				handle_client_message(handle, buffer);
			});
			return;
		}
		case cmix_proto::CMixMessage::ContentsCase::kBye: {
			BOOST_LOG_TRIVIAL(trace) << "Handling bye";
			
			handle->close();
			clients.erase(handle);
			
			if(clients.size() == 0) {
				cmix_proto::Bye bye;
				next_node.send(bye);
				
				io_service.post([this]{
					BOOST_LOG_TRIVIAL(trace) << "Shutting down";
					io_service.stop();
				});
				
			}
			
			return;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_client_message: CMixMessage contains unknown contents.";
		}
	}
	handle->close();
	clients.erase(handle);
}

void Node::handle_imanode(std::list<Client>::iterator handle) {
	handle->on_done([]{});
	prev_node = PrevNode(std::move(*handle));
	purgatory.erase(handle);
	prev_node.receive([this](std::vector<uint8_t> const& message_buffer){
		handle_node_message(message_buffer);
	});
}

void Node::handle_imaclient(std::list<Client>::iterator handle) {
	std::list<Client>::iterator it = clients.emplace(clients.end(), std::move(*handle));
	it->on_done([this, it]{
		clients.erase(it);
	});
	purgatory.erase(handle);
	it->receive([this, it](std::vector<uint8_t> buffer) {
		handle_client_message(it, buffer);
	});
}

void Node::handle_message(std::list<Client>::iterator handle, const std::vector<uint8_t>& message_buffer)
{
	cmix_proto::CMixMessage message;
	try {
		 message = parse_cmix_message(message_buffer);
	} catch(std::runtime_error const& e) {
		purgatory.erase(handle);
		return;
	}

	switch(message.contents_case()) {
		case cmix_proto::CMixMessage::ContentsCase::kImanode: {
			handle_imanode(handle);
			return;
		}
		case cmix_proto::CMixMessage::ContentsCase::kImaclient: {
			handle_imaclient(handle);
			return;
		}
		default: {
			BOOST_LOG_TRIVIAL(error) << "handle_message: CMixMessage contains unknown contents.";
		}
	}
	
	purgatory.erase(handle);
}

void Node::start_precomputation() {
}