#include <sys/time.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <assert.h>
#include <string.h>
#include <limits.h>

#include <sched.h>
#include <sys/time.h>
#include <sys/resource.h>

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#include "litmus.h"
#include "common.h"

#define UINT_WANT 4294967295
#if UINT_MAX < UINT_WANT
#error Unsigned int not big enough.
#endif

/*
 * Limit the execution times.
 */
#define EXEC_MIN (2 * __NS_PER_MS)
#define EXEC_MAX (100 * __NS_PER_MS)

/*
 * Limit interarrival times.
 */
#define IA_MIN (0 * __NS_PER_MS)
#define IA_MAX (200 * __NS_PER_MS)


static void usage(char *error) {
	fprintf(stderr, "Error: %s\n", error);
	fprintf(stderr,
		"Usage:\n"
		"	bespin [-w] [-o] [-n] [-p partition] "
		"[-s seed] [-f res-time file] OS_TYPE DURATION\n"
		"OS_TYPE is litmus or linux\n"
		"DURATION is milliseconds.\n");
	exit(EXIT_FAILURE);
}

#define NUMS 4096
static int num[NUMS];
static char* progname;

static int task_colors, avg_ways;

static gsl_rng *exec_rng;
static gsl_rng *ia_rng;
static struct timespec sleep_ts = { .tv_sec = 0 };

static void setup_rng(unsigned long seed)
{
	exec_rng = gsl_rng_alloc(gsl_rng_taus);
	ia_rng = gsl_rng_alloc(gsl_rng_taus);
	if (!exec_rng || !ia_rng)
		bail_out("Could not initialize RNG");
	gsl_rng_set(exec_rng, seed);
	gsl_rng_set(ia_rng, seed + 1);
}

static lt_t round_double_lt(const double d)
{
	/* if needed, add one so round up */
	if (d - ((lt_t)d) > 0.0)
		return ((lt_t)(d + 1.0));
	else
		return ((lt_t)d);
}

static lt_t get_exec_time(const lt_t exec_mean)
{
	double e = gsl_ran_exponential(exec_rng, exec_mean);
	e = (e < EXEC_MIN) ? EXEC_MIN : e;
	e = (e > EXEC_MAX) ? EXEC_MAX : e;
	return (round_double_lt(e));
}

static lt_t get_ia_time(const lt_t ia_mean)
{
	double ia = gsl_ran_exponential(ia_rng, ia_mean);
	ia = (ia < IA_MIN) ? IA_MIN : ia;
	ia = (ia > IA_MAX) ? IA_MAX : ia;
	return (round_double_lt(ia));
}

#define LITMUS_STR	"litmus"
#define LINUX_STR	"linux"
static int using_os(const char *const os, const char *const test)
{
	/* LITMUS_STR > LINUX_STR */
	return (0 == strncasecmp(os, test, sizeof(LITMUS_STR)));
}

static int valid_os(char *os)
{
	return (using_os(LITMUS_STR, os) || using_os(LINUX_STR, os));
}

static inline lt_t ts_to_ns(const struct timespec *ts)
{
	return ( (((lt_t)ts->tv_sec) * 1e9) + ts->tv_nsec );
}

static void mono_time_ts(struct timespec *ts)
{
	int ret = clock_gettime(CLOCK_MONOTONIC, ts);
	if (ret)
		bail_out("could not get monotonic time");
}

static lt_t mono_time(void)
{
	struct timespec ts;
	mono_time_ts(&ts);
	return (ts_to_ns(&ts));
}

static void add_to_ts(struct timespec *ts, const lt_t t)
{
	/* convert the lt_t to seconds and nanoseconds */
	const lt_t sec = t / 1e9;
	long nsec = t - sec;

	/* figure out carry for nsec field */
	const long nsec_sum_large = ts->tv_nsec + nsec;
	const time_t sec_carry    = nsec_sum_large / 1e9;

	ts->tv_sec	+= sec + sec_carry;
	ts->tv_nsec	= nsec_sum_large - sec_carry * 1e9;
	assert(ts->tv_nsec < 1e9);
}

static lt_t etime(void)
{
	struct timespec ts;
	int err;
	err = clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts);
	if (err)
		bail_out("could not get thread CPU time");
	return (ts_to_ns(&ts));
}

static int loop_once(void)
{
	int i, j = 0;
	for (i = 0; i < NUMS; i++)
		j += num[i]++;
	return j;
}

static int loop_for(lt_t exec_time, lt_t emergency_exit)
{
	double last_loop = 0, loop_start;
	int job_finished = 1;
	int tmp = 0;

	lt_t start = etime();
	lt_t now = etime();

	while (now + last_loop < start + exec_time) {
		loop_start = now;
		tmp += loop_once();
		now = etime();
		last_loop = now - loop_start;
		if (emergency_exit && mono_time() > emergency_exit)
			job_finished = 0;
	}

	return job_finished;
}

static void do_sleep(const lt_t ia_mean)
{
	const lt_t ia = get_ia_time(ia_mean);
	int err;

	add_to_ts(&sleep_ts, ia);
	err = clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME,
			&sleep_ts, NULL);

	/*
	 * For some reason, nanosleep returns non-zero even when things are
	 * okay the last time. */
	if (err)
		bail_out("could not nanosleep!");
}

static int job(const lt_t exec_mean, const lt_t ia_mean, FILE *out_f,
		const lt_t program_end)
{
	const lt_t now	= mono_time();
	const lt_t exec	= get_exec_time(exec_mean);
	int job_finished;
	lt_t end_time;

	if (now > program_end || now + exec > program_end)
	{
		/* this job is either after we want to stop, or would be,
		 * so just quit now */
		return 0;
	} else {
		/* add NS per MS as an error margin */
		job_finished = loop_for(exec, program_end + __NS_PER_MS);
		end_time = mono_time();

		/* record response time if wanted and the job finished */
		if (out_f && job_finished)
		{
			/* print start time, end time, exec time */
			fprintf(out_f, "%llu, %llu, %llu\n",
					ts_to_ns(&sleep_ts), end_time, exec);
		}

		if (end_time < program_end) {
			/* we can go to sleep some time */
			do_sleep(ia_mean);
			return 1;
		} else {
			/* that was the last job */
			return 0;
		}
	}
}

static void setup_litmus_task(const double phase_ms, const int cpu)
{
	const lt_t phase = phase_ms * __NS_PER_MS;
	int ret;

	/* Best-Effort task has no WCET or period, just use one */
	ret = sporadic_task_ns(1, 1, phase, cpu, RT_CLASS_BEST_EFFORT,
			NO_ENFORCEMENT, 0);
	if (ret < 0)
		bail_out("could not setup rt task params");

	request_resources(task_colors, avg_ways);

	init_litmus();

	ret = task_mode(LITMUS_RT_TASK);
	if (ret != 0)
		bail_out("could not become RT task");
}

#define OPTSTR "p:wl:r:nos:f:h:"

int main(int argc, char** argv)
{
	unsigned long seed = 1;
	double duration_ms;
	double exec_mean_ms = 10.0, ia_mean_ms = 100.0, phase_ms = 0.0;
	lt_t duration, start, exec_mean, ia_mean;
	char *os_type, *out_fname = NULL;
	FILE *out_f = NULL;
	int opt;
	int ret;
	int wait = 0;
	int nice = 0;
	int fifo = 0;
	int cpu = 0;
	int migrate = 0;
	struct sched_param param;

	progname = argv[0];

	while ((opt = getopt(argc, argv, OPTSTR)) != -1) {
		switch (opt) {
		case 'w':
			wait = 1;
			break;
		case 'p':
			cpu = atoi(optarg);
			migrate = 1;
			break;
		case 'n':
			nice = 1;
			break;
	        case 'o':
			fifo = 1;
			break;
		case 's':
			seed = atoi(optarg);
			break;
		/* case 'h': */
		/* 	phase_ms = atof(optarg); */
		/* 	break; */
		case 'e':
			exec_mean_ms = atof(optarg);
			break;
		case 'i':
			ia_mean_ms = atof(optarg);
			break;
		case 'f':
			out_fname = optarg;
			break;
	        case 'h':
			sscanf(optarg, "%d,%d", &task_colors, &avg_ways);
			break;
		case ':':
			usage("Argument missing.");
			break;
		case '?':
		default:
			usage("Bad argument.");
			break;
		}
	}

	if (argc - optind != 2)
		usage("Arguments missing.");

	os_type		= argv[optind + 0];
	duration_ms	= atof(argv[optind + 1]);

	if (!valid_os(os_type))
		bail_out("Invalid OS type");

	if (duration_ms < 0)
		bail_out("duration < 0 not allowed");

	if (exec_mean_ms < 0)
		bail_out("execution time mean < 0 not allowed");

	if (ia_mean_ms < 0)
		bail_out("IA time mean < 0 not allowed");

	if (exec_mean_ms * __NS_PER_MS > UINT_MAX >> 1)
		bail_out("Exec time too big");

	if (seed < 1)
		bail_out("seed < 1 not allowed");

	if (phase_ms < 0)
		bail_out("phase < 0 not allowed");

	if (wait && using_os(LINUX_STR, os_type))
		bail_out("Linux cannot use the wait flag");

	duration	= duration_ms	* __NS_PER_MS;
	exec_mean	= exec_mean_ms	* __NS_PER_MS;
	ia_mean		= ia_mean_ms	* __NS_PER_MS;

	if (out_fname)
	{
		out_f = fopen(out_fname, "w");
		if (!out_f)
			bail_out("Could not open response-time file");
	}

	setup_rng(seed);

	if (using_os(LITMUS_STR, os_type)) {

		if (migrate) {
			ret = be_migrate_to(cpu);
			if (ret < 0)
				bail_out("could not migrate to target partition");
		}

		setup_litmus_task(phase_ms, cpu);

		if (wait) {
			ret = wait_for_ts_release();
			if (ret != 0)
				bail_out("wait_for_ts_release()");
		}
	} else if (fifo) {
		param.sched_priority = sched_get_priority_max(SCHED_FIFO) - 2;
		ret = sched_setscheduler(0, SCHED_FIFO, &param);
		if (ret)
			bail_out("Could not switch to FIFO scheduler");
	} else if (nice) {
		ret = setpriority(PRIO_PROCESS, 0, 5);
		if (ret)
			bail_out("Could not set task priority");
	}

	/* initialize the start time */
	mono_time_ts(&sleep_ts);
	start = ts_to_ns(&sleep_ts);

	while (job(exec_mean, ia_mean, out_f, start + duration));

	if (using_os(LITMUS_STR, os_type)) {
		ret = task_mode(BACKGROUND_TASK);
		if (ret != 0)
			bail_out("could not become regular task (huh?)");
	}

	return 0;
}