#include "cmix.h"

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

struct CMixContext initialize_cmix_context(struct Api api) {
	return (struct CMixContext){
		.api = api,
		.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));
	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) {
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		ctx->api.free_group_element(ctx->r[i]);
		ctx->api.free_group_element(ctx->s[i]);
		ctx->api.free_group_element(ctx->decryption_shares[i]);
		ctx->api.free_group_element(ctx->EPiRS[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(unsigned int i = 0; i < ctx->nr_participants; ++i) {
		ctx->permutation[i] = 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];
		ctx->permutation[i] = ctx->permutation[rand];
		ctx->permutation[rand] = temp;
	}

	return no_error;
}

enum cmix_error initialize_mix_randomness(struct CMixContext* ctx) {
	for(size_t i = 0; i < ctx->nr_participants; ++i) {
		ctx->r[i] = ctx->api.get_group_element(true);
		ctx->s[i] = 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) {
		GroupElement random_element;
		GroupElement message_element;
	
		ctx->api.encrypt(&random_element, &message_element, ctx->r[i], ctx->network_key);
	
		element_to_buffer(ctx, random_buffer[i], random_element);
		element_to_buffer(ctx, message_buffer[i], 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) {
		GroupElement enc_random_element;
		GroupElement enc_message_element;
		
		ctx->api.encrypt(&enc_random_element, &enc_message_element, ctx->r[i], 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], el_size, true);
		GroupElement other_message_element = ctx->api.array_to_element(message_element[i], 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], new_random_element);
		element_to_buffer(ctx, message_buffer[i], 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) {
		unsigned int new_pos = ctx->permutation[i];
		GroupElement random_r = ctx->api.array_to_element(random_element[i], el_size, true);
		GroupElement message_r = ctx->api.array_to_element(message_element[i], el_size, true);
	
		GroupElement random_s;
		GroupElement message_s;
		
		ctx->api.encrypt(&random_s, &message_s, ctx->s[new_pos], 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], random_pirs);
		element_to_buffer(ctx, message_buffer[new_pos], 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) {
		unsigned int new_pos = ctx->permutation[i];
		
		GroupElement message_el = ctx->api.array_to_element(message[i], el_size, false);
		ctx->PiMRS[new_pos] = ctx->api.combine(message_el, ctx->s[new_pos], false);
		
		element_to_buffer(ctx, out_buffer[new_pos], ctx->PiMRS[new_pos]);
		
		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) {
	size_t el_size = get_group_element_array_size(ctx);
	
	GroupElement message_el = ctx->api.array_to_element(message, el_size, false);
	GroupElement mult = ctx->api.combine(message_el, ctx->s[ctx->permutation[index]], false);
	
	element_to_buffer(ctx, out_buffer, 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) {
		GroupElement x = ctx->api.array_to_element(r_epirs[i], get_group_element_array_size(ctx), true);
		ctx->decryption_shares[i] = ctx->api.get_decryption_share(x, ctx->keypair.sec);
		
		ctx->EPiRS[i] = ctx->api.array_to_element(m_epirs[i], 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;
	}
	
	intermediates[0] = ctx->api.message_to_element(message, len, true);
	
	for(size_t i = 0; i < nr_nodes; ++i) {
		intermediates[i+1] = ctx->api.combine(intermediates[i], keys[i], false);
	}
	
	element_to_buffer(ctx, m_out, intermediates[nr_nodes]);
	
	for(size_t i = 0; i < nr_nodes + 1; ++i) {
		ctx->api.free_group_element(intermediates[i]);
	}
	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) {
		GroupElement mes = ctx->api.array_to_element(message[i], len, false);
		GroupElement inv_key = ctx->api.invert(key[i]);
		
		GroupElement unblinded = ctx->api.uncombine(mes, key[i], false);
		GroupElement blinded = ctx->api.combine(unblinded, ctx->r[i], false);
		
		element_to_buffer(ctx, out_buffer[i], blinded);
		ctx->api.free_group_element(mes);
		ctx->api.free_group_element(inv_key);
		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) {
		element_to_buffer(ctx, out_buffer[i], ctx->EPiRS[i]);
	}
	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) {
		GroupElement msg = ctx->api.array_to_element(epirs[i], get_group_element_array_size(ctx), true);
		GroupElement pirs = ctx->api.combine(ctx->decryption_shares[i], msg, true);
		element_to_buffer(ctx, out_buffer[i], 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) {
		GroupElement g_pirs = ctx->api.array_to_element(pirs[i], len, false);
		
		GroupElement mult = ctx->api.uncombine(ctx->PiMRS[i], g_pirs, false);
		
		element_to_buffer(ctx, out_buffer[i], 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) {
	
	GroupElement el = ctx->api.array_to_element(element_buffer, ctx->api.get_group_element_array_size(), true);
	ctx->api.element_to_message(buffer, el);
	ctx->api.free_group_element(el);
	
	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 = 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 = 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;
}