#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/stat.h>

#define CPUSET(x) "/sys/fs/cgroup/cpuset/"x

#define ACTION_ENABLE 1
#define ACTION_DISABLE 2

char *readLine(FILE *fp, char *buffer) {
	int ch;
	int i = 0;
	size_t buff_len = 0;
	buffer = malloc(buff_len + 1);
	if(!buffer) {
		return NULL;
	}
	while((ch = fgetc(fp)) != '\n' && ch != EOF) {
	buff_len++;
	void *tmp = realloc(buffer, buff_len + 1);
	if(tmp == NULL) {
		free(buffer);
		return NULL;
	}
	buffer = tmp;
	buffer[i] = (char) ch;
	i++;
	}
	buffer[i] = '\0';
	if(ch == EOF && (i == 0 || ferror(fp))) {
		free(buffer);
		return NULL;
	}
	return buffer;
}

int moveTasks(char *src, char *dst) {
	FILE *fd1, *fd2;
	char *pid;
	fd1 = fopen(src, "rb");
	if(fd1 == NULL) {
		return 255;
	}
	while((pid = readLine(fd1, 0)) != NULL) {
		if(strtol(pid, (char **)NULL, 10) != getppid()) {
			fd2 = fopen(dst, "wb");
			if(fd2 == NULL) {
				return 255;
			}
			fprintf(fd2, "%s\n", pid);
			fclose(fd2);
			free(pid);
		}
	}
	fclose(fd1);
	return 0;
}

int setGroup(char *dest, char *data) {
	FILE *fd;
	fd = fopen(dest, "wb");
	if(fd == NULL) {
		return 255;
	}
	fprintf(fd, "%s", data);
	fclose(fd);
	return 0;
}

int disableGroup() {
	struct stat sb = {0};
	printf("Disabling CGROUPS ... ");
	if(stat(CPUSET("host"), &sb) == 0) {
		if(moveTasks(CPUSET("host/tasks"), CPUSET("tasks")) != 0) {
			return 255;
		}
		rmdir(CPUSET("host"));
		printf("OK!\n");
	} else {
		printf("SKIPPED!\n");
	}
	return 0;
}

int enableGroup(char *cpu, char *mem) {
	struct stat sb = {0};
	printf("Enabling CGROUPS ... ");
	if(stat(CPUSET("host"), &sb) == -1) {
		mkdir(CPUSET("host"), 0755);
	}
	if(setGroup(CPUSET("host/cpuset.cpus"), cpu) != 0) {
		printf("FAIL!\n");
		return 255;
	}
	if(setGroup(CPUSET("host/cpuset.mems"), mem) != 0) {
		printf("FAIL!\n");
		return 255;
	}
	if(moveTasks(CPUSET("tasks"), CPUSET("host/tasks")) != 0) {
		printf("FAIL!\n");
		return 255;
	}
	printf("OK!\n");
	return 0;
}

int main(int argc, char *argv[]) {
	int opt;
	int action = 0;
	char *cpu = {0};
	char *mem = {0};
	while((opt = getopt(argc, argv, ":edc:m:")) != -1) {
		switch(opt) {
			case 'e':
				action = ACTION_ENABLE;
				break;
			case 'd':
				action = ACTION_DISABLE;
				break;
			case 'c':
				cpu = optarg;
				break;
			case 'm':
				mem = optarg;
				break;
			case ':':
				printf("Option needs a value!\n");
				return 1;
			default:
				printf("Invalid option: -%c!\n", optopt);
				return 1;
		}
	}
	if(action == ACTION_ENABLE) {
		if(!cpu || !mem) {
			printf("No CPU or memory pinning specified!\n");
			return 1;
		}
		return enableGroup(cpu, mem);
	} else if(action == ACTION_DISABLE) {
		return disableGroup();
	} else {
		printf("No action specified!\n");
	}
	return 0;
}