#include "ed25519.h"

#include "api.h"
#include "impl_macros.h"

#include "gcrypt.h"

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <assert.h>
#include <math.h>

#ifdef POINTER_TRACING
	#define TRACE(FILE, LINE, PTR) \
		fprintf(stderr, "%s %d %p\n", FILE, LINE, PTR);
#else
	#define TRACE(FILE, LINE, PTR)
#endif

static gcry_ctx_t ctx;

static gcry_mpi_point_t g;
static gcry_mpi_t p;

static unsigned int nr_bits = 256;
static unsigned int nr_point_bits = 256 * 2;
static unsigned int nr_bytes = 256 / 8;
static unsigned int stride = 16;

void check(gcry_error_t error) {
	if (error) {
		fprintf (stderr, "Error: %s/%s\n", gcry_strsource (error), gcry_strerror (error));
		exit(1);
	}
}

void print_sexp(gcry_sexp_t exp) {
	size_t required_size = gcry_sexp_sprint(exp, GCRYSEXP_FMT_ADVANCED, NULL, 0); //passing NULL as the buffer will return the size needed.

	char* str = (char *) malloc(sizeof(char)*required_size);

	gcry_sexp_sprint(exp, GCRYSEXP_FMT_ADVANCED, str, required_size); //So we ignore the size here.
	
	printf("size = %zu\n", required_size);
	for(size_t i = 0; i < required_size; i++)
		printf("%c", str[i]);
	printf("\n");

	free(str);
}

void DEF(initialize)(void) {
	if (!gcry_check_version (GCRYPT_VERSION)) {
		fprintf(stderr, "libgcrypt version mismatch\n");
		exit(-1);
	}

	gcry_control (GCRYCTL_SUSPEND_SECMEM_WARN);
	gcry_control (GCRYCTL_INIT_SECMEM, 32768, 0);
	gcry_control (GCRYCTL_RESUME_SECMEM_WARN);
	gcry_control (GCRYCTL_INITIALIZATION_FINISHED, 0);
	
	gcry_error_t error;
	
	error = gcry_mpi_ec_new(&ctx, NULL, "Ed25519");
	check(error);
	
	g = gcry_mpi_ec_get_point("g", ctx, false);
	p = gcry_mpi_ec_get_mpi("p", ctx, false);
	
	TRACE(__FILE__, __LINE__, g)
	TRACE(__FILE__, __LINE__, p)
}

struct KeyPair DEF(create_keypair)(void) {
	size_t parse_error_offset;
	gcry_error_t error;
	
	gcry_sexp_t key_gen_sexp;
	error = gcry_sexp_build(&key_gen_sexp, &parse_error_offset, "(genkey (ecc (curve Ed25519)))", NULL);
	check(error);
	
	gcry_sexp_t key_pair;
	error = gcry_pk_genkey(&key_pair, key_gen_sexp);
	check(error);
	
	gcry_mpi_t d;
	gcry_mpi_t q_mpi;
	
	gcry_sexp_extract_param(key_pair, "key-data!private-key", "q d", &q_mpi, &d, NULL);
	
	gcry_mpi_point_t q = gcry_mpi_point_new(nr_point_bits);
	error = gcry_mpi_ec_decode_point(q, q_mpi, ctx);
	check(error);
	
	assert(gcry_mpi_ec_curve_point(q, ctx));
	
	gcry_mpi_release(q_mpi);
	gcry_sexp_release(key_gen_sexp);
	gcry_sexp_release(key_pair);
	
	TRACE(__FILE__, __LINE__, d)
	TRACE(__FILE__, __LINE__, q)
	
	return (struct KeyPair) {
		.sec = d,
		.pub = q
	};
}

void DEF(delete_keypair)(struct KeyPair* pair) {
	gcry_mpi_release((gcry_mpi_t)pair->sec);
	gcry_mpi_point_release((gcry_mpi_point_t)pair->pub);
	pair->sec = NULL;
	pair->pub = NULL;
}

void DEF(element_to_array)(unsigned char** buffer, size_t* len, GroupElement element) {
	gcry_error_t error;
	
	gcry_mpi_t x = gcry_mpi_new(nr_bits);
	gcry_mpi_t y = gcry_mpi_new(nr_bits);
	
	int ret = gcry_mpi_ec_get_affine(x, y, (gcry_mpi_point_t) element, ctx);
	
	if(ret != 0) {
		abort();
	}
	
	*len = nr_bytes * 2;
	*buffer = (unsigned char*) calloc(*len, sizeof(unsigned char));
	
	TRACE(__FILE__, __LINE__, buffer)
	
	unsigned char* temp = (unsigned char*) calloc(nr_bytes, sizeof(unsigned char));
	size_t nwritten = 0;
	
	error = gcry_mpi_print(GCRYMPI_FMT_USG, temp, nr_bytes, &nwritten, (gcry_mpi_t) x);
	check(error);
	memcpy(*buffer + (nr_bytes - nwritten), temp, nwritten);

	memset(temp, '\0', nr_bytes);
	error = gcry_mpi_print(GCRYMPI_FMT_USG, temp, nr_bytes, &nwritten, (gcry_mpi_t) y);
	check(error);
	memcpy(*buffer + nr_bytes + (nr_bytes - nwritten), temp, nwritten);
	
	free(temp);
	gcry_mpi_release(y);
	gcry_mpi_release(x);
}

GroupElement DEF(message_to_element)(char const* buffer, size_t len, bool secure) {
	(void) secure;
	gcry_error_t error;
	
	assert(len == 31);
	
	gcry_mpi_t mpi;
	gcry_mpi_scan(&mpi, GCRYMPI_FMT_USG, buffer, len, NULL);
	
	//mpi * stride
	gcry_mpi_lshift(mpi, mpi, (unsigned) lround(log2(stride)));
	
	char* encoded_point = (char*) calloc(nr_bytes+1, sizeof(char));
	gcry_mpi_point_t point = gcry_mpi_point_new(0);
	unsigned int count = 0;
	do {
		encoded_point[0] = 0x40;
		memset(encoded_point + 1, '\0', 32);
	
		unsigned char* gcry_buffer;
		size_t nr_written;

		gcry_mpi_aprint(GCRYMPI_FMT_USG, &gcry_buffer, &nr_written, mpi);
		
		memcpy(encoded_point + 1 + (nr_bytes - nr_written), gcry_buffer, nr_written);
		gcry_free(gcry_buffer);
		
		for(int i = 0; i < 16; i++) {
			char temp = encoded_point[i+1];
			encoded_point[i+1] = encoded_point[32-i];
			encoded_point[32-i] = temp;
		}
	
		gcry_mpi_t encoded_mpi = gcry_mpi_new(0);
		gcry_mpi_set_opaque_copy(encoded_mpi, encoded_point, (nr_bytes + 1) * sizeof(char) * 8);
		
		error = gcry_mpi_ec_decode_point(point, encoded_mpi, ctx);
		check(error);
		
		gcry_mpi_release(encoded_mpi);
	
		assert(count < stride);
		gcry_mpi_add_ui(mpi, mpi, 1);
		count++;
		
		if(!gcry_mpi_ec_curve_point(point, ctx)) {
			gcry_mpi_point_release(point);
			point = gcry_mpi_point_new(0);
		}
		
	} while(!gcry_mpi_ec_curve_point(point, ctx));
	
	gcry_mpi_release(mpi);
	free(encoded_point);
	
	TRACE(__FILE__, __LINE__, point)
	
	return point;
}

size_t DEF(get_message_size)(void);

void DEF(element_to_message)(unsigned char** buffer, const GroupElement el) {
	gcry_mpi_t y = gcry_mpi_new(0);
	gcry_mpi_ec_get_affine(NULL, y, (gcry_mpi_point_t) el, ctx);
	
	gcry_mpi_rshift(y, y, (unsigned) lround(log2(stride)));
	
	size_t mes_len = DEF(get_message_size)();
	*buffer = (unsigned char*) calloc(mes_len, sizeof(unsigned char));
	
	TRACE(__FILE__, __LINE__, buffer)
	
	unsigned char* buffer2;
	size_t nr_written;
	gcry_mpi_aprint(GCRYMPI_FMT_USG, &buffer2, &nr_written, y);
	
	memcpy(*buffer + (mes_len - nr_written), buffer2, nr_written);
	
	gcry_free(buffer2);
	gcry_mpi_release(y);
}

void DEF(free_buffer)(void* buffer) {
	gcry_free(buffer);
}

GroupElement DEF(array_to_element)(char const* buffer, size_t len, bool secure) {
	(void) secure;
	size_t error_pos;
	gcry_error_t error;
	
	gcry_mpi_t x;
	error = gcry_mpi_scan(&x, GCRYMPI_FMT_USG, buffer, len/2, &error_pos);
	check(error);
	
	gcry_mpi_t y;
	error = gcry_mpi_scan(&y, GCRYMPI_FMT_USG, buffer + len/2, len/2, &error_pos);
	
	gcry_mpi_point_t point = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_point_set(point, x, y, GCRYMPI_CONST_ONE);
	
	gcry_mpi_release(x);
	gcry_mpi_release(y);
	
	TRACE(__FILE__, __LINE__, point)
	
	return point;
}

size_t DEF(get_pub_key_hash_length)(void) {
	return 20;
}

//TODO: fix this up
void DEF(get_pub_key_hash)(char** buffer, size_t* len, GroupElement const pub) {
	(void)pub;
	//gcry_sexp_t pubkey;
	//gcry_error_t error;
	//size_t err_off;
	
	//error = gcry_sexp_build(&pubkey, &err_off, "(public-key (ecc (curve Curve25519) (flags eddsa) (q %M)))", (gcry_mpi_point_t) pub);
	//check(error);
	
	//*buffer = (char*) gcry_pk_get_keygrip(pubkey, NULL);
	//TODO: see if we can instantiate a public key sexp;
	
	gcry_mpi_t x = gcry_mpi_new(0);
	gcry_mpi_t y = gcry_mpi_new(0);
	
	gcry_mpi_ec_get_affine(x, y, (gcry_mpi_point_t)pub, ctx);
	
	*len = DEF(get_pub_key_hash_length)();
	*buffer = (char*) calloc(*len, sizeof(char));
	
	unsigned char* temp;
	size_t len2;
	DEF(element_to_array(&temp, &len2, pub));
	
	if(len2 < *len) {
		memcpy((*buffer) + (*len - len2), temp, len2);
	} else {
		memcpy(*buffer, temp, *len);
	}
	gcry_free(temp);
	
	TRACE(__FILE__, __LINE__, *buffer);
	
	gcry_mpi_release(x);
	gcry_mpi_release(y);
}

GroupElement DEF(get_group_element)(bool secure) {
	(void) secure;
	size_t parse_error_offset;
	gcry_error_t error;
	
	gcry_mpi_t a = NULL;
	void* bytes = NULL;
	do {
		gcry_free(bytes);
		gcry_mpi_release(a);
		bytes = gcry_random_bytes_secure(nr_bytes, GCRY_VERY_STRONG_RANDOM);
		TRACE(__FILE__, __LINE__, bytes);
		error = gcry_mpi_scan(&a, GCRYMPI_FMT_USG, bytes, nr_bytes, &parse_error_offset);
		TRACE(__FILE__,__LINE__, a);
		check(error);
	} while(gcry_mpi_cmp_ui(a, 0) == 0 || gcry_mpi_cmp(a, p) != -1);
	
	char* encoded_point = (char*) calloc(nr_bytes+1, sizeof(char));
	TRACE(__FILE__, __LINE__, encoded_point)
	gcry_mpi_point_t point = gcry_mpi_point_new(0);
	do {
		encoded_point[0] = 0x40;
		
		unsigned int byte_size = (unsigned) ceil(gcry_mpi_get_nbits(a) / 8.0);
		gcry_mpi_print(GCRYMPI_FMT_USG, (unsigned char*) encoded_point + 1, 32, NULL, a);
	
		for(int i = 0; i < 16; i++) {
			char temp = encoded_point[i+1];
			encoded_point[i+1] = encoded_point[32-i];
			encoded_point[32-i] = temp;
		}
	
		gcry_mpi_t encoded_mpi = gcry_mpi_new(0);
		TRACE(__FILE__, __LINE__, encoded_mpi)
		gcry_mpi_set_opaque_copy(encoded_mpi, encoded_point, (nr_bytes + 1) * sizeof(char) * 8);
		
		error = gcry_mpi_ec_decode_point(point, encoded_mpi, ctx);
		check(error);
			
		gcry_mpi_release(encoded_mpi);
		
		if(!gcry_mpi_ec_curve_point(point, ctx)) {
			gcry_mpi_point_release(point);
			point = gcry_mpi_point_new(0);
			
			gcry_mpi_add_ui(a, a, 1);
		}
	
	} while(!gcry_mpi_ec_curve_point(point, ctx));
	
	gcry_mpi_release(a);
	gcry_free(bytes);
	free(encoded_point);
	
	TRACE(__FILE__, __LINE__, point);
	
	return point;
}

GroupElement DEF(get_key_exchange_value)(GroupElement group_el) {	
	gcry_mpi_t x = gcry_mpi_new(nr_bits);
	gcry_mpi_point_get(x, NULL, NULL, (gcry_mpi_point_t) group_el);
	
	gcry_mpi_point_t mult = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_mul(mult, x, g, ctx);
	
	gcry_mpi_release(x);
	
	TRACE(__FILE__, __LINE__, mult)
	
	return mult;
}

GroupElement DEF(combine)(GroupElement lh, GroupElement rh, bool secure) {
	(void) secure;
	gcry_mpi_point_t addition = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_add(addition, (gcry_mpi_point_t)lh, (gcry_mpi_point_t)rh, ctx);
	
	TRACE(__FILE__, __LINE__, addition)
	
	return addition;
}

GroupElement DEF(invert)(GroupElement const el) {
	gcry_mpi_t x = gcry_mpi_new(nr_bits);
	gcry_mpi_t y = gcry_mpi_new(nr_bits);
	
	gcry_mpi_ec_get_affine(x, y, (gcry_mpi_point_t) el, ctx);
	
	gcry_mpi_t neg_x = gcry_mpi_new(nr_bits);
	gcry_mpi_subm(neg_x, p, x, p);
	
	gcry_mpi_point_t point = gcry_mpi_point_set(NULL, neg_x, y, GCRYMPI_CONST_ONE);
	
	gcry_mpi_release(neg_x);
	gcry_mpi_release(x);
	gcry_mpi_release(y);
	
	TRACE(__FILE__, __LINE__, point)
	
	return point;
}

void DEF(delete_group_element)(GroupElement element) {
	gcry_mpi_point_release((gcry_mpi_point_t)element);
}

GroupElement DEF(uncombine)(GroupElement lh, GroupElement rh, bool secure) {
	(void) secure;
	
	GroupElement inv_rh = DEF(invert)((gcry_mpi_point_t)rh);
	GroupElement ret = DEF(combine)(lh, inv_rh, secure);
	DEF(delete_group_element)(inv_rh);
	
	TRACE(__FILE__, __LINE__, ret)
	
	return ret;
}

GroupElement DEF(get_decryption_share)(GroupElement r, GroupElement e) {
	gcry_mpi_point_t inv_d = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_mul(inv_d, (gcry_mpi_t)e, (gcry_mpi_point_t) r, ctx);
	
	GroupElement d = DEF(invert)(inv_d);
	gcry_mpi_point_release(inv_d);
	
	TRACE(__FILE__, __LINE__, d)
	return d;
}

size_t DEF(get_group_element_array_size)(void) {
	return nr_bytes*2;
}

size_t DEF(get_message_size)(void) {
	return nr_bytes-1;
}

void DEF(add_public_share)(GroupElement* el, char const* share, size_t in_len, GroupElement pubkey) {
	GroupElement share_el = DEF(array_to_element)(share, in_len, false);
	*el = DEF(combine)(share_el, pubkey, false);
	TRACE(__FILE__, __LINE__, *el);
	gcry_mpi_point_release((gcry_mpi_point_t)share_el);
}

GroupElement DEF(derive_shared_key)(struct KeyPair keypair, unsigned char const* other_pub, size_t pub_len, unsigned char const* value, size_t value_len, void* priv_value, bool swap) {
	(void)keypair;
	(void)other_pub;
	(void)pub_len;
	(void)swap;
	
	GroupElement ga = DEF(array_to_element)((const char*) value, value_len, false);
	
	gcry_mpi_point_t gab = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_mul(gab, (gcry_mpi_t) priv_value, (gcry_mpi_point_t) ga, ctx);
	
	gcry_mpi_point_release((gcry_mpi_point_t) ga);
	
	TRACE(__FILE__, __LINE__, gab);
	
	return gab;
}

void DEF(delete_shared_key)(struct SharedKey* s) {
	gcry_mpi_point_release((gcry_mpi_point_t)s->shared);
	s->shared = NULL;
}

void DEF(encrypt)(GroupElement* random_element, GroupElement* message_element, GroupElement value, GroupElement key) {
	GroupElement random = DEF(get_group_element)(true);
	
	gcry_mpi_t random_mpi = gcry_mpi_new(nr_bits);
	gcry_mpi_ec_get_affine(random_mpi, NULL, (gcry_mpi_point_t)random, ctx);
	
	*random_element = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_mul((gcry_mpi_point_t)*random_element, random_mpi, g, ctx);
	
	gcry_mpi_point_t c2 = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_mul(c2, random_mpi, (gcry_mpi_point_t) key, ctx);
	
	*message_element = gcry_mpi_point_new(nr_point_bits);
	gcry_mpi_ec_add((gcry_mpi_point_t)* message_element, c2, (gcry_mpi_point_t) value, ctx);
	
	gcry_mpi_release(random_mpi);
	gcry_mpi_point_release(c2);
	gcry_mpi_point_release(random);
	
	TRACE(__FILE__, __LINE__, *random_element);
	TRACE(__FILE__, __LINE__, *message_element);
	
}

unsigned int DEF(get_uniform_int)(unsigned int upper) {
	unsigned int random_value;
	
	if(upper == 0) {
		void* buffer = gcry_random_bytes_secure(sizeof(unsigned int), GCRY_STRONG_RANDOM);
		TRACE(__FILE__, __LINE__, buffer);
		memcpy(&random_value, buffer, sizeof(unsigned int));
		gcry_free(buffer);
		
		return random_value;
	}
	
	unsigned int lower = -upper % upper;
	do {
		void* buffer = gcry_random_bytes_secure(sizeof(unsigned int), GCRY_STRONG_RANDOM);
		TRACE(__FILE__, __LINE__, buffer);
		memcpy(&random_value, buffer, sizeof(unsigned int));
		gcry_free(buffer);
		
		if (random_value >= lower)
			return random_value % upper;
		
		
	} while(true);
}

void DEF(deinitialize)(void) {
	gcry_mpi_point_release(g);
	gcry_mpi_release(p);
	gcry_ctx_release(ctx);
}

LINK_IMPLEMENTATION