#define _GNU_SOURCE

#include <assert.h>
#include <stdio.h>
#include <pthread.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <sched.h>
#include <errno.h>
#include <sys/syscall.h>

#define PCOUNT 20
#define EXPECTED_PRIO -21

// Copied from related header.
struct sched_attr {
	uint32_t size;

	uint32_t sched_policy;
	uint64_t sched_flags;

	/* SCHED_NORMAL, SCHED_BATCH */
	int32_t sched_nice;

	/* SCHED_FIFO, SCHED_RR */
	uint32_t sched_priority;

	/* SCHED_DEADLINE */
	uint64_t sched_runtime;
	uint64_t sched_deadline;
	uint64_t sched_period;

	/* Utilization hints */
	uint32_t sched_util_min;
	uint32_t sched_util_max;
};

static pthread_mutex_t mutexes[PCOUNT];

static void init_mutex(pthread_mutex_t *mutex) {
	pthread_mutexattr_t attr;
	int ret;

	ret = pthread_mutexattr_init(&attr);
	assert(!ret);
	ret = pthread_mutexattr_setprotocol(&attr, PTHREAD_PRIO_INHERIT);
	assert(!ret);

	ret = pthread_mutex_init(mutex, &attr);
	assert(!ret);

	pthread_mutexattr_destroy(&attr);
}

enum thread_type {
	WAITER,
	PRIO_CHANGER,
};

static struct thread_spec {
	enum thread_type type;
	int name_index;
	pthread_barrier_t *barrier;

	// Used by waiter.
	pthread_mutex_t *wait_mutex;
	pid_t pid_save;
	
	// Used by prio changer.
	pid_t *waiter_pid;
	uint32_t prio;
} thread_specs[PCOUNT * 2];

static pthread_t threads[PCOUNT * 2];
static int thread_count = 0;

static int barrier_wait(pthread_barrier_t *barrier) {
	if (!barrier) {
		return 0;
	}

	int ret = pthread_barrier_wait(barrier);
	assert(!ret || ret == PTHREAD_BARRIER_SERIAL_THREAD);

	return ret;
}

static int sched_getattr(pid_t pid, struct sched_attr *attr, int flags) {
	return syscall(SYS_sched_getattr, pid, attr, sizeof(*attr), flags);
}

static int sched_setattr(pid_t pid, struct sched_attr *attr, int flags) {
	return syscall(SYS_sched_setattr, pid, attr, flags);
}

static void *prio_change_loop(struct thread_spec *spec) {
	struct sched_attr attr;
	int ret;

	for (;;) {
		barrier_wait(spec->barrier);

		ret = sched_getattr(*spec->waiter_pid, &attr, 0);
		if (ret < 0) {
			perror("sched_getattr");
			assert(0);
		}
		
		attr.sched_policy = SCHED_RR;
		attr.sched_priority = spec->prio;
	
		ret = sched_setattr(*spec->waiter_pid, &attr, 0);
		if (ret < 0) {
			perror("sched_setattr");
			exit(1);
		}

		barrier_wait(spec->barrier);
	}
}

static void *thread(void *arg) {
	struct thread_spec *spec = arg;
	char name[64];
	int ret;
	
	snprintf(name, sizeof(name), "%s-%d", 
			spec->type == WAITER ? "waiter" : "changer", 
			spec->name_index);
	ret = pthread_setname_np(pthread_self(), name);
	assert(!ret);

	switch (spec->type) {
	case WAITER:
		spec->pid_save = gettid();

		barrier_wait(spec->barrier);

		ret = pthread_mutex_lock(spec->wait_mutex);
		assert(!ret);
		break;
	case PRIO_CHANGER:
		prio_change_loop(spec);
		break;
	default:
		assert(0);
		break;
	}

	return NULL;
}

static void create_thread(struct thread_spec *spec) {
	int ret;

	ret = pthread_create(&threads[thread_count++], NULL, thread, spec);
	assert(!ret);
}

static int get_prio(int need_print) {
	FILE *file = fopen("/proc/self/stat", "r");
	assert(file);

	char comm[64];
	int tmp, prio, ret;
	ret = fscanf(file, "%d %s %c %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d",
			&tmp, comm, comm, &tmp, &tmp, &tmp, &tmp, &tmp, &tmp,
			&tmp, &tmp, &tmp, &tmp, &tmp, &tmp, &tmp, &tmp, &prio);
	assert(ret == 18);

	if (need_print) {
		printf("prio: %d\n", prio);
	}

	fclose(file);

	return prio;
}

static void print_threads(pid_t process_pid) {
	char cmd[128];
	snprintf(cmd, sizeof(cmd), "ps -L -p %d", process_pid);
	system(cmd);
}

int main(void) {
	pid_t pid;
	int ret;

	pid = getpid();
	printf("pid: %d\n", pid);

	for (int i = 0; i < sizeof(mutexes) / sizeof(mutexes[0]); ++i) {
		init_mutex(&mutexes[i]);
		ret = pthread_mutex_lock(&mutexes[i]);
		assert(!ret);
	}

	pthread_barrier_t barrier;
	ret = pthread_barrier_init(&barrier, NULL, PCOUNT + 1);
	assert(!ret);

	for (int i = 0; i < PCOUNT; ++i) {
		struct thread_spec *spec = &thread_specs[i];
		
		spec->type = WAITER;
		spec->name_index = i;
		spec->wait_mutex = &mutexes[i];
		spec->barrier = &barrier;

		create_thread(spec);
	}

	// Wait for filling of pid_save.
	barrier_wait(&barrier);

	for (int i = 0; i < PCOUNT; ++i) {
		struct thread_spec *spec = &thread_specs[i + PCOUNT];

		spec->type = PRIO_CHANGER; 
		spec->name_index = i;
		spec->prio = 99;
		spec->barrier = &barrier;
		spec->waiter_pid = &thread_specs[i].pid_save;

		create_thread(spec);
	}

	printf("Before: ");
	get_prio(1);
	print_threads(pid);

	time_t ts, start_ts;
	srandom(time(&ts));
	start_ts = ts;
	int iter = 0;
	for (;;) {
		++iter;
		for (int i = 0; i < PCOUNT; ++i) {
			thread_specs[i + PCOUNT].prio = (iter % 2) ? i + 1 : PCOUNT - i;
		}

		for (int i = 0; i < PCOUNT; ++i) {
			int pos = random() % PCOUNT;
			int tmp = thread_specs[pos + PCOUNT].prio;
			thread_specs[pos + PCOUNT].prio = thread_specs[PCOUNT].prio;
			thread_specs[PCOUNT].prio = tmp;
		}

		barrier_wait(&barrier);
		barrier_wait(&barrier);

		//system("echo iteration > /sys/kernel/debug/tracing/trace_marker");

		for (int try = 0; ; ++try) {
			int prio = get_prio(0);
			time_t now = time(NULL);
			if (now - ts >= 120) {
				printf("Prio: %d, elapsed: %.1f minutes.\n", prio, (double)(now - start_ts) / 60);
				ts = now;
			}

			if (prio == EXPECTED_PRIO) {
				// Try a new iteration.
				break;
			}

			if (try >= 10) {
				print_threads(pid);
				printf("found race, hang...\n");
				while (1) {
					sleep(3600);
				}
			}
		}
	}

	for (int i = 0; i < thread_count; ++i) {
		pthread_join(threads[i], NULL);
	}

	return 0;
}