#define _GNU_SOURCE
#include <pthread.h>
#include <sys/prctl.h>
#include <stdio.h>
#include <string.h>
#include <sched.h>

#define TRAINING		(50)
#define ITERATIONS		(10^3)
#define HIGH_THD_ITERATIONS	(TRAINING + ITERATIONS)

#define NS_PER_SEC	(10^9)

#define HIGH_PRIO	(59)

#define CORE0_ID	(0)
#define CORE1_ID	(1)

pthread_mutex_t lock;

/* the busyloop & lowpriority threads will loop on this while the high priority
   thread will set it after it finishes all its iterations */
int threads_stop_flag = 0;

/* the high-priority thread uses this struct to find the maximum latency between
   payloads from all its iterations */
struct timespec max_latency = {0, 0};

struct loads
{
	struct timespec hold_lock_load;
	struct timespec no_lock_load;
};

int subtract_timespec(struct timespec *x,
		      struct timespec *y,
		      struct timespec *result)
{
	if (x->tv_nsec < y->tv_nsec) {
		int sec = (y->tv_nsec - x->tv_nsec) / NS_PER_SEC + 1;
		y->tv_nsec -= NS_PER_SEC * sec;
		y->tv_sec += sec;
	}
	if (x->tv_nsec - y->tv_nsec > NS_PER_SEC) {
		int sec = (x->tv_nsec - y->tv_nsec) / NS_PER_SEC;
		y->tv_nsec += NS_PER_SEC * sec;
		y->tv_sec -= sec;
	}

	result->tv_sec = x->tv_sec - y->tv_sec;
	result->tv_nsec = x->tv_nsec - y->tv_nsec;

	return x->tv_sec < y->tv_sec;
}

int compare_timespec(struct timespec *x,
		     struct timespec *y)
{
	if (x->tv_sec > y->tv_sec)
		return 1;
	if (x->tv_sec == y->tv_sec)
	{
		if (x->tv_nsec == y->tv_nsec)
			return 0;
		if (x->tv_nsec > y->tv_nsec)
			return 1;
	}
	return -1;
}

int setup_rt_pthread(int prio, int cpu, const char* name)
{
	struct sched_param schedp;
	cpu_set_t cpuset;
	pthread_t thread = pthread_self();
	int error = 0;

	memset(&schedp, 0, sizeof(schedp));
	CPU_ZERO(&cpuset);

	schedp.sched_priority = prio;
	error = pthread_setschedparam(thread, SCHED_FIFO, &schedp);
	if (error)
		fprintf(stderr, "%s: Couldn't setschedparam, returned %d\n", name, error);

	CPU_SET(cpu, &cpuset);
	error = pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset);
	if (error)
		fprintf(stderr, "%s: Couldn't setaffinity, returned %d\n", name, error);

	error = prctl(PR_SET_NAME, name, NULL, NULL, NULL,NULL);
	if (error)
		fprintf(stderr, "%s: Couldn't set proc name, returned %d\n", name, error);

	return -error;
}

inline void do_load(struct timespec load)
{
	struct timespec loop_start, current, result;

	clock_gettime(CLOCK_MONOTONIC, &loop_start);

	do {
		clock_gettime(CLOCK_MONOTONIC, &current);
		subtract_timespec(&current, &loop_start, &result);
	} while (compare_timespec(&result, &load) <= 0);
}

void *busy_loop_threadfun(void *busy_loads)
{
	if (setup_rt_pthread(HIGH_PRIO - 1, CORE1_ID, "Busy-loop******")) {
		fprintf(stderr, "Busy-loop****** setup error");
		return (void*) -1;
	}
	
	while (!threads_stop_flag); /* loop until high-prio thread signals */

	return 0;
}

void *low_prio_threadfun(void *low_loads)
{
	struct loads *load = (struct loads*) low_loads;

	if (prctl(PR_SET_NAME, "Low-prio******", NULL, NULL, NULL,NULL) < 0) {
		fprintf(stderr, "Low--prio****** setup error");
		return (void*) -1;
	}

	while (!threads_stop_flag) {
		pthread_mutex_lock(&lock);
		do_load(load->hold_lock_load);
		pthread_mutex_unlock(&lock);
		do_load(load->no_lock_load);
	}

	return 0;
}

void *high_prio_threadfun(void* high_loads)
{
	struct loads *load = (struct loads*)high_loads;
	struct timespec prev, cur, latency;
	int i;

	if (setup_rt_pthread(HIGH_PRIO, CORE1_ID, "High-prio******")) 
	{
		fprintf(stderr, "High-prio****** setup error");
		threads_stop_flag = 1;
		return (void*) -1;
	}

	clock_gettime(CLOCK_MONOTONIC, &prev);

	/* runs the payloads and computes latencies between them, stores max;
	   the first iterations are for training, so no latency is stored */
	for (i = HIGH_THD_ITERATIONS; i > 0; --i) {
		pthread_mutex_lock(&lock);
		do_load(load->hold_lock_load);
		pthread_mutex_unlock(&lock);
		do_load(load->no_lock_load);

		clock_gettime(CLOCK_MONOTONIC, &cur);
		subtract_timespec(&cur, &prev, &latency);
		subtract_timespec(&latency, &load->hold_lock_load, &latency);
		subtract_timespec(&latency, &load->no_lock_load, &latency);

		if (i <= ITERATIONS &&
		    compare_timespec(&latency, &max_latency) > 0) {
			max_latency.tv_sec = latency.tv_sec;
			max_latency.tv_nsec = latency.tv_nsec;
		}

		prev = cur;
	}

	threads_stop_flag = 1;

	return 0;
}

unsigned int initialize_mutex()
{
	pthread_mutexattr_t lock_attr;
	unsigned int error = 0;

	memset(&lock, 0, sizeof(pthread_mutex_t));
	pthread_mutexattr_init(&lock_attr);
	error = pthread_mutexattr_settype(&lock_attr, PTHREAD_MUTEX_NORMAL);
	error |= pthread_mutexattr_setprotocol(&lock_attr, PTHREAD_PRIO_INHERIT);
	error |= pthread_mutex_init(&lock, &lock_attr);
	pthread_mutexattr_destroy(&lock_attr);

	return error;
}

int main (int argc, char* argv[])
{
	pthread_t low_prio_thd, high_prio_thd, busy_loop_thd;
	struct loads high_prio_load = {{0,1001},{0,1001}};
	struct loads low_prio_load = {{0,1001},{0,1001}};

	if (initialize_mutex())
		return -1;

	pthread_create(&low_prio_thd, NULL, &low_prio_threadfun, (void *) &low_prio_load);
	pthread_create(&high_prio_thd, NULL, &high_prio_threadfun, (void *) &high_prio_load);
	pthread_create(&busy_loop_thd, NULL, &busy_loop_threadfun, NULL);

	pthread_join(low_prio_thd, NULL);
	pthread_join(high_prio_thd, NULL);
	pthread_join(busy_loop_thd, NULL);

	pthread_mutex_destroy(&lock);

	printf("Highest latency: %li\n", max_latency.tv_sec * NS_PER_SEC + max_latency.tv_nsec);

	return 0;
}