#include "smg_rsa.h"
#include "mpi.h"

#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <stdio.h>

void err(char *msg)
{
  fprintf(stderr, "%s\n", msg);
  exit(1);
}

void time_entropy_source(int nruns, int noctets) {
	unsigned char buffer[noctets];
	int read, i;
	struct timespec tstart, tend;
	long int diff;

	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i=0; i<nruns; i++) {
		read = get_random_octets(noctets,buffer);
		if (read != noctets)
			err("Failed reading from entropy source!");
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;
	double kbps = (nruns*noctets) / (diff*1000.0);
	printf("ENTROPY source timing: %d kB in %ld seconds, at an average speed of %f kB/s over %d runs of %d octets each\n", nruns*noctets, diff, kbps, nruns, noctets);
}

void test_entropy_output(unsigned int noctets, char * filename) {
	FILE * out;
	int source;
	unsigned int nread, total_read, to_read;
	const int buffer_length = 1000;
	unsigned char buffer[buffer_length];

	source = open_entropy_source(ENTROPY_SOURCE);
	if (source <= 0)
		err("unable to access entropy source!");

	out = fopen(filename, "wb");
	if ( !out )
		err("unable to open output file for test_entropy_output!");

	printf("TEST_ENTROPY_SOURCE: reading %u octets from %s ", noctets, ENTROPY_SOURCE);
	total_read = 0;
	while (total_read < noctets) {
		to_read = noctets - total_read;
		if (to_read > buffer_length)
			to_read = buffer_length;

		nread = get_random_octets_from(to_read, buffer, source);
		if (nread > 0) {
			total_read = total_read + nread;
			fwrite(buffer, 1, nread, out);
			fflush(out);
			printf(".");
			fflush(stdout);
		}
	}
	printf("done.\n");

	fclose(out);
	close(source);
}

void test_is_composite(int nruns, char *hex_number, int expected) {
	int i;
	int output;
	int count_ok = 0;
	int source = open_entropy_source(ENTROPY_SOURCE);
	MPI p = mpi_alloc(0);

	mpi_fromstr(p, hex_number);
	printf("TEST is_composite on MPI(hex) ");
	mpi_print(stdout, p, 1);
	for (i=0; i < nruns; i++) {
		printf(".");
		fflush(stdout);
		output = is_composite(p, M_R_ITERATIONS, source);
		if (output == expected)
			count_ok = count_ok + 1;
	}
	printf("done, with %d out of %d correct runs for expected=%d: %s\n", count_ok, nruns, expected, count_ok==nruns? "PASS":"FAIL");
	mpi_free(p);
	close(source);
}

void time_mr(int nruns) {
	struct timespec tstart, tend;
	long int diff;
	int i;
	MPI prime;
	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);

	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	/* first generate a prime of half key length, to make sure M-R will run max number of iterations */
	printf("Generating a prime number of %d octets length for M-R timing test\n", noctets);
	prime = mpi_alloc(nlimbs);
	gen_random_prime(noctets, prime);

	printf("Running timing test for Miller-Rabin with %d repetitions and %d witnesses on prime number ", nruns, M_R_ITERATIONS);
	mpi_print(stdout, prime, 1);
	printf("\n");
	/* now do the actual runs and time it all */
	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i=0; i<nruns; i++) {
		if (is_composite(prime, M_R_ITERATIONS, entropy_source))
			printf("FAIL");
		else printf(".");
		fflush(stdout);
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;
	printf("\nTimings on prime number %d octets long, %d runs of MR with %d iterations (witnesses checked) each\n", \
		noctets, nruns, M_R_ITERATIONS);
	printf("Total time: %ld seconds\nTime per MR run: %f seconds\nTime per MR iteration: %f seconds\n",\
			 diff, diff / (1.0*nruns), diff / (1.0*nruns * M_R_ITERATIONS));

	mpi_free(prime);
	close(entropy_source);
}

void test_rpng(int nruns) {
	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);
	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	MPI prime = mpi_alloc(nlimbs);
	int i;

	printf("TEST: random prime number generator with %d runs\n", nruns);
	for (i = 0;i < nruns; i++) {
		gen_random_prime(noctets, prime);
		printf("Run %d: ", i+1);
		mpi_print(stdout, prime, 1);
		if (is_composite(prime, M_R_ITERATIONS, entropy_source))
			printf("  **FAIL**\n");
		else
			printf("  **PASS**\n");
	}

	mpi_free(prime);
	close(entropy_source);
}

void time_rpng(int nruns) {
	struct timespec tstart, tend;
	long int diff;

	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);

	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	MPI prime = mpi_alloc(nlimbs);
	int i;

	printf("TIMING: random prime number generator with %d runs\n", nruns);
	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i = 0;i < nruns; i++) {
		gen_random_prime(noctets, prime);
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;

	printf("TOTAL: %ld seconds\n", diff);
	printf("Average: %f seconds to generate one random prime of %d octets length\n", diff / (1.0*nruns), noctets);
	mpi_free(prime);
	close(entropy_source);
}

int main(int ac, char **av)
{
	int nruns;
	int id;

	if (ac<2) {
		printf("Usage: %s number_of_runs/octets [testID]\n", av[0]);
		return -1;
	}
	nruns = atoi(av[1]);

	if (ac < 3)
		id = -1;
	else
		id = atoi(av[2]);

	switch ( id ) {
		case 0:
			printf("Timing entropy source...\n");
			time_entropy_source(nruns, 4096);
			break;
		case 1:
			test_entropy_output(nruns, "entropy_source_output.txt");
			break;
		case 2:
			/* tests on miller-rabin */
			/* a few primes (decimal): 65537, 116447, 411949103, 20943302231 */
			test_is_composite(nruns, "0x10001", 0);
			test_is_composite(nruns, "0x1C6DF", 0);
			test_is_composite(nruns, "0x188DD82F", 0);
			test_is_composite(nruns, "0x4E0516E57", 0);
			/* a few mersenne primes (decimal): 2^13 - 1 = 8191, 2^17 - 1 = 131071, 2^31 - 1 = 2147483647 */
			test_is_composite(nruns, "0x1FFF", 0);
			test_is_composite(nruns, "0x1FFFF", 0);
			test_is_composite(nruns, "0x7FFFFFFF", 0);
			/* a few carmichael numbers, in decimal: 561, 60977817398996785 */
			test_is_composite(nruns, "0x231", 1);
			test_is_composite(nruns, "0xD8A300793EEF31", 1);
			/* an even number */
			test_is_composite(nruns, "0x15A9E672864B1E", 1);
			/* a phuctor-found non-prime public exponent: 170141183460469231731687303715884105731 */
			test_is_composite(nruns, "0x80000000000000000000000000000003", 1);
			break;
		case 3:
			time_mr(nruns);
			break;
		case 4:
			test_rpng(nruns);
			break;
		case 5:
			time_rpng(nruns);
			break;
		default:
			printf("Current test ids:\n");
			printf("0 for timing entropy source\n");
			printf("1 for entropy output test\n");
			printf("2 for is_composite (Miller-Rabin) test\n");
			printf("3 for timing Miller-Rabin\n");
			printf("4 for random prime number generator test\n");
			printf("5 for timing random prime number generator\n");
	}

  return 0;
}