#include <iostream>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <cstring>
#include <cstdlib>
#include <ctime>
#include <errno.h>
using namespace std;

struct Shared{
	int next_task;
	int completed[64];
};

void random_task(int id){
	srand(time(nullptr)^getpid());
	int delay=1+rand()%3;
	cout<<"[WORKER] PID "<<getpid()<<" take up task "<<id<<" (Delay: "<<delay<<"s)\n";
	sleep(delay);
}

int main(){
	cout<<"================================================================\n";
	cout<<" [***] ADVANCED SKO SCHEDULER & PROCESS POOL [***]       \n";
	cout<<"================================================================\n\n";
	int n=3;
	int total_tasks=10;
	pid_t root=getpid();
	cout<<"[ROOT] The starting process started with PID = "<<root<<"\n";
	if(prctl(PR_SET_CHILD_SUBREAPER,1)!=0){
		cerr<<"Error in prctl subreaper.\n";
		return 1;
	}
	Shared* shm=(Shared*) mmap(nullptr,sizeof(Shared),PROT_READ|PROT_WRITE,MAP_SHARED|MAP_ANONYMOUS,-1,0);
	memset(shm,0,sizeof(Shared));
	setpgid(0,0);
	pid_t pgid=getpgrp();
	cout<<"[ROOT] Creating Pocess Group: "<<pgid<<"\n";
	cout<<"[ROOT] Begining Exponential Forking...\n\n";
	for(int i=0;i<n;i++){
		fork();
	}
	if(getpid()!=root){
		setpgid(0,pgid);
		while(1){
			int task=__sync_fetch_and_add(&shm->next_task,1);
			if(task>=total_tasks){break;}
			random_task(task);
			shm->completed[task]=getpid();
		}
		_exit(0);
	}
	int status;
	pid_t pid;
	cout<<"\n[ROOT] Waiting for every worker to finish...\n";
	while(1){
		pid=waitpid(-1,&status,0);
		if(pid>0){
			cout<<"[ROOT] Reaped the process with PID "<<pid<<"\n";
			continue;
		}
		if(pid==-1&&errno==ECHILD){
			break;
		}
		break;
	}
	cout<<"\n======================================================\n";
	cout<<" [***] TASK COMPLETION REPORT [***]       \n";
	cout<<"======================================================\n";
	for(int i=0;i<total_tasks;i++){
		if(shm->completed[i]!=0){
			cout<<"Task "<<i<<" -> Completed by PID: "<<shm->completed[i]<<"\n";
		}
	}
	munmap(shm,sizeof(Shared));
	killpg(pgid,SIGTERM);
	cout<<"\n[ROOT] Terminating system. All clear!\n";
	return 0;
}
