#include "gen.hh"
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <fstream>
#include "ntr.h"
#include "nms.h"
#include <getopt.h>
#include <unistd.h>
#include <topo.h>

static void
usage()
{
	fprintf(stdout,
	    "Usage:\n"
		"    -v: verbose mode\n"
	    "    -b: memory block size\n"
	    "    -d: destination domain index\n"
	    "    -s: worker threads cpu list\n"
		"    -S: enable shared memory block\n"
		"    -t: time to run\n"
		"    -o: output file path\n");
	fflush(stdout);
}

static char output_file[256] = "memloadgen_samples.txt";

int main(int argc, char * argv[])
{
	ntr_init();
	ntr_set_level(NTR_DEP_USER1, NTR_LEVEL_WARNING);

	size_t arr_sz = 1024 * 1024;
	uint32_t time = -1;
	cpuset_t threads;
	CPU_ZERO(&threads);
	CPU_SET(0, &threads);

	int shared_buffer = 0;
	cpuset_t domain_mask;
	CPU_ZERO(&domain_mask);
	CPU_SET(0, &domain_mask);
	{
		int c;
		// parse arguments
		while ((c = getopt(argc, argv, "vhb:d:s:So:t:")) != -1) {
			switch (c) {
			case 'v':
				ntr_set_level(NTR_DEP_USER1, ntr_get_level(NTR_DEP_USER1) + 1);
				break;
			case 'h':
				usage();
				exit(0);
			case 'b':
				arr_sz = strtoull(optarg, nullptr, 10);
				break;
			case 'd':
				cpulist_to_cpuset(optarg, &domain_mask);
				break;
			case 's':
				cpulist_to_cpuset(optarg, &threads);
				break;
			case 'S':
				shared_buffer = 1;
				break;
			case 'o':
				strncpy(output_file, optarg, 256);
				break;
			case 't':
				time = (uint32_t)strtol(optarg, nullptr, 10);
			default:
				usage();
				exit(0);
			}
		}
	}

	ntr(NTR_DEP_USER1, NTR_LEVEL_INFO, "MLG: [size: %ld, # threads: 0x%d, domain: 0x%ld]\n", arr_sz, CPU_COUNT(&threads), CPU_FFS(&domain_mask) - 1);

	// init topo
	if (topo_init(ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT)) {
		fprintf(stderr, "libtopo init failed!\n");
		exit(1);
	}

	// init 
	if (nms_init(ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT)) {
		fprintf(stderr, "libnms init failed!\n");
		exit(1);
	}

	bool success = false;
	memload_generator::memload_generator_options opts;
	opts.chunk_size = arr_sz;
	opts.shared_buffer = shared_buffer;
	opts.verbose = ntr_get_level(NTR_DEP_USER1) != NTR_LEVEL_DEFAULT;

	std::ofstream ofile;
	ofile.open(output_file, std::ios::out | std::ios::trunc);

    auto mgen = new memload_generator(&threads, &domain_mask, &opts, &success);
	if (!mgen->start()) {
		fprintf(stderr, "failed to start memloadgen!\n");
		exit(1);
	}

	uint64_t prev_ts = topo_uptime_ns();
	uint64_t prev_trans = mgen->get_transactions();
	uint32_t cur_time = 0;
	while(cur_time < time) {
		usleep(S2US);
		uint64_t cur_ts = topo_uptime_ns();
		uint64_t trans = mgen->get_transactions();
		uint64_t bps = (uint64_t)((double)((trans - prev_trans) * arr_sz) / ((double)(cur_ts - prev_ts) / (double)S2NS));
		
		ntr(NTR_DEP_USER1, NTR_LEVEL_INFO, "main: MLG bps = %ld ~= %ldM\n", bps, bps / 1024 / 1024);
		ofile << bps << std::endl;

		prev_ts = cur_ts;
		prev_trans = trans;
		cur_time++;
	}
	mgen->stop();
	delete mgen;
	ofile.close();

	return 0;
}