#include "rsa.h"

void bigint_set_random_bytes(bigint_t n) {
	int fd = open("/dev/urandom", O_RDONLY);
	read(fd, n.data, n.len * sizeof(uint32_t));
	close(fd);
}

void bigint_set_msb_and_lsb_to_one(bigint_t n) {
	n.data[0] |= 1;
	n.data[n.len - 1] |= 1 << 31;
}

void bigint_bitwise_right_shift(bigint_t n) {
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (size_t i = 0; i < n.len - 1; i++) {
		n.data[i] = n.data[i] >> 1 | (n.data[i + 1] & 1) << size;
	}
	n.data[n.len - 1] >>= 1;
}

void bigint_bitwise_left_shift(bigint_t n) {
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (int i = n.len - 1; i > 0; i--) {
		n.data[i] = n.data[i] << 1 | ((n.data[i - 1] & (1 << size)) >> size);
	}
	n.data[0] <<= 1;
}

bigint_t assignable_bigint_bitwise_left_shift(bigint_t n) {
	bigint_t result = bigint_clone(n);
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (int i = result.len - 1; i > 0; i--) {
		result.data[i] = result.data[i] << 1 | ((result.data[i - 1] & (1 << size)) >> size);
	}
	result.data[0] <<= 1;
	return result;
}

// Will underflow
void bigint_decrement(bigint_t n) {
	size_t cursor = 0;
	size_t size = sizeof(uint32_t) * 8;
	while (cursor < size * n.len) {
		n.data[cursor / size] = n.data[cursor / size] ^ (1 << (cursor % size));
		if (((n.data[cursor / size] >> (cursor % size)) & 1) == 0) {
			return;
		}
		cursor += 1;
	}
}

// TODO refactor/clean assume same length ?
int bigint_cmp(bigint_t a, bigint_t b) {
	uint32_t size = sizeof(uint32_t) * 8;
	uint32_t acursor = size * a.len - 1;
	uint32_t bcursor = size * b.len - 1;
	while (acursor > bcursor) {
		if (a.data[acursor / size] & (1 << acursor % size)) {
			return 1;
		}
		acursor -= 1;
	}
	while (bcursor > acursor) {
		if (b.data[bcursor / size] & (1 << bcursor % size)) {
			return -1;
		}
		bcursor -= 1;
	}
	int cursor = acursor;
	while (cursor >= 0) {
		uint32_t abit = a.data[cursor / size] & (1 << (cursor % size));
		uint32_t bbit = b.data[cursor / size] & (1 << (cursor % size));
		//printf("cursor %d abit %ud bbit %ud\n", cursor, abit, bbit);
		if (abit > bbit) {
			return 1;
		}
		if (bbit > abit) {
			return -1;
		}
		cursor -= 1;
	}
	return 0;
}

// TODO check opti
bigint_t assignable_bigint_substraction(bigint_t a, bigint_t b) {
	if (a.len != b.len) {
		printf("error: attempting to substract numbers of different length\n");
		exit(1);
	}

	bigint_t result = bigint_clone(a);
	bigint_t borrow = bigint_clone(b);
	bigint_t y = bigint_clone(b);
	bigint_t zero = bigint_zero(a.len);
	while (bigint_cmp(borrow, zero)) {
		for (size_t i = 0; i < a.len; i++) {
			borrow.data[i] = ~result.data[i] & y.data[i];
			result.data[i] = result.data[i] ^ y.data[i];
		}
		bigint_destroy(y);
		y = assignable_bigint_bitwise_left_shift(borrow);
	}
	bigint_destroy(y);
	bigint_destroy(borrow);
	bigint_destroy(zero);
	return result;
}

// TODO check opti
void bigint_substraction(bigint_t a, bigint_t bb) {
	bigint_t b = bigint_clone(bb);
	if (a.len > bb.len) {
		b = bigint_zero(a.len);
		memcpy(b.data, bb.data, b.len * sizeof(uint32_t));
	}
	if (a.len != b.len) {
		printf("error: attempting to substract numbers of different length\n");
		exit(1);
	}

	bigint_t borrow = bigint_clone(b);
	bigint_t y = bigint_clone(b);
	bigint_t zero = bigint_zero(a.len);
	while (bigint_cmp(borrow, zero)) {
		for (size_t i = 0; i < a.len; i++) {
			borrow.data[i] = ~a.data[i] & y.data[i];
			a.data[i] = a.data[i] ^ y.data[i];
		}
		bigint_destroy(y);
		y = assignable_bigint_bitwise_left_shift(borrow);
	}
	bigint_destroy(y);
	bigint_destroy(borrow);
	bigint_destroy(zero);
}

// TODO check opti
bigint_t assignable_bigint_modulo(bigint_t a, bigint_t b) {
	bigint_t result = bigint_clone(a);
	bigint_t mod = bigint_clone(b);
	if (a.len > b.len) {
		mod = bigint_zero(a.len);
		memcpy(mod.data, b.data, b.len * sizeof(uint32_t));
	}
	if (bigint_cmp(result, b) == -1) {
		bigint_destroy(mod);
		return result;
	}
	bigint_bitwise_left_shift(mod);
	while (bigint_cmp(b, mod) == -1) {
		while (bigint_cmp(result, mod) == 1) {
			bigint_bitwise_left_shift(mod);
		}
		bigint_bitwise_right_shift(mod);
		if (bigint_cmp(result, mod) == 1) {
			bigint_substraction(result, mod);
		}
	}
	while (bigint_cmp(result, b) == 1) {
		bigint_substraction(result, b);
	}
	bigint_destroy(mod);
	return result;
}

void bigint_add(bigint_t a, bigint_t b) {
	bigint_t result = bigint_zero(a.len);
	size_t size = sizeof(uint32_t) * 8;
	size_t width = a.len * size;
	uint32_t carriage = 0;

//	printf("hello add\n");
	for (size_t cursor = 0; cursor < width; cursor++) {
//		printf("hahaha %ld %ld\n", cursor, width);
		uint32_t a_bit = a.data[cursor / size] >> (cursor % size) & 1;
		uint32_t b_bit = b.data[cursor / size] >> (cursor % size) & 1;
		result.data[cursor / size] |= (a_bit ^ b_bit ^ carriage) << (cursor % size);
		carriage = (a_bit & b_bit) | ((a_bit ^ b_bit) & carriage);
	}
//	printf("im out\n");
	bigint_destroy(a);
	a = bigint_clone(result);
	bigint_destroy(result);
}

void bigint_set_zeros(bigint_t n) {
//	printf("hello set zeros\n");
	for (size_t i = 0; i < n.len; i++) {
		n.data[i] = 0;
	}
//	printf("goodbye set zeros\n");
}

bigint_t assignable_bigint_mul(bigint_t a, bigint_t b) {
	bigint_t result = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 4);
	bigint_t b_tool = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 4);
	/*if (a.len > b.len) {
		result = bigint_zero(a.len);
		b_tool = bigint_zero(a.len);
	} else {
		result = bigint_zero(a.len + b.len);
		b_tool = bigint_zero(a.len + b.len);
	}*/
	size_t size = sizeof(uint32_t) * 8;
	size_t width = a.len * size;

	printf("multiplying %d and %d\n", a.data[0], b.data[0]);

//	printf("hello mul\n");
	for (size_t cursor = 0; cursor < width; cursor++) {
//		printf("hello BIG LOOP ls %ld %ld\n", cursor, width);
		if (a.data[cursor / 32] >> (cursor % 32) & 1) {
			bigint_set_zeros(b_tool);
			printf("bef %d\n", b_tool.data[0]);
//			printf("hello memcpy\n");
			memcpy(b_tool.data, b.data, b.len * sizeof(uint32_t));
			printf("aft %d\n", b_tool.data[0]);
//			printf("goodbye memcpy\n");
			for (size_t i = 0; i < cursor; i++) {
//				printf("hello bitwise ls %ld %ld\n", i, cursor);
				bigint_bitwise_left_shift(b_tool);
//				printf("goodbye bitwise ls\n");
			}
//			printf("before hello add\n");
			bigint_add(result, b_tool);
		}
	}
//	printf("GOODBYE BIG LOOP ls \n");
	bigint_destroy(b_tool);
	return result;
}

// a^e mod n
// clean memory tricks !!!
bigint_t assignable_bigint_pow_mod(bigint_t a, bigint_t e, bigint_t n) {
	printf("print a\n");
	bigint_print(a);
	printf("print e\n");
	bigint_print(e);
	printf("print n\n");
	bigint_print(n);
	bigint_t result = bigint_clone(a);
	size_t size = sizeof(uint32_t) * 8;
	int cursor = e.len * size - 1;
	while (!(e.data[cursor / 32] & 1 << (cursor % 32))) {
		cursor--;
	}
	cursor--;
	/*
		printf("SQUARE\n");
		bigint_t tmp_result2 = assignable_bigint_mul(result, result);
		bigint_destroy(result);
		result = bigint_clone(tmp_result2);
		bigint_destroy(tmp_result2);
		tmp_result2 = assignable_bigint_modulo(result, n);
		bigint_destroy(result);
		result = bigint_clone(tmp_result2);
		bigint_destroy(tmp_result2);
		*/
	printf("cursor %d\n", cursor);
	while (cursor >= 0) {
		printf("SQUARE\n");
		bigint_t tmp_result2 = assignable_bigint_mul(result, result);
		bigint_destroy(result);
		result = bigint_clone(tmp_result2);
		bigint_destroy(tmp_result2);
		tmp_result2 = assignable_bigint_modulo(result, n);
		bigint_destroy(result);
		result = bigint_clone(tmp_result2);
		bigint_destroy(tmp_result2);
		if (e.data[cursor / 32] & 1 << (cursor % 32)) {
		printf("MULTIPLY\n");
			bigint_t tmp_result = assignable_bigint_mul(result, a);
			bigint_destroy(result);
			result = bigint_clone(tmp_result);
			bigint_destroy(tmp_result);
			tmp_result = assignable_bigint_modulo(result, n);
			bigint_destroy(result);
			result = bigint_clone(tmp_result);
			bigint_destroy(tmp_result);
		}
		cursor -= 1;
	}
	printf("this time its over\n");
	return result;
}

void bigint_print(bigint_t n) {
	for (int i = n.len - 1; i >= 0; i--) {
		printf("bigint %ud\n", n.data[i]);
	}
}

bigint_t bigint_clone(bigint_t src) {
	bigint_t dst;

	dst.len = src.len;
	dst.data = (uint32_t *)protected_malloc(src.len * sizeof(uint32_t));
	memcpy(dst.data, src.data, src.len * sizeof(uint32_t));

	return dst;
}

bigint_t bigint_prime(size_t len) {
	bigint_t n = bigint_new(len);

	printf("new\n");
	bigint_print(n);
	bigint_set_random_bytes(n);
	printf("random bytes\n");
	bigint_print(n);
	bigint_set_msb_and_lsb_to_one(n);
	printf("msb and lsb set to tone\n");
	bigint_print(n);

	bigint_t d = bigint_clone(n);
	d.data[0] -= 1;
	uint32_t s = 0;
	while (!d.data[0] & 1) {
		bigint_bitwise_right_shift(d);
		s += 1;
	}

	bigint_t two = bigint_zero(len);
	two.data[0] = 2;
	bigint_t n_minus_two = bigint_clone(d);
	bigint_decrement(n_minus_two);

	for (uint32_t k = 0; k < 128; k++) {
		bigint_t a = bigint_zero(len);
		while (bigint_cmp(a, two) == -1 || bigint_cmp(a, n_minus_two) == 1) {
			printf("this is good %d\n", k);
			bigint_set_random_bytes(a);
		}
		bigint_destroy(a);
	}
	return n;
}

bigint_t bigint_new(size_t len) {
	bigint_t bigint;

	bigint.len = len;
	bigint.data = (uint32_t *)protected_malloc(len * sizeof(uint32_t));

	return bigint;
}

bigint_t bigint_zero(size_t len) {
	bigint_t bigint;

	bigint = bigint_new(len);
	for (size_t i = 0; i < len; i++) {
		bigint.data[i] = 0;
	}

	return bigint;
}

void bigint_destroy(bigint_t n) {
	free(n.data);
	n.data = NULL;
}