#include "cmix.h"

#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>

struct CMixContext initialize_cmix_context(struct Api api, unsigned int nr_mixes) {
	return (struct CMixContext){
		.api = api,
		.nr_mixes = nr_mixes,
		.nr_participants = 0,
		.network_key = NULL,
		.r = NULL,
		.s = NULL,
		.permutation = NULL,
		.decryption_shares = NULL,
		.EPiRS = NULL,
		.PiMRS = NULL
	};
}

enum cmix_error alloc_mix(struct CMixContext* ctx) {
	ctx->r = (GroupElement**) calloc(ctx->nr_participants, sizeof(GroupElement*));
	ctx->s = (GroupElement**) calloc(ctx->nr_participants, sizeof(GroupElement*));
	ctx->permutation = (unsigned int**) calloc(ctx->nr_participants, sizeof(unsigned int*));
	ctx->decryption_shares = (GroupElement**) calloc(ctx->nr_participants, sizeof(GroupElement*));
	ctx->EPiRS = (GroupElement**) calloc(ctx->nr_participants, sizeof(GroupElement*));
	ctx->PiMRS = (GroupElement**) calloc(ctx->nr_participants, sizeof(GroupElement*));

	for(size_t m = 0; m < ctx->nr_participants; ++m) {
		ctx->r[m] = (GroupElement*) calloc(ctx->nr_mixes, sizeof(GroupElement));
		ctx->s[m] = (GroupElement*) calloc(ctx->nr_mixes, sizeof(GroupElement));
		ctx->permutation[m] = (unsigned int*) calloc(ctx->nr_mixes, sizeof(unsigned int));
		ctx->decryption_shares[m] = (GroupElement*) calloc(ctx->nr_mixes, sizeof(GroupElement));
		ctx->EPiRS[m] = (GroupElement*) calloc(ctx->nr_mixes, sizeof(GroupElement));
		ctx->PiMRS[m] = (GroupElement*) calloc(ctx->nr_mixes, sizeof(GroupElement));
	}

	if(!ctx->r || !ctx->s || !ctx->permutation || !ctx->decryption_shares || !ctx->EPiRS || !ctx->PiMRS ) {
		return out_of_memory;
	}
	return no_error;
}

void release_mix(struct CMixContext* ctx) {

	if(ctx->r && ctx->s && ctx->permutation && ctx->decryption_shares && ctx->EPiRS) {
		for(size_t m = 0; m < ctx->nr_participants; ++m) {
			if(ctx->r[m]) {
				for(size_t i = 0; i < ctx->nr_mixes; ++i) {
					ctx->api.free_group_element(ctx->r[m][i]);
				}
				free(ctx->r[m]);
			}
			if(ctx->s[m]){
				for(size_t i = 0; i < ctx->nr_mixes; ++i) {
					ctx->api.free_group_element(ctx->s[m][i]);
				}
			}
			free(ctx->permutation[m]);

			if(ctx->decryption_shares[m]){
				for(size_t i = 0; i < ctx->nr_mixes; ++i) {
					ctx->api.free_group_element(ctx->decryption_shares[m][i]);
				}
			}
			if(ctx->EPiRS[m]) {
				for(size_t i = 0; i < ctx->nr_mixes; ++i) {
					ctx->api.free_group_element(ctx->EPiRS[m][i]);
				}
			}
		}
	}


	for(size_t i = 0; i < ctx->nr_participants; ++i) {

		ctx->api.free_group_element(ctx->PiMRS[i]);
	}

	free(ctx->r);
	ctx->r = NULL;
	
	free(ctx->s);
	ctx->s = NULL;
	
	free(ctx->permutation);
	ctx->permutation = NULL;
	
	free(ctx->decryption_shares);
	ctx->decryption_shares = NULL;
	
	free(ctx->EPiRS);
	ctx->EPiRS = NULL;
	
	free(ctx->PiMRS);
	ctx->PiMRS = NULL;
	
	ctx->nr_participants = 0;
}

void deinitialize(struct CMixContext* ctx)
{
	ctx->api.free_keypair(&ctx->keypair);
	if(ctx->network_key != NULL) {
		ctx->api.free_group_element(ctx->network_key);
		ctx->network_key = NULL;
	}
	release_mix(ctx);
	ctx->api.deinitialize();
}

void element_to_buffer(struct CMixContext const* ctx, char* out_buffer, GroupElement const element) {
	size_t el_size = get_group_element_array_size(ctx);
	
	unsigned char* buffer;
	size_t len;
	
	ctx->api.element_to_array(&buffer, &len, element);
	
	size_t diff = el_size - len;
	memcpy(out_buffer + diff, buffer, len);
	ctx->api.free_buffer(buffer);
}

enum cmix_error initialize_keypair(struct CMixContext* ctx)
{
	ctx->keypair = ctx->api.create_keypair();
	return no_error;
}

enum cmix_error get_public_key(struct CMixContext const* ctx, char* buffer)
{
	element_to_buffer(ctx, buffer, ctx->keypair.pub);
	return no_error;
}

enum cmix_error add_public_share(struct CMixContext const* ctx, char* buffer, char const* share) {
	size_t len = get_group_element_array_size(ctx);
	
	GroupElement el;
	ctx->api.add_public_share(&el, share, len, ctx->keypair.pub);
	
	element_to_buffer(ctx, buffer, el);
	ctx->api.free_group_element(el);
	
	return no_error;
}

enum cmix_error start_mix(struct CMixContext* ctx, size_t nr_participants) {
	release_mix(ctx);
	ctx->nr_participants = nr_participants;
	return alloc_mix(ctx);
}

enum cmix_error generate_permutation(struct CMixContext* ctx) {
	if (ctx->nr_participants == 0) {
		return no_error;
	}
	
	for(size_t m = 0; m < ctx->nr_mixes; ++m) {
		for(unsigned int i = 0; i < ctx->nr_participants; ++i) {
			ctx->permutation[i][m] = i;
		}

		//Fisher-Yates shuffle
		unsigned int temp = 0;
		for(unsigned int i = ctx->nr_participants - 1; i > 0; --i) {
			unsigned int rand = ctx->api.get_uniform_int(i+1);
			temp = ctx->permutation[i][m];
			ctx->permutation[i][m] = ctx->permutation[m][rand];
			ctx->permutation[rand][m] = temp;
		}
	}


	return no_error;
}

enum cmix_error initialize_mix_randomness(struct CMixContext* ctx) {

	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			ctx->r[i][m] = ctx->api.get_group_element(true);
			ctx->s[i][m] = ctx->api.get_group_element(true);
		}
	}
	
	return generate_permutation(ctx);
}

enum cmix_error generate_random_message(struct CMixContext* ctx, char* buffer) {
	GroupElement el = ctx->api.get_group_element(true);
	element_to_buffer(ctx, buffer, el);
	ctx->api.free_group_element(el);
	
	return no_error;
}

size_t get_group_element_array_size(struct CMixContext const* ctx) {
	return ctx->api.get_group_element_array_size();
}

size_t get_message_size(struct CMixContext const* ctx)
{
	return ctx->api.get_message_size();
}

enum cmix_error set_network_key(struct CMixContext* ctx, char const* buffer) {
	ctx->network_key = ctx->api.array_to_element(buffer, get_group_element_array_size(ctx), true);
	return no_error;
}

enum cmix_error encrypt_r(struct CMixContext const* ctx, char** random_buffer, char** message_buffer) {
	

	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; m++) {
			size_t offset = m * get_group_element_array_size(ctx);

			GroupElement random_element;
			GroupElement message_element;

			ctx->api.encrypt(&random_element, &message_element, ctx->r[i][m], ctx->network_key);

			element_to_buffer(ctx, random_buffer[i] + offset, random_element);
			element_to_buffer(ctx, message_buffer[i] + offset, message_element);

			ctx->api.free_group_element(random_element);
			ctx->api.free_group_element(message_element);
		}
	}
	return no_error;
}

enum cmix_error encrypt_r_and_multiply(struct CMixContext const* ctx, char** random_buffer, char** message_buffer, const char** random_element, const char** message_element) {
	
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			GroupElement enc_random_element;
			GroupElement enc_message_element;

			size_t offset = m * get_group_element_array_size(ctx);

			ctx->api.encrypt(&enc_random_element, &enc_message_element, ctx->r[i][m], ctx->network_key);

			size_t el_size = get_group_element_array_size(ctx);

			GroupElement other_random_element = ctx->api.array_to_element(random_element[i] + offset, el_size, true);
			GroupElement other_message_element = ctx->api.array_to_element(message_element[i] + offset, el_size, true);

			GroupElement new_random_element = ctx->api.combine(enc_random_element, other_random_element, true);
			GroupElement new_message_element = ctx->api.combine(enc_message_element, other_message_element, true);

			element_to_buffer(ctx, random_buffer[i] + offset, new_random_element);
			element_to_buffer(ctx, message_buffer[i] + offset, new_message_element);

			ctx->api.free_group_element(enc_random_element);
			ctx->api.free_group_element(enc_message_element);
			ctx->api.free_group_element(other_random_element);
			ctx->api.free_group_element(other_message_element);
			ctx->api.free_group_element(new_random_element);
			ctx->api.free_group_element(new_message_element);
		}
	}
	
	return no_error;
}

enum cmix_error permute_and_multiply_encrypted_s(struct CMixContext const* ctx, char** random_buffer, char** message_buffer, char const** random_element, char const** message_element) {
	size_t el_size = get_group_element_array_size(ctx);

	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			unsigned int new_pos = ctx->permutation[i][m];
			size_t offset = m * get_group_element_array_size(ctx);

			GroupElement random_r = ctx->api.array_to_element(random_element[i] + offset, el_size, true);
			GroupElement message_r = ctx->api.array_to_element(message_element[i] + offset, el_size, true);

			GroupElement random_s;
			GroupElement message_s;

			ctx->api.encrypt(&random_s, &message_s, ctx->s[new_pos][m], ctx->network_key);

			GroupElement random_pirs = ctx->api.combine(random_r, random_s, true);
			GroupElement message_pirs = ctx->api.combine(message_r, message_s, true);

			element_to_buffer(ctx, random_buffer[new_pos] + offset, random_pirs);
			element_to_buffer(ctx, message_buffer[new_pos] + offset, message_pirs);

			ctx->api.free_group_element(random_r);
			ctx->api.free_group_element(message_r);
			ctx->api.free_group_element(random_s);
			ctx->api.free_group_element(message_s);
			ctx->api.free_group_element(random_pirs);
			ctx->api.free_group_element(message_pirs);
		}
	}
	
	return no_error;
}

enum cmix_error permute_and_multiply_s(struct CMixContext* ctx, char** out_buffer, char const** message) {
	size_t el_size = get_group_element_array_size(ctx);
	
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * get_group_element_array_size(ctx);
			unsigned int new_pos = ctx->permutation[i][m];

			GroupElement message_el = ctx->api.array_to_element(message[i] + offset, el_size, false);
			ctx->PiMRS[new_pos][m] = ctx->api.combine(message_el, ctx->s[new_pos][m], false);

			element_to_buffer(ctx, out_buffer[new_pos] + offset, ctx->PiMRS[new_pos][m]);

			ctx->api.free_group_element(message_el);
		}
	}
	
	return no_error;
}


enum cmix_error multiply_s(struct CMixContext const* ctx, char* out_buffer, char const* message, size_t index) {
	for(size_t m = 0; m < ctx->nr_mixes; ++m) {
		size_t el_size = get_group_element_array_size(ctx);
		size_t offset = get_group_element_array_size(ctx);

		GroupElement message_el = ctx->api.array_to_element(message + offset, el_size, false);
		GroupElement mult = ctx->api.combine(message_el, ctx->s[m][ctx->permutation[m][index]], false);

		element_to_buffer(ctx, out_buffer + offset, mult);

		ctx->api.free_group_element(message_el);
		ctx->api.free_group_element(mult);
	}
	
	return no_error;
}

enum cmix_error key_exchange_init(struct CMixContext const* ctx, char* pubkey_buffer, char* value_buffer, GroupElement* priv_el) {
	*priv_el = ctx->api.get_group_element(true);
	GroupElement value = ctx->api.get_key_exchange_value(*priv_el);
	
	get_public_key(ctx, pubkey_buffer);
	element_to_buffer(ctx, value_buffer, value);
	
	ctx->api.free_group_element(value);
	return no_error;
}

enum cmix_error key_exchange_responder(struct CMixContext const* ctx, GroupElement* shared_key, char* public_key_buffer, char* exchange_value_buffer, char const* pubkey, char const* value) {
	GroupElement priv_el = ctx->api.get_group_element(true);
	GroupElement ex_val = ctx->api.get_key_exchange_value(priv_el);
	
	size_t el_len = get_group_element_array_size(ctx);
	*shared_key = ctx->api.derive_shared_key(ctx->keypair, (unsigned char const*)pubkey, el_len, (unsigned char const*)value, el_len, priv_el, true);
	
	element_to_buffer(ctx, public_key_buffer, ctx->keypair.pub);
	element_to_buffer(ctx, exchange_value_buffer, ex_val);
	
	ctx->api.free_group_element(priv_el);
	ctx->api.free_group_element(ex_val);
	
	return no_error;
}

enum cmix_error key_exchange_initiator(struct CMixContext const* ctx, GroupElement* shared_key, char const* pubkey, char const* value, GroupElement* priv_el) {
	size_t el_len = get_group_element_array_size(ctx);
	*shared_key = ctx->api.derive_shared_key(ctx->keypair, (unsigned char const*)pubkey, el_len, (unsigned char const*)value, el_len, *priv_el, false);
	
	ctx->api.free_group_element(*priv_el);
	*priv_el = NULL;
	
	return no_error;
}

/*
enum cmix_error precomputation_post_process_unsafe(struct CMixContext* ctx, char** r_out, char** m_out, const char** r_epirs, const char** m_epirs) {
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		GroupElement x = ctx->api.array_to_element(r_epirs[i], get_group_element_array_size(ctx), true);
		GroupElement D = ctx->api.get_decryption_share(x, ctx->keypair.sec);
		
		GroupElement msg = ctx->api.array_to_element(m_epirs[i], get_group_element_array_size(ctx), true);
		GroupElement pirs = ctx->api.multiply(D, msg, true);
		element_to_buffer(ctx, m_out[i], pirs);
		
		memcpy(r_out[i], r_epirs[i], get_group_element_array_size(ctx));
		
		ctx->pirs[i] = pirs; // this is not always usable as only the last node will be able to use this effectively, but we store it anyways.
		
		ctx->api.free_group_element(x);
		ctx->api.free_group_element(D);
		ctx->api.free_group_element(msg);
	}
	
	return no_error;
}
*/

size_t get_commitment_length(struct CMixContext const* ctx) {
	(void) ctx;
	return 1;
}

enum cmix_error precomputation_post_process(struct CMixContext* ctx, char* commitment, const char** r_epirs, const char** m_epirs) {
	size_t len = get_group_element_array_size(ctx);
	
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * ctx->api.get_group_element_array_size();

			GroupElement x = ctx->api.array_to_element(r_epirs[i] + offset, get_group_element_array_size(ctx), true);
			ctx->decryption_shares[i][m] = ctx->api.get_decryption_share(x, ctx->keypair.sec);

			ctx->EPiRS[i][m] = ctx->api.array_to_element(m_epirs[i] + offset, len, true);

			ctx->api.free_group_element(x);
		}
	}

	
	//ToDo: actually calculate commitment.
	
	memset(commitment, 0, get_commitment_length(ctx));
	
	return no_error;
}

enum cmix_error commit_precomputation_ciphertext(struct CMixContext* ctx, char* commitment, const char** m_epirs) {
	(void) m_epirs;
	//ToDo: actually calculate commitment.
	memset(commitment, 0, get_commitment_length(ctx));

	return no_error;
}

enum cmix_error blind_message(struct CMixContext const* ctx, char* m_out, char const* message, GroupElement const* keys, size_t const nr_nodes) {
	size_t len = get_message_size(ctx); 
	
	GroupElement* intermediates = (GroupElement*) calloc(nr_nodes + 1, sizeof(GroupElement));
	if(intermediates == NULL) {
		return out_of_memory;
	}

	for (size_t i = 0; i < ctx->nr_mixes; ++i) {
		size_t offset = i * ctx->api.get_group_element_array_size();

		intermediates[0] = ctx->api.message_to_element(message + i*len, len, true);

		for(size_t j = 0; j < nr_nodes; ++j) {
			intermediates[j+1] = ctx->api.combine(intermediates[j], keys[j], false);
		}

		element_to_buffer(ctx, m_out + offset, intermediates[nr_nodes]);

		for(size_t j = 0; j < nr_nodes + 1; ++j) {
			ctx->api.free_group_element(intermediates[j]);
		}
	}
	
	free(intermediates);
	
	return no_error;
}

enum cmix_error swap_k_for_r(struct CMixContext const* ctx, char** out_buffer, const char** message, const GroupElement* key) {
	size_t len = get_group_element_array_size(ctx);

	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * get_group_element_array_size(ctx);
			GroupElement mes = ctx->api.array_to_element(message[i] + offset, len, false);

			GroupElement unblinded = ctx->api.uncombine(mes, key[i], false);
			GroupElement blinded = ctx->api.combine(unblinded, ctx->r[i][m], false);

			element_to_buffer(ctx, out_buffer[i] + offset, blinded);
			ctx->api.free_group_element(mes);
			ctx->api.free_group_element(unblinded);
			ctx->api.free_group_element(blinded);
		}
	}
	
	return no_error;
}


enum cmix_error get_epirs(struct CMixContext const* ctx, char** out_buffer) {
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * get_group_element_array_size(ctx);
			element_to_buffer(ctx, out_buffer[i] + offset, ctx->EPiRS[i][m]);
		}
	}

	return no_error;
}

enum cmix_error get_pimrs_commitment(struct CMixContext const* ctx, char* out_buffer, const char** pimrss) {
	(void) pimrss;
	memset(out_buffer, '\0', get_commitment_length(ctx));
	return no_error;
}

enum cmix_error decrypt_epirs(struct CMixContext const* ctx, char** out_buffer, const char** epirs) {
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * get_group_element_array_size(ctx);

			GroupElement msg = ctx->api.array_to_element(epirs[i] + offset, get_group_element_array_size(ctx), true);
			GroupElement pirs = ctx->api.combine(ctx->decryption_shares[i][m], msg, true);
			element_to_buffer(ctx, out_buffer[i] + offset, pirs);

			ctx->api.free_group_element(msg);
			ctx->api.free_group_element(pirs);
		}
	}

	return no_error;
}

enum cmix_error remove_pirs(struct CMixContext const* ctx, char** out_buffer, const char** pirs) {
	size_t len = get_group_element_array_size(ctx);

	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		for(size_t m = 0; m < ctx->nr_mixes; ++m) {
			size_t offset = m * get_group_element_array_size(ctx);
			GroupElement g_pirs = ctx->api.array_to_element(pirs[i] + offset, len, false);

			GroupElement mult = ctx->api.uncombine(ctx->PiMRS[i][m], g_pirs, false);

			element_to_buffer(ctx, out_buffer[i] + offset, mult);

			ctx->api.free_group_element(g_pirs);
			ctx->api.free_group_element(mult);
		}
	}
	
	return no_error;
}

enum cmix_error element_to_message(struct CMixContext const* ctx, unsigned char** buffer, char const* element_buffer) {

	*buffer = (unsigned char*) calloc(ctx->nr_mixes * get_message_size(ctx), sizeof(unsigned char));
	for(size_t m = 0; m < ctx->nr_mixes; ++m) {
		unsigned char* temp;
		size_t offset = m * get_group_element_array_size(ctx);
		GroupElement el = ctx->api.array_to_element(element_buffer + offset, get_group_element_array_size(ctx), true);
		ctx->api.element_to_message(&temp, el);
		ctx->api.free_group_element(el);
		memcpy((*buffer) + m * get_message_size(ctx), temp, get_message_size(ctx));
	}

	return no_error;
}

enum cmix_error get_pub_key_hash(struct CMixContext const* ctx, char** buffer, size_t* len) {
	ctx->api.get_pub_key_hash(buffer, len, ctx->keypair.pub);
	
	return no_error;
}

size_t get_pub_key_hash_length(struct CMixContext const* ctx)
{
	return ctx->api.get_pub_key_hash_length();
}

enum cmix_error create_message(struct CMixContext const* ctx, char* out_buffer, const char* dest, size_t dest_len, const char* payload, size_t payload_len)
{
	size_t el_arr_len = ctx->nr_mixes * get_message_size(ctx);
	
	memcpy(out_buffer, dest, dest_len);
	
	size_t message_limit = el_arr_len - dest_len;
	
	memcpy(out_buffer + dest_len, payload, (message_limit < payload_len ? message_limit : payload_len));
	
	return no_error;
}

enum cmix_error split_message(struct CMixContext const* ctx, char** dest_buffer, size_t* dest_len, char** payload_buffer, size_t* payload_len, char const* message)
{
	size_t message_len = get_message_size(ctx);
	*dest_len = get_pub_key_hash_length(ctx);
	
	*dest_buffer = (char*) calloc(*dest_len, sizeof(char));
	if(*dest_buffer == NULL) {
		return out_of_memory;
	}
	memcpy(*dest_buffer, message, *dest_len);
	
	*payload_len = ctx->nr_mixes * message_len - *dest_len;
	*payload_buffer = (char*) calloc(*payload_len, sizeof(char));
	if(*payload_buffer == NULL) {
		return out_of_memory;
	}

	memcpy(*payload_buffer, &message[*dest_len], *payload_len);
	
	return no_error;
}