#include "node.hpp"

template <typename T>
cmix_proto::PrePre fill_precomputation_pre_message(CMixContext& ctx, T const& rs, T const& ms) {
	if(start_mix(&ctx, rs.size()) != no_error) {
		exit(-1);
	}
	
	if(initialize_mix_randomness(&ctx) != no_error) {
		exit(-1);
	}
	
	cmix_proto::PrePre prepre;
	
	for(size_t i = 0; i < ms.size(); ++i) {
		prepre.add_m_er();
		prepre.add_r_er();
	
		size_t len = get_group_element_array_size(&ctx);
		
		prepre.mutable_m_er(i)->resize(len);
		prepre.mutable_r_er(i)->resize(len);
		
		{
			std::stringstream ss;
			ss << "r: ";
			std::string r = to_string(ctx.r[i], 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_and_multiply(
			&ctx, 
			&(*prepre.mutable_r_er(i))[0],
			&(*prepre.mutable_m_er(i))[0],
			rs.Get(i).data(),
			ms.Get(i).data(),
			i
		) != no_error) {
			exit(-1);
		}
	}
	return prepre;
}

template <typename T>
cmix_proto::PreMix fill_precomputation_mix_message(CMixContext const& ctx, T const& rs, T const& ms) {
	cmix_proto::PreMix premix;
	
	for(size_t i = 0; i < ms.size(); ++i) {
		premix.add_r_epirs();
		premix.add_m_epirs();
	}
	for(size_t i = 0; i < ms.size(); ++i) {
		auto new_pos = ctx.permutation[i];
		size_t el_len = get_group_element_array_size(&ctx);
		
		premix.mutable_r_epirs(new_pos)->resize(el_len);
		premix.mutable_m_epirs(new_pos)->resize(el_len);
		
		{
			std::stringstream ss;
			ss << "s: ";
			std::string s = to_string(ctx.s[i], ctx);
			for(auto&& c : s) {
				ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
			}
			BOOST_LOG_TRIVIAL(trace) << ss.str();
		}
		
		multiply_encrypted_s(
			&ctx,
			&(*premix.mutable_r_epirs(new_pos))[0],
			&(*premix.mutable_m_epirs(new_pos))[0],
			rs.Get(i).data(),
			ms.Get(i).data(),
			i
		);
	}
	return premix;
}

template <typename T>
cmix_proto::PrePost fill_precomputation_post_message(CMixContext& ctx, T const& rs, T const& ms) {
	cmix_proto::PrePost prepost;
	
	for(size_t i = 0; i < ms.size(); ++i) {
		prepost.add_r_epirs();
		prepost.add_m_epirs();
		
		size_t el_len = get_group_element_array_size(&ctx);
		
		prepost.mutable_r_epirs(i)->resize(el_len);
		prepost.mutable_m_epirs(i)->resize(el_len);
		
		post_process(
			&ctx,
			&(*prepost.mutable_r_epirs(i))[0],
			&(*prepost.mutable_m_epirs(i))[0],
			rs.Get(i).data(),
			ms.Get(i).data(),
			i
		);
		
		{
			std::stringstream ss;
			ss << "pirs: ";
			std::string pirs = to_string(ctx.pirs[i], ctx);
			for(auto&& c : pirs) {
				ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
			}
			BOOST_LOG_TRIVIAL(trace) << ss.str();
		} 
		
		*prepost.mutable_r_epirs(i) = rs.Get(i);
		*prepost.mutable_m_epirs(i) = ms.Get(i);
	}
	
	return prepost;
}

template <typename T>
cmix_proto::RealPre fill_realtime_pre_message(CMixContext& ctx, T const& hs, T const& ms, Node::ClientData const& data) {
	cmix_proto::RealPre realpre;
	
	size_t len = get_group_element_array_size(&ctx);
	for(int i = 0; i < ms.size(); ++i) {
		realpre.add_h();
		realpre.add_m();
		realpre.mutable_m(i)->resize(len);
		
		swap_k_for_r(
			&ctx,
			&(*realpre.mutable_m(i))[0],
			ms.Get(i).data(),
			data.at(hs.Get(i)).shared_value,
			i
		);
		*realpre.mutable_h(i) = hs.Get(i);
	}
	
	return realpre;
}

template <typename T>
cmix_proto::RealMix fill_realtime_mix_message(CMixContext& ctx, T const& ms) {
	cmix_proto::RealMix realmix;
	
	size_t len = get_group_element_array_size(&ctx);
	
	for(int i = 0; i < ms.size(); ++i) {
		realmix.add_m();
	}
	
	for(int i = 0; i < ms.size(); ++i) {
		auto new_pos = ctx.permutation[i];
		
		realmix.mutable_m(i)->resize(len);
		
		multiply_s(
			&ctx,
			&(*realmix.mutable_m(new_pos))[0],
			ms.Get(i).data(),
			i
		);
		
	}
	
	return realmix;
}

void Node::handle_node_initialization(const cmix_proto::Initialization& init)
{
	if(network_settings.is_first) {
		set_network_key(&cmix_ctx, init.public_share().data(), init.public_share().size());
	
		cmix_proto::SecretKey sec;
		sec.set_secret_key(init.public_share().data(), init.public_share().size());
		
		next_node.async_send(sec);
	} else {
		size_t len = get_group_element_array_size(&cmix_ctx);
		
		cmix_proto::Initialization n_init;
		n_init.mutable_public_share()->resize(len);
		add_public_share(&cmix_ctx, &(*n_init.mutable_public_share())[0], init.public_share().data());
	
		BOOST_LOG_TRIVIAL(trace) << "sending message: " << n_init.ShortDebugString();
		next_node.async_send(n_init);
	}
}

void Node::handle_node_secretkey(cmix_proto::SecretKey const& secret)
{
	std::string share = secret.secret_key();
	
	set_network_key(&cmix_ctx, secret.secret_key().data(), secret.secret_key().size());
	
	if(network_settings.is_first) {
		timer.expires_from_now(boost::posix_time::seconds(4));
		timer.async_wait([this](boost::system::error_code const& ec) {
			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 = fill_precomputation_mix_message(cmix_ctx, pre.r_er(), pre.m_er());
		next_node.async_send(premix);
	} else {
		cmix_proto::PrePre prepre = fill_precomputation_pre_message(cmix_ctx, pre.r_er(), pre.m_er());
		next_node.async_send(prepre);
	}
}

void Node::handle_node_premix(cmix_proto::PreMix const& premix) {
	if(network_settings.is_first) {
		cmix_proto::PrePost prepost = fill_precomputation_post_message(cmix_ctx, premix.r_epirs(), premix.m_epirs());
		next_node.async_send(prepost);	
	} else {
		cmix_proto::PreMix n_premix = fill_precomputation_mix_message(cmix_ctx, premix.r_epirs(), premix.m_epirs());
		next_node.async_send(n_premix);
	}
}

void Node::handle_node_prepost(cmix_proto::PrePost const& prepost) {
	if(network_settings.is_first) {
		start_realtime_phase();
	} else {
		cmix_proto::PrePost n_prepost = fill_precomputation_post_message(cmix_ctx, prepost.r_epirs(), prepost.m_epirs());
		next_node.async_send(n_prepost);
	}
}

void Node::handle_node_realpre(cmix_proto::RealPre const& realpre) {
	if(network_settings.is_first) {
		cmix_proto::RealMix n_realmix = fill_realtime_mix_message(cmix_ctx, realpre.m());
		next_node.async_send(n_realmix);
	} else {
		cmix_proto::RealPre n_realpre = fill_realtime_pre_message(cmix_ctx, realpre.h(), realpre.m(), data);
		next_node.async_send(n_realpre);
	}
}

void Node::handle_node_realmix(cmix_proto::RealMix const& realmix) {
	if(network_settings.is_last) {
		BOOST_LOG_TRIVIAL(trace) << "Doing the last step:";
		
		cmix_proto::RealMix n_realmix = fill_realtime_mix_message(cmix_ctx, realmix.m());
		
		size_t len = get_group_element_array_size(&cmix_ctx);
		std::string str;
		str.resize(len);
		
		for(int i = 0; i < n_realmix.m_size(); i++) {
			remove_r_and_s(&cmix_ctx, &str[0], n_realmix.m(i).data(), i);
			
			{
				std::stringstream ss;
				for(auto&& c : str) {
					ss << "\\" << std::setw(3) << std::setfill('0') << std::oct << (unsigned int) c;
				}
				BOOST_LOG_TRIVIAL(trace) << ss.str();
			} 
		}
		
	} else {
		cmix_proto::RealMix n_realmix = fill_realtime_mix_message(cmix_ctx, realmix.m());
		next_node.async_send(n_realmix);
	}
}

void Node::handle_node_message(cmix_proto::CMixMessage message)
{
	BOOST_LOG_TRIVIAL(trace) << "Message: " << message.ShortDebugString();
	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;
		}
		case cmix_proto::CMixMessage::ContentsCase::kPremix: {
			BOOST_LOG_TRIVIAL(trace) << "Handling PreMix";
			handle_node_premix(message.premix());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kPrepost: {
			BOOST_LOG_TRIVIAL(trace) << "Handling PrePost";
			handle_node_prepost(message.prepost());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kRealpre: {
			BOOST_LOG_TRIVIAL(trace) << "Handling RealPre";
			handle_node_realpre(message.realpre());
			break;
		}
		case cmix_proto::CMixMessage::ContentsCase::kRealmix: {
			BOOST_LOG_TRIVIAL(trace) << "Handling RealMix";
			handle_node_realmix(message.realmix());
			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);
	});
}