#include "node.hpp"

template <typename T>
void fill_precomputation_pre_message(CMixContext& ctx, cmix_proto::PrePre& prepre, 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);
	}
	
	std::stringstream ss; 
	ss << "permutation:";
	for(auto i = 0u; i < ctx.nr_participants; ++i) {
		ss << " " << ctx.permutation[i];
	}
	BOOST_LOG_TRIVIAL(trace) << ss.str();
	
	size_t len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	std::vector<char*> r_ers(ms.size(), nullptr);
	std::vector<char*> m_ers(ms.size(), nullptr);
	std::vector<char const*> rsv(ms.size(), nullptr);
	std::vector<char const*> msv(ms.size(), nullptr);
	
	for(auto i = 0; i < ms.size(); ++i) {
		std::string* r = prepre.add_r_er();
		r->resize(len);
		r_ers[i] = &(*r)[0];
	
		std::string* m = prepre.add_m_er();
		m->resize(len);
		m_ers[i] = &(*m)[0];
		
		rsv[i] = rs.Get(i).data();
		msv[i] = ms.Get(i).data();
	}
	
	if(encrypt_r_and_multiply(
		&ctx, 
		r_ers.data(),
		m_ers.data(),
		rsv.data(),
		msv.data()
	) != no_error) {
		exit(-1);
	}
}

template <typename T>
void fill_precomputation_mix_message(CMixContext const& ctx, cmix_proto::PreMix& premix, T const& rs, T const& ms) {
	size_t el_len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	std::vector<char*> r_epirs(ms.size(), nullptr);
	std::vector<char*> m_epirs(ms.size(), nullptr);
	std::vector<char const*> rsv(ms.size(), nullptr);
	std::vector<char const*> msv(ms.size(), nullptr);
	
	for(auto i = 0; i < ms.size(); ++i) {
		std::string* r = premix.add_r_epirs();
		r->resize(el_len);
		r_epirs[i] = &(*r)[0];
		
		std::string* m = premix.add_m_epirs();
		m->resize(el_len);
		m_epirs[i] = &(*m)[0];
		
		rsv[i] = rs.Get(i).data();
		msv[i] = ms.Get(i).data();
	}
	
	permute_and_multiply_encrypted_s(
		&ctx,
		r_epirs.data(),
		m_epirs.data(),
		rsv.data(),
		msv.data()
	);
}

template <typename T>
void fill_precomputation_post_message(CMixContext& ctx, cmix_proto::PrePost& prepost, T const& rs, T const& ms, T const& cs = T{}) {

	std::vector<char const*> rsv(ms.size(), nullptr);
	std::vector<char const*> msv(ms.size(), nullptr);
	for(auto i = 0; i < ms.size(); ++i) {		
		rsv[i] = rs.Get(i).data();
		msv[i] = ms.Get(i).data();
	}
	
	std::string* decryption_commitment = prepost.add_decryption_commitments();
	decryption_commitment->resize(get_commitment_length(&ctx));
	
	precomputation_post_process(
		&ctx,
		&(*decryption_commitment)[0],
		rsv.data(),
		msv.data()
	);
	
	*prepost.mutable_r_epirs() = rs;
	*prepost.mutable_m_epirs() = ms;
	*prepost.mutable_decryption_commitments() = cs;
}

template <typename T>
void fill_commitments_message(CMixContext& ctx, cmix_proto::Commitments& commits, T const& ms, T const& cs) {

	std::vector<char const*> msv(ms.size(), nullptr);
	for(auto i = 0; i < ms.size(); ++i) {		
		msv[i] = ms.Get(i).data();
	}
	
	std::vector<char> commitment(get_commitment_length(&ctx), '\0');
	
	commit_precomputation_ciphertext(
		&ctx,
		commitment.data(),
		msv.data()
	);
	
	*commits.mutable_m_epirs() = ms;
	*commits.mutable_decryption_commitments() = cs;
	commits.set_ciphertext_commitment(commitment.data(), commitment.size());
}

template <typename T, typename U>
void fill_realtime_pre_message(CMixContext& ctx, cmix_proto::RealPre& realpre, T const& hs, T const& ms, U const& data) {
	size_t len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	
	std::vector<char*> msv1(ms.size(), nullptr);
	std::vector<char const*> msv2(ms.size(), nullptr);
	std::vector<GroupElement> gs(ms.size(), nullptr);
	
	for(int i = 0; i < ms.size(); ++i) {
		std::string* m = realpre.add_m();
		m->resize(len);
		msv1[i] = &(*m)[0];
		
		msv2[i] = ms.Get(i).data();
		
		realpre.add_h(hs.Get(i));
		
		gs[i] = data.at(hs.Get(i)).shared_value;
	}
	
	swap_k_for_r(
		&ctx,
		msv1.data(),
		msv2.data(),
		gs.data()
	);
}

template <typename T>
void fill_realtime_mix_message(CMixContext& ctx, cmix_proto::RealMix& realmix, T const& ms) {
	size_t len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	std::vector<char*> mv(ms.size(), nullptr);
	std::vector<char const*> msv(ms.size(), nullptr);
	for(int i = 0; i < ms.size(); ++i) {
		std::string* m = realmix.add_m();
		m->resize(len);
		mv[i] = &(*m)[0];
		
		msv[i] = ms.Get(i).data();
	}
	
	permute_and_multiply_s(
		&ctx,
		mv.data(),
		msv.data()
	);
}

void fill_realtime_post_message(CMixContext& ctx, cmix_proto::RealPost& realpost, cmix_proto::RealMix const& realmix) {
	
	size_t len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	std::vector<char*> mv(ctx.nr_participants, nullptr);
	for(size_t i = 0; i < ctx.nr_participants; ++i) {
		std::string* m = realpost.add_m_epirs();
		m->resize(len);
		mv[i] = &(*m)[0];
	}
	
	get_epirs(
		&ctx,
		mv.data()
	);
	
	std::string* m = realpost.mutable_pimrs_commitment();
	m->resize(get_commitment_length(&ctx));
	
	std::vector<const char*> mvs(ctx.nr_participants, nullptr);
	for(size_t i = 0; i < ctx.nr_participants; ++i) {
		mvs[i] = realmix.m(int(i)).data();
	}
	
	get_pimrs_commitment(
		&ctx,
		&(*m)[0],
		mvs.data()
	);
}

void fill_realtime_post_message(CMixContext& ctx, cmix_proto::RealPost& n_realpost, cmix_proto::RealPost const& realpost) {
	size_t len = ctx.nr_mixes * get_group_element_array_size(&ctx);
	std::vector<char*> mv(ctx.nr_participants, nullptr);
	std::vector<char const*> mvs(ctx.nr_participants, nullptr);
	for(size_t i = 0; i < ctx.nr_participants; ++i) {
		std::string* m = n_realpost.add_m_epirs();
		m->resize(len);
		mv[i] = &(*m)[0];
		
		mvs[i] = realpost.m_epirs(i).data();
	}
	
	decrypt_epirs(
		&ctx,
		mv.data(),
		mvs.data()
	);
}

void Node::handle_node_initialization(const cmix_proto::Initialization& init)
{
	if(network_settings.is_first) {
		set_network_key(&cmix_ctx, init.public_share().data());
	
		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();
	
	if(network_settings.is_first) {
		start_timer_delayed_mix();
	} else {
		set_network_key(&cmix_ctx, secret.secret_key().data());
		next_node.async_send(secret);
	}
}

void Node::handle_node_prepre(cmix_proto::PrePre const& pre) {
	if(network_settings.is_first) {
		if(performance) {
			performance->send("pre_mix_start");
		}
	
		ArenaMessage<cmix_proto::PreMix> arena;
		auto& premix = arena.get();
		fill_precomputation_mix_message(cmix_ctx, premix, pre.r_er(), pre.m_er());
		next_node.async_send(premix);
		
		if(performance) {
			performance->send("pre_mix_end");
		}
	} else {
		if(performance) {
			performance->send("pre_pre_start");
		}
		
		ArenaMessage<cmix_proto::PrePre> arena;
		auto& prepre = arena.get();
		fill_precomputation_pre_message(cmix_ctx, prepre, pre.r_er(), pre.m_er());
		next_node.async_send(prepre);
		
		if(performance) {
			performance->send("pre_pre_end");
		}
	}
}

void Node::handle_node_premix(cmix_proto::PreMix const& premix) {
	if(network_settings.is_first) {
		if(performance) {
			performance->send("pre_post_start");
		}
	
		ArenaMessage<cmix_proto::PrePost> arena;
		auto& prepost = arena.get();
		fill_precomputation_post_message(cmix_ctx, prepost, premix.r_epirs(), premix.m_epirs());
		next_node.async_send(prepost);	
		
		if(performance) {
			performance->send("pre_post_end");
		}
		
	} else {
		if(performance) {
			performance->send("pre_mix_start");
		}
		
		ArenaMessage<cmix_proto::PreMix> arena;
		auto& n_premix = arena.get();
		fill_precomputation_mix_message(cmix_ctx, n_premix, premix.r_epirs(), premix.m_epirs());
		next_node.async_send(n_premix);
		
		if(performance) {
			performance->send("pre_mix_end");
		}
	}
}

void Node::handle_node_prepost(cmix_proto::PrePost const& prepost) {
	if(performance) {
		performance->send("pre_post_start");
	}
	
	if(network_settings.is_last) {
		ArenaMessage<cmix_proto::PrePost> arena;
		auto& n_prepost = arena.get();
		fill_precomputation_post_message(cmix_ctx, n_prepost, prepost.r_epirs(), prepost.m_epirs(), prepost.decryption_commitments());
	
		ArenaMessage<cmix_proto::Commitments> arena2;
		auto& commits = arena2.get();
		fill_commitments_message(cmix_ctx, commits, n_prepost.m_epirs(), n_prepost.decryption_commitments());
		next_node.async_send(commits);
	} else {
		ArenaMessage<cmix_proto::PrePost> arena;
		auto& n_prepost = arena.get();
		fill_precomputation_post_message(cmix_ctx, n_prepost, prepost.r_epirs(), prepost.m_epirs(), prepost.decryption_commitments());
		next_node.async_send(n_prepost);
	}
	if(performance) {
		performance->send("pre_post_end");
	}
}

void Node::handle_node_commitments(cmix_proto::Commitments const& comm) {
	
	//Do your commitments check/storage.
	
	if(!network_settings.is_first && !network_settings.is_last) {
		next_node.async_send(comm);
	}
	if(network_settings.is_first) {
		next_node.async_send(comm, [this]{start_realtime_phase();});
	}
}

void Node::handle_node_realpre(cmix_proto::RealPre const& realpre) {
	if(network_settings.is_first) {
		if(performance) {
			performance->send("real_mix_start");
		}
	
		ArenaMessage<cmix_proto::RealMix> arena;
		auto& realmix = arena.get();
		fill_realtime_mix_message(cmix_ctx, realmix, realpre.m());
		next_node.async_send(realmix);
		
		if(performance) {
			performance->send("real_mix_end");
		}
	} else {
		if(performance) {
			performance->send("real_pre_start");
		}
			
		ArenaMessage<cmix_proto::RealPre> arena;
		auto& n_realpre = arena.get();
		fill_realtime_pre_message(cmix_ctx, n_realpre, realpre.h(), realpre.m(), data);
		next_node.async_send(n_realpre);
		
		if(performance) {
			performance->send("real_pre_end");
		}
	}
}

void Node::handle_node_realmix(cmix_proto::RealMix const& realmix) {	
	if(network_settings.is_last) {
		BOOST_LOG_TRIVIAL(trace) << "Doing the last step:";
		
		if(performance) {
			performance->send("real_mix_start");
		}
		
		ArenaMessage<cmix_proto::RealMix> arena;
		auto& n_realmix = arena.get();
		fill_realtime_mix_message(cmix_ctx, n_realmix, realmix.m());
		
		ArenaMessage<cmix_proto::RealPost> arena2;
		auto& realpost = arena2.get();
		fill_realtime_post_message(cmix_ctx, realpost, n_realmix);
		
		next_node.async_send(realpost);
		
		if(performance) {
			performance->send("real_mix_end");
		}
		
	} else {
		if(performance) {
			performance->send("real_mix_start");
		}
	
		ArenaMessage<cmix_proto::RealMix> arena;
		auto& n_realmix = arena.get();
		fill_realtime_mix_message(cmix_ctx, n_realmix, realmix.m());
		next_node.async_send(n_realmix);
		
		if(performance) {
			performance->send("real_mix_end");
		}
	}
}

void Node::handle_node_realpost(cmix_proto::RealPost const& realpost) {
	if(network_settings.is_last) {
		if(performance) {
			performance->send("real_post_start");
		}
		
		ArenaMessage<cmix_proto::RealPost> arena;
		auto& n_realpost = arena.get();
		fill_realtime_post_message(cmix_ctx, n_realpost, realpost);
		
		size_t len = cmix_ctx.nr_mixes * get_group_element_array_size(&cmix_ctx);
		std::vector<std::string> message_strings(cmix_ctx.nr_participants);
		std::vector<const char*> pirs(cmix_ctx.nr_participants, nullptr);
		std::vector<char*> msgs(cmix_ctx.nr_participants, nullptr);
		for(size_t i = 0; i < cmix_ctx.nr_participants; ++i) {
			message_strings[i].resize(len);
			msgs[i] = &message_strings[i][0];
			
			pirs[i] = n_realpost.m_epirs(i).data();
		}
		
		remove_pirs (
			&cmix_ctx,
			msgs.data(),
			pirs.data()
		);
		
		//for (auto&& mes : msgs) {
		//	std::cout << std::string(mes+32, 32) << std::endl;
		//}
		
		
		for(auto i = 0u; i < cmix_ctx.nr_participants; i++) {
			char* message;
			element_to_message(&cmix_ctx, reinterpret_cast<unsigned char**>(&message), msgs[i]);
			
			char* dest;
			size_t dest_len;
			
			char* payload = nullptr;
			size_t payload_len;
			
			split_message(&cmix_ctx, &dest, &dest_len, &payload, &payload_len, message);
			
			std::string dest_s = std::string(dest, dest_len);
			
			cmix_proto::Payload pay;
			pay.set_payload(payload, payload_len);
			
			clients.at(dest_s).async_send(pay);
			
			free(dest);
			free(payload);
			free(message);
		}
		
		if(performance) {
			performance->send("real_post_end");
		}
	} else {
		if(performance) {
			performance->send("real_post_start");
		}
		
		ArenaMessage<cmix_proto::RealPost> arena;
		auto& n_realpost = arena.get();
		fill_realtime_post_message(cmix_ctx, n_realpost, realpost);
		next_node.async_send(n_realpost);
		
		if(performance) {
			performance->send("real_post_end");
		}
	}
}

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";
			if(send_bye(true)) {
				return;
			} else {
				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::kCommitments: {
			BOOST_LOG_TRIVIAL(trace) << "Handling Commitments";
			handle_node_commitments(message.commitments());
			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;
		}
		case cmix_proto::CMixMessage::ContentsCase::kRealpost: {
			BOOST_LOG_TRIVIAL(trace) << "Handling RealPost";
			handle_node_realpost(message.realpost());
			break;
		}

		case cmix_proto::CMixMessage::ContentsCase::CONTENTS_NOT_SET:
		case cmix_proto::CMixMessage::ContentsCase::kImaclient:
		case cmix_proto::CMixMessage::ContentsCase::kImanode:
		case cmix_proto::CMixMessage::ContentsCase::kKeyexchange:
		case cmix_proto::CMixMessage::ContentsCase::kNodeready:
		case cmix_proto::CMixMessage::ContentsCase::kUsermessage:
		case cmix_proto::CMixMessage::ContentsCase::kPayload:
		case cmix_proto::CMixMessage::ContentsCase::kPerformance:
		{
			BOOST_LOG_TRIVIAL(error) << "handle_node_message: CMixMessage contains unknown contents.";
			break;
		}
	}
	prev_node.async_receive([this](cmix_proto::CMixMessage message2) {
		handle_node_message(message2);
	});
}