#include <obliv.oh>
#include "oram.oh"
#include <copy.oh>
#include "test_generic.h"

static const char TESTNAME[] = "oram_readwrite_benchmark";

#define TEXT_HELP_SUPPLEMENTARY "\
  -e \x1b[4mNUMBER\x1b[0m, --element-count=\x1b[4mNUMBER\x1b[0m \n\t\tuse ORAMs of \x1b[4mNUMBER\x1b[0m elements\n\n\
  -s \x1b[4mNUMBER\x1b[0m, --element-size=\x1b[4mNUMBER\x1b[0m \n\t\tuse ORAMs with elements containing \x1b[4mNUMBER\x1b[0m 32-bit words\n\n\
  -o \x1b[4mTYPE\x1b[0m, --oram-type=\x1b[4mTYPE\x1b[0m \n\t\tforce all ORAMs to be \x1b[4mTYPE\x1b[0m ORAMs. Valid types are \033[1msqrt\033[0m, \033[1mcircuit\033[0m, \033[1mfssl\033[0m, \033[1mfssl_cprg\033[0m, and \033[1mlinear\033[0m.\n\n\
  -i \x1b[4mNUMBER\x1b[0m, --samples=\x1b[4mNUMBER\x1b[0m \n\t\trun \x1b[4mNUMBER\x1b[0m iterations of the benchmark\n\n"

static const char options_string[] = "e:s:o:i:";
static struct option long_options[] = {
	{"element-count", required_argument, NULL, 'e'},
	{"element-size", required_argument, NULL, 's'},
	{"oram-type", required_argument, NULL, 'o'},
	{"samples", required_argument, NULL, 'i'},
	{0, 0, 0, 0}
};

char* get_test_name() {
	return TESTNAME;
}

char* get_supplementary_options_string() {
	return options_string;
}

struct option* get_long_options() {
	return long_options;
}

void print_supplementary_help() {
	fprintf(stderr, TEXT_HELP_SUPPLEMENTARY);
}

void test_main(void*varg) {

	#ifdef ORAM_OVERRIDE
	oram_set_default_type(ORAM_OVERRIDE);
	#endif

	size_t elct = 4;
	size_t elsz = 1;
	int samples = 1;

	args_t * args_pass = varg;
	int arg;
	optind = 0; // this allows us to getopt a second time
	while ((arg = getopt_long(args_pass->argc, args_pass->argv, options_string, long_options, NULL)) != -1) {
		if (arg == 'e') {
			elct = atoll(optarg);
			if (elct <= 0) {
				fprintf (stderr, "Argument for -%c must be positive.\n", arg);
				return;
			}
		} else if (arg == 's') {
			elsz = atoll(optarg);
			if (elsz <= 0) {
				fprintf (stderr, "Argument for -%c must be positive.\n", arg);
				return;
			}
		} else if (arg == 'o') {
			if (strcmp(optarg,"sqrt") == 0) {
				oram_set_default_type(ORAM_TYPE_SQRT);
			} else if (strcmp(optarg,"circuit") == 0) {
				oram_set_default_type(ORAM_TYPE_CIRCUIT);
			} else if (strcmp(optarg,"linear") == 0) {
				oram_set_default_type(ORAM_TYPE_LINEAR);
			} else if (strcmp(optarg,"fssl") == 0) {
				oram_set_default_type(ORAM_TYPE_FSSL);
			} else if (strcmp(optarg,"fssl_cprg") == 0) {
				oram_set_default_type(ORAM_TYPE_FSSL_CPRG);
			} else {
				fprintf (stderr, "Invalid argument for -%c.\n", arg);
				return;
			}
		} else if (arg == 'i') {
			samples = atoi(optarg);
			if (samples < 0) {
				fprintf (stderr, "Argument for -%c must be nonnegative.\n", arg);
				return;
			}
		} else if (arg == '?' || arg == ':') {
			if (optopt == 'e' || optopt == 's' || optopt == 'o' || optopt == 'i') {
				fprintf (stderr, "Option -%c requires an argument.\n", optopt);
				return;
			} else {
				fprintf (stderr, "Option -%c not recognized.\n", optopt);
				return;
			}
		} else {
			abort();
		}
	}

	OcCopy cpy = ocCopyIntN(elsz);

	uint64_t tally = 0;
	uint64_t tallygates = 0;
	uint64_t tallybytes = 0;
	obliv uint32_t * input = calloc(1, elsz * sizeof(obliv uint32_t));

	int64_t startuptime = -current_timestamp();
	int64_t startupgates = -yaoGateCount();
	int64_t startupbytes = -tcp2PBytesSent(ocCurrentProto());
	oram * o = oram_new(ORAM_TYPE_AUTO, &cpy, elct);
	startuptime += current_timestamp();
	startupgates += yaoGateCount();
	startupbytes += tcp2PBytesSent(ocCurrentProto());

	#ifndef ORAM_PROFILE_SCHEDULING
	fprintf(stdout, "# ORAM ACCESS (SETUP) (element count, element size, startup microseconds, startup gates, startup bytes)\n");
	fprintf(stdout, "%lld,%lld", elct, elsz);
	fprintf(stdout, ",%llu,%llu,%llu\n", startuptime,startupgates,startupbytes);
	fprintf(stdout, "# ORAM ACCESS (READ/WRITE) (element count, element size, sample 1 microseconds, sample 1 gates, sample 1 bytes, ...)\n");
	fprintf(stdout, "%lld,%lld", elct, elsz);
	fflush(stdout);
	#endif

	for (int ii = 0; ii < samples; ii++) {
		uint32_t index_raw = ocBroadcastInt(rand() % elct, 2);
		obliv uint32_t index = feedOblivInt(index_raw, 2);
		uint32_t rindex_raw = ocBroadcastInt(rand() % elct, 2);
		obliv uint32_t rindex = feedOblivInt(rindex_raw, 2);
		for (int kk = 0; kk < elsz; kk++) input[kk] = feedOblivInt(rand(), 1);
		int64_t runtime = -current_timestamp();
		int64_t rungates = -yaoGateCount();
		int64_t runbytes = -tcp2PBytesSent(ocCurrentProto());
		oram_write(o, input, index);
		oram_read(input, o, rindex);
		runtime += current_timestamp();
		rungates += yaoGateCount();
		runbytes += tcp2PBytesSent(ocCurrentProto());
		#ifndef ORAM_PROFILE_SCHEDULING
		fprintf(stdout, ",%llu,%llu,%llu", runtime,rungates, runbytes);
		fflush(stdout);
		#endif
		tally += runtime;
		tallygates += rungates;
		tallybytes += runbytes;
	}

	free(input);
	oram_free(o);
	#ifndef ORAM_PROFILE_SCHEDULING
	fprintf(stdout, "\n");
	#endif
	if (samples > 0) {
	fprintf(stderr, "ORAM Write (count:%lld, size: %lld): %llu microseconds avg, %llu gates avg, %llu bytes avg\n", elct, elsz, tally / samples, tallygates/samples, tallybytes/samples);
	}

}