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

static pid_t gettid()
{
	return syscall(SYS_gettid);
}

#define LOOP_FOREVER while (1) sleep(99999999);

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

const char USAGE[] = "\
cs_clone <options>\n\
\n\
Version: 1.0\n\
\n\
options:\n\
	-v <num> : verbose\n\
	-P <num> : num processes\n\
	-T <num> : num threads\n\
\n\
	-s       : core sched\n\
	-d <num> : delay\n\
\n\
	-h : help\n\
";


#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);
}


#define STACK_SIZE (1024 * 1024)

struct child_args {
	int sync_fd;
	long child_num;
	int sleep_amt;
	int cnt;
	int core_sched;
	int verbose;
};


static int child_func_thread(void *arg)
{
	struct child_args *ca = (struct child_args *)arg;
	int res;
	int s;

	if (ca->verbose)
		printf("Child_Thread[%lu] (pid=%d/tid=%d) sleeping for '%d' seconds\n", ca->child_num, getpid(), gettid(), ca->sleep_amt);

	for (s = 0; s < ca->sleep_amt - 2; ++s) {
		volatile uint64_t i = 0;
		int j;
		for (j = 0; j < 89990000; ++j)
			i = (i + 1 + j) & 0x0fffffffffffffff;
		sleep(1);
	}
	if (ca->verbose)
		printf("Child [%lu] done\n", ca->child_num);

	if ((res = eventfd_write(ca->sync_fd, 1)) == -1)
		handle_error("(child) write eventfd");

	LOOP_FOREVER
	return 0;
}


static void create_threads(int sync_fd, int verbose, int delay, unsigned long num_threads)
{
	unsigned long thd_num;
	for (thd_num = 0; thd_num < num_threads; ++thd_num) {
		pid_t tid;

		void *child_stack = malloc(STACK_SIZE);
		if (!child_stack)
			handle_error("child stack allocate");

		struct child_args *ca = calloc(1, sizeof(struct child_args));
		if (!ca)
			handle_error("child args allocate");

		ca->child_num = thd_num;
		ca->sync_fd = sync_fd;
		ca->sleep_amt = delay;
		ca->verbose = verbose;

		if (verbose)
			printf("creating thread %lu\n", thd_num);
		tid = clone(child_func_thread, child_stack + STACK_SIZE, CLONE_THREAD|CLONE_SIGHAND|CLONE_FS|CLONE_VM|CLONE_FILES , ca);
		if (tid == -1)
			handle_error("clone thread");
		sleep(0);
	}
	sleep(0);
}

static int wait_for_children(int sync_fd, int verbose, int num_children)
{
	unsigned long num = 0;

	if (verbose) {
		sleep(1);
		printf("\nParent (pid=%d/tid=%d) waiting for '%d' children\n", getpid(), gettid(), num_children);
	}

	while (num < (unsigned int)num_children) {
		int res;
		eventfd_t u;

		if ((res = eventfd_read(sync_fd, &u)) == -1)
			handle_error("(parent) read eventfd");
		num += u;
	}
	return num;
}

static int child_func_process(void *arg)
{
	int res;
	struct child_args *ca = (struct child_args *)arg;

	if (ca->verbose)
		printf("Child_Process[%lu] (pid=%d) sleeping for '%d' seconds\n", ca->child_num, getpid(), ca->sleep_amt);

	if (ca->core_sched) {
		if (prctl(PR_SCHED_CORE_SHARE, PR_SCHED_CORE_CREATE, getpid(), 0, 0) < 0)
			handle_error("spid sched_core failed");

		if (ca->verbose)
			printf("[%d] - core sched enabled\n", getpid());
	}

	if (ca->cnt) {
		int sync_fd;
		if ((sync_fd = eventfd(0, 0)) == -1)
			handle_error("eventfd sync_fd");

		create_threads(sync_fd, ca->verbose, ca->sleep_amt, ca->cnt);

		sleep(ca->sleep_amt ? ca->sleep_amt : 1);
		wait_for_children(sync_fd, ca->verbose, ca->cnt);
	} else {
		sleep(ca->sleep_amt);
	}

	if ((res = eventfd_write(ca->sync_fd, 1)) == -1)
		handle_error("(child) write eventfd");

	sleep(20);
	return 0;
}


static unsigned char child_func_process_stack[STACK_SIZE];
static struct child_args cpa;

static void create_processes(int sync_fd, int verbose, int core_sched, int delay, int num_processes, int num_threads)
{
	cpa.sync_fd = sync_fd;
	cpa.child_num = -1;
	cpa.sleep_amt = delay;
	cpa.cnt = num_threads;
	cpa.core_sched = core_sched;
	cpa.verbose = verbose;

	int proc_num;
	for (proc_num=0; proc_num < num_processes; ++proc_num) {
		cpa.child_num = proc_num;
		if (verbose)
			printf("creating process %d\n", proc_num);
		clone(child_func_process, child_func_process_stack + STACK_SIZE, SIGCHLD, &cpa);
	}
}

int main(int argc, char *argv[])
{
	int delay = 10;
	int verbose = 0;
	int core_sched = 0;
	int num_processes = 1;
	unsigned long num_threads = 1;

	int opt;
	while ((opt = getopt(argc, argv, ":hv:sT:P:d:")) != -1) {
		switch (opt) {
			case 'P': num_processes = strtol(optarg, NULL, 10); break;
			case 'T': num_threads = strtoul(optarg, NULL, 10); break;
			case 's': core_sched = 1; break;
			case 'v': verbose = strtol(optarg, NULL, 10); break;
			case 'd': delay = strtol(optarg, NULL, 10); break;
			case 'h': printf(USAGE);  exit(EXIT_SUCCESS);
			default: handle_usage(20, "unknown option");
		}
	}

	if (optind != argc)
		handle_usage(1, "too many args");

	if (core_sched) {
		if (prctl(PR_SCHED_CORE_SHARE, PR_SCHED_CORE_CREATE, getpid(), 0, 0) < 0)
			handle_error("sched_core failed");

		if (verbose)
			printf("[%d] - core sched enabled\n", getpid());
	}

	int sync_fd;
	if ((sync_fd = eventfd(0, 0)) == -1)
		handle_error("eventfd sync_fd");

	if (num_processes)
		create_processes(sync_fd, verbose, core_sched, delay, num_processes, num_threads);

	if (num_threads)
		create_threads(sync_fd, verbose, delay, num_threads);

	wait_for_children(sync_fd, verbose, num_threads + num_processes);
	exit(EXIT_SUCCESS);
}
