#define _GNU_SOURCE

#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <pthread.h>
#include <sys/eventfd.h>
#include <sys/prctl.h>
#include <signal.h>
#include <err.h>

int to_thread, from_thread;
volatile int state;  // 0 = warmup.  1 = benchmark.  2 = exit.

int use_xstate = 0;
int thread_cpu;

static void maybe_use_xstate(void)
{
	if (use_xstate)
		asm volatile ("pxor %%xmm0, %%xmm0" : : : "xmm0");
}

void *threadproc(void *x)
{
	cpu_set_t cpuset;
	CPU_ZERO(&cpuset);
	CPU_SET(thread_cpu, &cpuset);
	if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
		printf("[WARN]\tsched_setaffinity failed\n");

	while(1)
	{
		uint64_t buf;

		if (read(to_thread, &buf, 8) != 8)
			err(1, "read eventfd");

		if (state == 2)
			return 0;

		maybe_use_xstate();

		buf = 1;
		if (write(from_thread, &buf, 8) != 8)
			err(1, "write eventfd");
	}
}

static void bounce()
{
	uint64_t buf = 1;
	if (write(to_thread, &buf, 8) != 8)
		err(1, "write eventfd");
  
	if (read(from_thread, &buf, 8) != 8)
		err(1, "read eventfd");
}

void killit()
{
	uint64_t buf = 1;
	state = 2;
	if (write(to_thread, &buf, 8) != 8)
		err(1, "write eventfd at exit");
}

int main(int argc, char **argv)
{
	struct timespec start, end;
	pthread_t thread;
	uint64_t i, iters = 100000;
	uint64_t ns;
	int main_cpu;
	bool use_thread = false;

	if (argc != 4) {
		printf("Usage: %s <0|1> <thread|process> <same|different>\n\nSet the parameter to 1 to use xstate\nUse 'same' for same CPU or 'different' for cross-CPU\n",
		       argv[0]);
		return 1;
	}

	if (!strcmp(argv[1], "1"))
		use_xstate = 1;
	else if (strcmp(argv[1], "0"))
		errx(1, "first arg is wrong");

	if (!strcmp(argv[3], "same")) {
		main_cpu = 0;
		thread_cpu = 0;
	} else if (!strcmp(argv[3], "different")) {
		main_cpu = 0;
		thread_cpu = 1;
	} else {
		errx(1, "third arg is wrong");
	}

	if (!strcmp(argv[2], "thread")) {
		use_thread = true;
	} else if (!strcmp(argv[2], "process")) {
		use_thread = false;
	} else {
		errx(1, "second arg is wrong");
	}

	printf("use_xstate = %d\n", use_xstate);

	to_thread = eventfd(0, 0);
	if (to_thread < 0)
		err(1, "eventfd");
	from_thread = eventfd(0, 0);
	if (from_thread < 0)
		err(1, "eventfd");

	if (use_thread) {
		printf("Using threads\n");
		if (pthread_create(&thread, 0, threadproc, 0) != 0)
			err(1, "pthread_create");
	} else {
		printf("Using a subprocess\n");
		pid_t child = fork();
		if (child < 0)
			err(1, "fork");

		if (child == 0) {
			prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0);
			threadproc(NULL);
		}
	}

	cpu_set_t cpuset;
	CPU_ZERO(&cpuset);
	CPU_SET(main_cpu, &cpuset);
	if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
		printf("[WARN]\tsched_setaffinity failed\n");

	/* Warm up (and burn the xstate heuristic) */
	for(i = 0; i < 10000; i++) {
		bounce();
		maybe_use_xstate();
	}

	/* Clear out the xstate hack. */
	state = 1;
	bounce();

	clock_gettime(CLOCK_MONOTONIC, &start);

	for (i = 0; i < iters; i++) {
		bounce();
		maybe_use_xstate();
	}

	clock_gettime(CLOCK_MONOTONIC, &end);

	killit();

	ns = 1000000000ULL * (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec);

	printf("%llu iters at %.1f ns/switch\n",
	       (unsigned long long)iters, (double)ns / (2*iters));

	if (use_thread)
		pthread_join(thread, 0);

	return 0;
}