#include "cmix.h"

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

enum cmix_error permutation(struct CMixBuffer b) {
	return no_error;
}

enum cmix_error get_message(char* message, struct CMixBuffer b, unsigned int index)
{
	if(index >= b.message_length) {
		return index_out_of_range;
	}
	
	strncpy(message, b.buffer + index * b.message_length, b.message_length);

	return no_error;
}

enum cmix_error set_message(char const* message, struct CMixBuffer b, unsigned int index)
{
	if(index >= b.message_length) {
		return index_out_of_range;
	}
	
	strncpy(b.buffer + index * b.message_length, message, b.message_length);
	
	return no_error;
}

enum cmix_error calculate_shared_key_part(struct Bignum* result, struct Bignum partial_shared, struct Bignum my_share, struct Bignum mod)
{
	if(multiply_mod(result, partial_shared, my_share, mod) != NoError) {
		return cmix_bignum_error;
	}
	return no_error;
}

struct CMixContext initialize_cmix_context(struct Api api) {
	return (struct CMixContext){
		.api = api,
		.nr_participants = 0,
		.r = NULL,
		.s = NULL,
		.permutation = NULL,
		.pirs = NULL,
	};
}

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->pirs[i]);
	}
	free(ctx->r);
	free(ctx->s);
	free(ctx->permutation);
	free(ctx->pirs);
}

void deinitialize(struct CMixContext* ctx)
{
	ctx->api.free_keypair(&ctx->keypair);
	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);
}

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->pirs = (GroupElement*) calloc(ctx->nr_participants, sizeof(GroupElement));
	if(!ctx->r || !ctx->s || !ctx->permutation || !ctx->pirs) {
		return out_of_memory;
	}
	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) {
	ctx->permutation = (unsigned int*) calloc(ctx->nr_participants, sizeof(unsigned int));
	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();
}

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

enum cmix_error encrypt_r(struct CMixContext const* ctx, char* random_buffer, char* message_buffer, size_t index) {
	GroupElement random_element;
	GroupElement message_element;
	
	ctx->api.encrypt(&random_element, &message_element, ctx->r[index], ctx->network_key);
	
	element_to_buffer(ctx, random_buffer, random_element);
	element_to_buffer(ctx, message_buffer, 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, size_t index) {
	GroupElement enc_random_element;
	GroupElement enc_message_element;
	
	ctx->api.encrypt(&enc_random_element, &enc_message_element, ctx->r[index], ctx->network_key);
	
	size_t el_size = get_group_element_array_size(ctx);
	
	GroupElement other_random_element = ctx->api.array_to_element(random_element, el_size, true);
	GroupElement other_message_element = ctx->api.array_to_element(message_element, el_size, true);
	
	GroupElement new_random_element = ctx->api.multiply(enc_random_element, other_random_element, true);
	GroupElement new_message_element = ctx->api.multiply(enc_message_element, other_message_element, true);
	
	element_to_buffer(ctx, random_buffer, new_random_element);
	element_to_buffer(ctx, message_buffer, 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 nr_elements) {
	size_t el_size = get_group_element_array_size(ctx);
	
	for(size_t i = 0; i < nr_elements; ++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.multiply(random_r, random_s, true);
		GroupElement message_pirs = ctx->api.multiply(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);
	}
	
	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.multiply(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*)pubkey, el_len, (unsigned char*)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*)pubkey, el_len, (unsigned char*)value, el_len, *priv_el, false);
	
	ctx->api.free_group_element(*priv_el);
	*priv_el = NULL;
	
	return no_error;
}

enum cmix_error post_process(struct CMixContext* ctx, char* r_out, char* m_out, char const* r_epirs, char const* m_epirs, size_t index) {
	GroupElement x = ctx->api.array_to_element(r_epirs, get_group_element_array_size(ctx), true);
	GroupElement D = ctx->api.get_decryption_share(x, ctx->keypair.sec);
	//element_to_buffer(ctx, r_out, D);
	
	GroupElement msg = ctx->api.array_to_element(m_epirs, get_group_element_array_size(ctx), true);
	GroupElement pirs = ctx->api.multiply(D, msg, true);
	element_to_buffer(ctx, m_out, pirs);
	//GroupElement new_r = ctx->api.multiply(x, D, true);
	//element_to_buffer(ctx, r_out, new_r);
	
	memcpy(r_out, r_epirs, get_group_element_array_size(ctx));
	
	ctx->pirs[index] = 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);
	//ctx->api.free_group_element(new_r);
	
	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_group_element_array_size(ctx); 
	
	GroupElement* intermediates = (GroupElement*) calloc(nr_nodes + 1, sizeof(GroupElement));
	
	intermediates[0] = ctx->api.array_to_element(message, len, true);
	
	for(size_t i = 0; i < nr_nodes; ++i) {
		intermediates[i+1] = ctx->api.multiply(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, char const* message, GroupElement const key, size_t index) {
	size_t len = get_group_element_array_size(ctx);
	
	GroupElement mes = ctx->api.array_to_element(message, len, false);
	GroupElement inv_key = ctx->api.invert(key);
	
	GroupElement unblinded = ctx->api.multiply(mes, inv_key, false);
	GroupElement blinded = ctx->api.multiply(unblinded, ctx->r[index], false);
	
	element_to_buffer(ctx, out_buffer, 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 remove_r_and_s(struct CMixContext const* ctx, char* out_buffer, char const* message, size_t index) {
	size_t len = get_group_element_array_size(ctx);
	
	GroupElement mes = ctx->api.array_to_element(message, len, false);
	GroupElement inv_pirs = ctx->api.invert(ctx->pirs[index]);
	
	GroupElement mult = ctx->api.multiply(mes, inv_pirs, false);
	
	element_to_buffer(ctx, out_buffer, mult);
	
	ctx->api.free_group_element(mes);
	ctx->api.free_group_element(inv_pirs);
	ctx->api.free_group_element(mult);
}

enum cmix_error get_pub_key_hash(struct CMixContext const* ctx, char** buffer)
{
	char* buffer2;
	size_t len;
	ctx->api.get_pub_key_hash(&buffer2, &len, ctx->keypair.pub);
	
	size_t hex_len = len*2+1;
	*buffer = (char*) calloc(hex_len, sizeof(char));
	
	for(size_t i = 0; i < len; i++) {
		sprintf(&(*buffer)[i*2],"%hhX",buffer2[i]);
	}
	
	ctx->api.free_buffer(buffer2);
	
	return no_error;
}