#define _GNU_SOURCE
//#include <sys/capability.h>
#include <string.h>
#include <sys/eventfd.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <limits.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/prctl.h>
#include <sys/syscall.h>

#ifndef PR_SCHED_CORE_SHARE
#define PR_SCHED_CORE_SHARE 		59
# define PR_SCHED_CORE_CLEAR		0
# define PR_SCHED_CORE_SHARE_FROM	1
# define PR_SCHED_CORE_SHARE_TO		2
#endif

const char USAGE[] = "\
cs_exec <options> [command]\n\
\n\
Version: 1.0\n\
\n\
options:\n\
	-p : specify pid\n\
	-m : mode, 0-clear, 1-from, 2-to, 3-pull from src_pid send to dest_pid\n\
	-M : mode, \"clear\", \"from\", \"to\", \"copy\" -copy cookie from src_pid to dest_pid\n\
\n\
	-E : test error conditions. implies '-s'\n\
	-h : help\n\
\n\
For mode == 0, specify a command to exec.\n\
\n\
Example of setting one pid cs to another pid cs\n\
	cs_exec -s -m 3 <src_pid> <dest_pid>\n\
";


#define handle_expected_error(msg) \
	__handle_error(__LINE__, msg)

#define handle_error(msg) \
	__handle_error(__LINE__, msg)

static void __handle_error(int ln, char *msg) {
	printf("Error: line: %d\n", ln);
	perror(msg);
	exit(EXIT_FAILURE);
}

static void handle_usage(int rc, char *msg) {
	puts(USAGE);
	puts(msg);
	putchar('\n');
	exit(rc);
}

int main(int argc, char *argv[])
{
	int mode = 0;
	int pid = -1;
	int opt;

	while ((opt = getopt(argc, argv, ":hSp:m:M:")) != -1) {
		switch (opt) {
			case 'p': pid = (int)strtoul(optarg, NULL, 10); break;
			case 'm': mode = (int)strtoul(optarg, NULL, 10); break;
			case 'M':
				if (!strncmp("clear", optarg, 5))
					mode = PR_SCHED_CORE_CLEAR;
				else if (!strncmp("from", optarg, 4))
					mode = PR_SCHED_CORE_SHARE_FROM;
				else if (!strncmp("to", optarg, 2))
					mode = PR_SCHED_CORE_SHARE_TO;
				else if (!strncmp("copy", optarg, 2))
					mode = 3;
				else
					handle_usage(30, "-M unknown option");
				break;
			case 'h': printf(USAGE);  exit(EXIT_SUCCESS);
			default: handle_usage(20, "unknown option");
		}
	}

	printf("Mode: %d\n", mode);

	if (pid < 0)
		pid = getpid();

	switch (mode) {
		case PR_SCHED_CORE_CLEAR:
			if (prctl(PR_SCHED_CORE_SHARE, PR_SCHED_CORE_CLEAR, pid, 0, 0) < 0)
				handle_error("(2) sched_core_clear failed");
			break;

		case PR_SCHED_CORE_SHARE_TO:
		case PR_SCHED_CORE_SHARE_FROM:
			if (prctl(PR_SCHED_CORE_SHARE, mode, pid, 0, 0) < 0)
				handle_error("(0 1) sched_core failed");
			break;

		case 3:
			if (optind + 2 != argc)
				handle_usage(2, "mode 3 needs <to_pid> <from_pid>");

			int spid = (int)strtol(argv[optind], NULL, 10);
			if (prctl(PR_SCHED_CORE_SHARE, PR_SCHED_CORE_SHARE_FROM, spid, 0, 0) < 0)
				handle_error("spid sched_core failed");

			pid = (int)strtol(argv[optind+1], NULL, 10);
			if (prctl(PR_SCHED_CORE_SHARE, PR_SCHED_CORE_SHARE_TO, pid, 0, 0) < 0)
				handle_error("dpid sched_core failed");
			break;

		default:
			handle_usage(1, "");
			break;
	}

	if (mode == PR_SCHED_CORE_SHARE_FROM) {
		if (optind < argc)
			execvp(argv[optind], &argv[optind]);
	}
}
