#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <sys/wait.h>
#include <signal.h>
#include <numa.h>
#include <fcntl.h>
#include <stdbool.h>
#include <sys/resource.h>
#include <sys/time.h>

pid_t child_pid;

const char* drm_sysfs_path = "/sys/class/drm";
const char* rx6800xt_id = "fd4107589d5b1bb0\n";
const char* auto_str = "auto\n";
const char* manual_str = "manual\n";

void sig_handler(int sig)
{
	if(child_pid != 0)
		kill(child_pid, sig);
}

int find_gpu(const char* id)
{
	char path[255];
	int i;

	for(i = 0; i < 9; ++i)
	{
		snprintf(path, sizeof(path), "%s/card%i/device/unique_id", drm_sysfs_path, i);
		if(access(path, F_OK) == 0)
		{
			char id[255] = {};
			int id_fd = open(path, O_RDONLY);
			if(id_fd < 0)
				continue;
			read(id_fd, id, sizeof(id)-1);
			close(id_fd);

			if(strcmp(id, rx6800xt_id) == 0)
			{
				snprintf(path, sizeof(path), "%s/card%i/device/power_dpm_force_performance_level", drm_sysfs_path, i);
				int fd = open(path, O_RDONLY);
				return i;
			}
		}
	}
	return -1;
}

bool set_gpu_mode(const char* level, const char* mode, int card)
{
	char path[255];

	snprintf(path, sizeof(path), "%s/card%i/device/power_dpm_force_performance_level", drm_sysfs_path, card);
	int level_fd = open(path, O_WRONLY);
	if(level_fd < 0)
	{
		printf("Unable to open %s: %s\n", path, strerror(errno));
		return false;
	}
	int ret = write(level_fd, level, strlen(level));
	close(level_fd);
	if(ret != strlen(level))
	{
		printf("Unable to write %s to %s: %s\n", level, path, strerror(errno));
		return false;
	}

	snprintf(path, sizeof(path), "%s/card%i/device/pp_power_profile_mode", drm_sysfs_path, card);
	int mode_fd = open(path, O_WRONLY);
	if(mode_fd < 0)
	{
		printf("Unable to open %s: %s\n", path, strerror(errno));
		return false;
	}
	ret = write(mode_fd, mode, strlen(mode));
	close(mode_fd);
	if(ret != strlen(mode))
	{
		printf("Unable to write %s to %s: %s\n", level, path, strerror(errno));
		return false;
	}
	return true;
}

bool enable_vr_mode()
{
	int card = find_gpu(rx6800xt_id);

	return set_gpu_mode(manual_str, "4\n", card);
}

void disable_vr_mode()
{
	int card = find_gpu(rx6800xt_id);
	set_gpu_mode(auto_str, "0\n", card);
}

int main(int argc, char** argv)
{
	uid_t user = getuid();
	uid_t effective_user = geteuid();
	printf("User id %u, effective user %u\n", user, effective_user);

	for(int i = 1; i < _NSIG; ++i)
		signal(i, &sig_handler);

	if(argc < 3)
	{
		printf("Usage: %s [base_numa_node] [command] [options...]\n", argv[0]);
		return 1;
	}

	if(!enable_vr_mode())
		return 1;

	child_pid = fork();
	if(child_pid == 0)
	{
		/*int ret = setpriority(PRIO_PROCESS, 0, -20);
		if(ret)
		{
			printf("Warning: Cant set process priority\n");
		}*/

		int ret = setuid(user);
		if(ret)
		{
			printf("Unable to setuid\n");
			return 1;
		}

		ret = numa_available();
		if(ret < 0)
		{
			printf("Warning: NUMA not available\n");
		}
		else
		{
			int requested_node = atoi(argv[1]);
			if(numa_max_node() < requested_node+1)
			{
				printf("Numa nodes %i and %i requested, but only %i nodes available\n", requested_node, requested_node+1, numa_max_node()+1);
				return 1;
			}

			struct bitmask* nodes = numa_allocate_nodemask();
			numa_bitmask_setbit(nodes, requested_node);
			numa_bitmask_setbit(nodes, requested_node+1);
			printf("Using numa mask: ");
			for(size_t i = 0; i < numa_max_node()+1; ++i)
				printf("%i, ", numa_bitmask_isbitset(nodes, i));
			putc('\n', stdout);
			numa_bind(nodes);
			numa_free_nodemask(nodes);
		}

		ret = execvp(argv[2], argv+2);
		if(ret)
		{
			printf("Unable to exec %s: %s\n", argv[2], strerror(errno));
			return 1;
		}
	}
	else if(child_pid < 0)
	{
		printf("Unable to fork: %s\n", strerror(errno));
		disable_vr_mode();
		return 1;
	}
	else
	{
		int ret;
		waitpid(child_pid, &ret, 0);
		printf("Exiting vr mode\n");
		disable_vr_mode();
		return ret;
	}

	return 0;
}