#include <sys/time.h>

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

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

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

#include <litmus/sched_mc.h>


static void usage(char *error) {
	fprintf(stderr, "Error: %s\n", error);
	fprintf(stderr,
		"Usage:\n"
		"	rt_spin [COMMON-OPTS] LVL_C_TIME WCET PERIOD DURATION\n"
		"	rt_spin -l\n"
		"\n"
		"COMMON-OPTS = [-w] [-p PARTITION] [-c CLASS] [-s SCALE]\n"
		"              [-r CRITICALITY = [a|b|c|d]] [-i MC-LVL-A-ID]\n"
		"              [-a ALPHA] [-b BETA] [-d SEED]\n"
		"\n"
		"WCET and PERIOD are nanoseconds, DURATION is seconds.\n");
	exit(EXIT_FAILURE);
}

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

static gsl_rng *beta_rng;
static gsl_rng *coin_rng;

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

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

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

	double start = cputime();
	double now = cputime();

	while (now + last_loop < start + exec_time) {
		loop_start = now;
		tmp += loop_once();
		now = cputime();
		last_loop = now - loop_start;
		if (emergency_exit && wctime() > emergency_exit) {
			/* Oops --- this should only be possible if the execution time tracking
			 * is broken in the LITMUS^RT kernel. */
			fprintf(stderr, "!!! rtspin/%d emergency exit!\n", getpid());
			fprintf(stderr, "Something is seriously wrong! Do not ignore this.\n");
			break;
		}
	}

	return tmp;
}


static void debug_delay_loop(void)
{
	double start, end, delay;

	while (1) {
		for (delay = 0.5; delay > 0.01; delay -= 0.01) {
			start = wctime();
			loop_for(delay, 0);
			end = wctime();
			printf("%6.4fs: looped for %10.8fs, delta=%11.8fs, error=%7.4f%%\n",
			       delay,
			       end - start,
			       end - start - delay,
			       100 * (end - start - delay) / delay);
		}
	}
}

static int job(double exec_time, double program_end)
{
	if (wctime() > program_end)
		return 0;
	else {
		loop_for(exec_time, program_end + 1);
		sleep_next_period();
		return 1;
	}
}

enum crit_level str2crit(const char* str)
{
	if (0 == strncasecmp("a", str, 1))
		return CRIT_LEVEL_A;
	else if (0 == strncasecmp("b", str, 1))
		return CRIT_LEVEL_B;
	else if (0 == strncasecmp("c", str, 1))
		return CRIT_LEVEL_C;
	/* failure */
	return NUM_CRIT_LEVELS;
}

#define OPTSTR "a:b:d:p:c:wlves:r:i:"

#define COIN_PROB 0.5

int main(int argc, char** argv)
{
	int ret;
	lt_t wcet;
	lt_t period;
	lt_t lvl_c_time;
	int migrate = 0;
	int cpu = 0;
	int opt;
	int wait = 0;
	int test_loop = 0;
	int want_enforcement = 0;
	unsigned long seed = 1;
	double duration = 0, start;
	double scale = 1.0;
	double alpha = -1, beta = -1, beta_sample, exec_time;
	task_class_t class = RT_CLASS_HARD;
	struct mc_task mc_task = { .crit = NUM_CRIT_LEVELS, .lvl_a_id = -1 };

	progname = argv[0];

	while ((opt = getopt(argc, argv, OPTSTR)) != -1) {
		switch (opt) {
		case 'a':
			alpha = atof(optarg);
			break;
		case 'b':
			beta = atof(optarg);
			break;
		case 'd':
			seed = atoi(optarg);
			break;
		case 'w':
			wait = 1;
			break;
		case 'p':
			cpu = atoi(optarg);
			migrate = 1;
			break;
		case 'c':
			class = str2class(optarg);
			if (class == -1)
				usage("Unknown task class.");
			break;
		case 'e':
			want_enforcement = 1;
			break;
		case 'l':
			test_loop = 1;
			break;
		case 's':
			scale = atof(optarg);
			break;
		case 'r':
			mc_task.crit = str2crit(optarg);
			if (NUM_CRIT_LEVELS == mc_task.crit)
				usage("Bad crit level.");
			break;
		case 'i':
			mc_task.lvl_a_id = atoi(optarg);
			break;
		case ':':
			usage("Argument missing.");
			break;
		case '?':
		default:
			usage("Bad argument.");
			break;
		}
	}

	if (test_loop) {
		debug_delay_loop();
		return 0;
	}

	if (argc - optind < 4)
		usage("Arguments missing.");

	lvl_c_time	= atoi(argv[optind + 0]);
	wcet		= atoi(argv[optind + 1]);
	period 		= atoi(argv[optind + 2]);

	if (seed < 1)
		usage("Seed < 1 not allowed.");
	if (alpha <= 0 || beta <= 0)
		usage("Need positive alpha and beta.");
	if (wcet <= 0)
		usage("The worst-case execution time must be a "
				"positive number.");
	if (period <= 0)
		usage("The period must be a positive number.");
	if (wcet > period) {
		usage("The worst-case execution time must not "
				"exceed the period.");
	}

	duration  = atof(argv[optind + 3]);

	setup_rng(seed);

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

	ret = sporadic_task_ns(wcet, period, 0, cpu, class,
			       want_enforcement ? PRECISE_ENFORCEMENT
			                        : NO_ENFORCEMENT,
			       migrate);
	if (ret < 0)
		bail_out("could not setup rt task params");

	if (NUM_CRIT_LEVELS != mc_task.crit) {
		ret = set_rt_task_mc_param(gettid(), &mc_task);
		if (ret < 0)
			bail_out("could not setup rt mixed criticality params");
	}

	init_litmus();

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

	if (wait) {
		ret = wait_for_ts_release();
		if (ret != 0)
			bail_out("wait_for_ts_release()");
	}

	start = wctime();

	do {
		double coin_flip = gsl_rng_uniform(coin_rng);
		if (coin_flip < COIN_PROB) {
			beta_sample = gsl_ran_beta(beta_rng, alpha, beta);
			exec_time = lvl_c_time * 20.0 * beta_sample * scale;
		} else
			exec_time = lvl_c_time;
		/* convert to seconds */
		exec_time = exec_time * 0.000000001;
	} while(job(exec_time, start + duration));

	ret = task_mode(BACKGROUND_TASK);
	if (ret != 0)
		bail_out("could not become regular task (huh?)");

	gsl_rng_free(beta_rng);
	gsl_rng_free(coin_rng);

	return 0;
}