//Deian Stefan
#include <sys/time.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <ucontext.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include "heap.h"
#include "sched.h"

struct sched_proc {
	volatile long state;
	int pid;
	
	ucontext_t context;

	sched_proc_t *parent;
	proc_list_t children;

	sched_waitq_t *wq; //pointer back to self on a waitq
	int is_waiting;    //flag - waiting for a child
	int exit_code;
	
	int nice; //static priority
	int prio; //dynamic priority

	unsigned long time_slice;
	unsigned long long timestamp;  //time right before sleep
	unsigned long long  s_time;    //ticks slept
	unsigned long long  cpu_time;
};



int sched_init(void (*init_fn)()) {
	int i;
	struct itimerval t_int;
	struct timeval t1;
	struct sigaction sig,sig_ps;

	// create runqueue
	if(!(_active=create_heap(SCHED_NPROC))) {
		fprintf(stderr,"sched_init(): heap init failed\n");
		exit(-1);
	}
	active=_active;
	if(!(_expired=create_heap(SCHED_NPROC))) {
		fprintf(stderr,"sched_init(): heap init failed\n");
		exit(-1);
	}
	expired=_expired;
	_rqueue=1;
	// ---------------

	for(i=0;i<SCHED_NPROC/32;i++) {
		pid_bitmap[i]=0;
	}
	g_ticks=0;
	
	// create swapper 
	if(!(swapper=malloc(sizeof(sched_proc_t)))) {
		fprintf(stderr,"sched_init():swapper malloc failed: %s\n",
				strerror(errno));
		free(_active);
		free(_expired);
		exit(-1);
	}
	swapper->time_slice=1;
	swapper->pid=0;
	if (getcontext(&swapper->context)) {
		fprintf(stderr,"sched_init(): getcontext() failed: %s\n",
				strerror(errno));
		free(swapper);
		free(_active);
		free(_expired);
		exit(-1);
	}
	if((swapper->context.uc_stack.ss_sp=mmap(0,STACK_SIZE,
	PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS,0,0))==MAP_FAILED){
		fprintf(stderr,"sched_init(): mmap() failed: %s\n",
				strerror(errno));
		free(swapper);
		free(_active);
		free(_expired);
		exit(-1);
	}
	swapper->context.uc_stack.ss_size=STACK_SIZE;
	makecontext(&swapper->context,__idle,0);
	//other stuff does not matter for swapper
	
	//create init
	if(!(init=malloc(sizeof(sched_proc_t)))) {
		fprintf(stderr,"sched_init(): init malloc failed: %s\n",
				strerror(errno));
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(_active);
		free(_expired);
		exit(-1);
	}
	init->state=SCHED_RUNNING;
	init->nice=init->prio=20;
	init->time_slice=__assign_time_slice(init);
	init->exit_code=0;
	init->s_time=0;
	init->cpu_time=0;
	init->parent=swapper;
	__proc_list_init(&init->children);
	init->wq=NULL;
	init->is_waiting=0;

	if(__assign_pid(init)<0) { 
		fprintf(stderr, "sched_init(): failed to get pid.\n");
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(init);
		free(_active);
		free(_expired);
		exit(-1);
	}
	current=init;
	// ---------------
	
	
	// waitq wait_list init, list of
	// any parent wating for a child
	sched_waitq_init(&wait_list);
	// ---------------
	
	// context related
	if (getcontext(&current->context)) {
		fprintf(stderr,"sched_init(): getcontext() failed: %s\n",
				strerror(errno));
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(init);
		free(_active);
		free(_expired);
		exit(-1);
	}
	if((current->context.uc_stack.ss_sp=mmap(0,STACK_SIZE,
	PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS,0,0))==MAP_FAILED){
		fprintf(stderr,"sched_init(): mmap() failed: %s\n",
				strerror(errno));
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(init);
		free(_active);
		free(_expired);
		exit(-1);
	}
	current->context.uc_stack.ss_size=STACK_SIZE;
	current->context.uc_link=NULL;
	makecontext(&current->context, init_fn,0);
	setcontext(&current->context);
	// ---------------
	

	// sinals related
	sigemptyset(&_mask1);
	sigaddset(&_mask1,SIGUSR1);
	sigaddset(&_mask1,SIGUSR2);
	sigaddset(&_mask1,SIGABRT);
	sigaddset(&_mask1,SIGVTALRM);

	sig.sa_handler=sched_tick;
	if(sigaction(SIGVTALRM,&sig,NULL)) {
		fprintf(stderr,"sched_init(): sigaction() failed: %s\n",
				strerror(errno));
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(init);
		free(_active);
		free(_expired);
		exit(-1);
	}
	sig_ps.sa_handler=sched_ps;
	if(sigaction(SIGABRT,&sig_ps,NULL)) {
		fprintf(stderr,"sched_init(): sigaction() failed: %s\n",
				strerror(errno));
		munmap(swapper->context.uc_stack.ss_sp,STACK_SIZE);
		free(swapper);
		free(init);
		free(_active);
		free(_expired);
		exit(-1);
	}

	t1.tv_sec=0;
	t1.tv_usec=SCHED_TICKS;
	t_int.it_interval=t1;
	t_int.it_value=t1;
	setitimer(ITIMER_VIRTUAL,&t_int,NULL);

	
	/*
	// ---------------
	// finish creating init
	// the code bellow was used before any of the ucontext related code 
	// -used setjmp and longjmp (does not work on current box 
	// linux kernel 2.6.17-rc6
	asm volatile ("leave\n\t"//ebp_main is restored
	    "popl %%eax\n\t"	 //ret_main
	    "popl %%ebx\n\t"	 //arg0 to sched_init
	    "movl %0,%%esp\n\t"  //new esp
	    "pushl $0\n\t"	 //dont need ret val
	    "pushl %%eax\n\t"	 //ret_main
	    "pushl %%ebx\n\t"	 //arg0 is init_fn
//	    "pushl %1\n"	 //init_fn
	    "ret"		 //return to 1st on stack
	    : :"r"(current->context.uc_stack.ss_sp)	 //,"r"(init_fn)
	    :"eax","ebx");
	// ---------------
	// */
}

int sched_fork() {
	volatile int inchild=0;
	sched_proc_t *child;
	long s_off;
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	
	if(!(child=malloc(sizeof(sched_proc_t)))) {
		fprintf(stderr,"pid %d sched_fork(): malloc failed: %s\n",
				current->pid,strerror(errno));
		exit(-1);
	}
	if(__assign_pid(child)<0) {
		free(child);
		fprintf(stderr,"pid %d sched_fork(): too many processes\n",
				current->pid);
				return -1;
	}

	child->parent=current;
	child->state=SCHED_READY;
	child->exit_code=0;
	child->nice=current->nice;
	child->prio=current->prio;
	child->time_slice=current->time_slice=MAX(current->time_slice/2,1);
	child->timestamp=0;
	child->wq=NULL;
	child->is_waiting=0;
	child->s_time=0;
	child->cpu_time=0;

	__proc_list_init(&child->children);
	if(!__proc_list_add(&current->children,child,1)) {
		fprintf(stderr,"pid %d sched_fork():can't add child to list\n",
				current->pid);
		return -1;
	}
	
	if((child->context.uc_stack.ss_sp=mmap(0,STACK_SIZE,
	PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANONYMOUS,0,0))==MAP_FAILED){
		fprintf(stderr,"pid %d map() failed: %s\n",child->pid,
				strerror(errno));
		free(child);
		exit(-1);
	}

	insert(active,child,child->prio,child->pid);
	inchild=1;
	memcpy(child->context.uc_stack.ss_sp,
			current->context.uc_stack.ss_sp,STACK_SIZE);
	inchild=0;
	if (getcontext(&child->context)) {
		fprintf(stderr,"pid %d getcontext() failed: %s\n",child->pid,
				strerror(errno));
		free(child);
		exit(-1);
	}
	if(inchild) { 
		sigprocmask(SIG_SETMASK,&old_mask,NULL);
		return 0;
	}

	//from /usr/include/sys/ucontext.h
	s_off=(unsigned long)child->context.uc_stack.ss_sp-
			(unsigned long)current->context.uc_stack.ss_sp;
	((unsigned long *)child->context.uc_mcontext.gregs)[6]+=s_off;
	((unsigned long *)child->context.uc_mcontext.gregs)[7]+=s_off;
	*(unsigned long *)child->context.uc_mcontext.gregs[6] += s_off;

	sigprocmask(SIG_SETMASK,&old_mask,NULL);
	return child->pid;
}

int sched_waitq_init(sched_waitq_t *wq) {
	__proc_list_init(wq);
	return 0;
}

int sched_exit(int code) {
	unsigned long long now=sched_gettick();
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	current->state=SCHED_ZOMBIE;
	current->exit_code=code;
	delete(active,current->pid,NULL);
	if(current->parent->is_waiting) {
		current->parent->is_waiting=0;
		__proc_list_del(current->parent->wq);
		current->parent->wq=NULL;
		current->parent->s_time+=now-current->parent->timestamp;
		current->parent->prio=__effective_prio(current->parent);
		current->parent->time_slice=
					__assign_time_slice(current->parent);
		current->parent->state=SCHED_READY;
		insert(active,current->parent,
				current->parent->prio,current->parent->pid);
	}
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
	sched_switch();
}

int sched_wait(int *code) {
	int rc;
	proc_list_t *tmp;
	sched_proc_t *child=NULL;
	sigset_t old_mask;

retry_wait:
	rc=0;
	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);

	if(current->children.next==&current->children) {
		return -1;
	}
	for(tmp=current->children.next;tmp!=&current->children;tmp=tmp->next) {
		if(tmp->proc->state==SCHED_ZOMBIE) {
			child=__proc_list_del(tmp);
			break;
		}
	}
	if(child) {
		rc=child->pid;
		if(__remove_proc(child,code)) {
			return -1;
		}
		//o snap! dead kid!
	} else {
		current->is_waiting=1;
		sigprocmask(SIG_SETMASK,&old_mask,NULL);
		sched_sleep(&wait_list);
		goto retry_wait; //woke up, should have a zombie kid
	}
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
	return rc;
}

int sched_sleep(sched_waitq_t *waitq) {
	unsigned long long now=sched_gettick();
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	if(!(current->wq=__proc_list_add(waitq,current,current->is_waiting))) {
		return -1;
	}
	current->timestamp=now;
	current->state=SCHED_SLEEPING;
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
	sched_switch();
}

int sched_wakeup(sched_waitq_t *waitq) {
	unsigned long long now=sched_gettick();
	sched_waitq_t *wq_c;
	sched_proc_t *tmp;
	sigset_t old_mask;
	int flag=0;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	for(wq_c=waitq->next;wq_c!=waitq;wq_c=waitq->next){
		tmp=__proc_list_del(wq_c);
		tmp->s_time+=now-tmp->timestamp;
		tmp->prio=__effective_prio(tmp);
		tmp->time_slice=__assign_time_slice(tmp);
		if(tmp->prio<current->prio) { flag=1 ;}
		tmp->state=SCHED_READY;
		insert(active,tmp,tmp->prio,tmp->pid);
	}

	sched_waitq_init(waitq);
	if(flag&&current->pid!=0) {
		current->prio=__effective_prio(current);
		current->time_slice=__assign_time_slice(current);
		current->state=SCHED_READY;
		insert(active,current,current->prio,current->pid);
	}
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
	if(flag) { sched_switch(); }
}

int sched_nice(int niceval) {
	current->nice+=niceval;
	current->nice=THRESH(current->nice,0,39);
	return current->nice;
}

int sched_getpid() {
	return current->pid;
}

int sched_getppid() {
	return current->parent->pid;
}

unsigned long long sched_gettick() {
	return g_ticks;
}

int sched_ps() {
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	printf("  pid\t ppid\t  state \t   stack \t  waitq\t");
	printf("       sprio   dprio\t     cpu  \n");
	__ps_all(init);
	printf("\n\n");
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
}

int sched_switch() {
	sched_proc_t *prev,*next;
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	prev=current;
	if(!(next=delete_min(active,NULL,NULL))) {
		__swap_rqueue();
		if(!(next=delete_min(active,NULL,NULL))) {
			next=swapper;
		}
	}
	next->state=SCHED_RUNNING;
	current=next;
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
//	uncomment the next two lines to see it all happen
//	if(current->prio)
//		sched_ps(); 
	swapcontext(&prev->context, &current->context);
}

int sched_tick() {
	sigset_t old_mask;

	sigprocmask(SIG_BLOCK,&_mask1,&old_mask);
	g_ticks++;
	current->time_slice--;
	if(!current->time_slice) {
		current->s_time>>=1;
		current->prio=__effective_prio(current);
		current->state=SCHED_READY;
		current->time_slice=__assign_time_slice(current);
		insert(expired,current,current->prio,current->pid);
		sigprocmask(SIG_SETMASK,&old_mask,NULL);
		sched_switch();
	}
	current->cpu_time++;
	sigprocmask(SIG_SETMASK,&old_mask,NULL);
}

int __assign_time_slice(sched_proc_t *proc) {
	return (40-proc->nice)*((proc->nice<20)?10:5);
}

int __effective_prio(sched_proc_t *proc) {
	return MAX(0,MIN(proc->nice-__bonus(proc)+5,39));
}

int __assign_pid(sched_proc_t *proc) {
	int i,j,pid;
	for(i=0;i<SCHED_NPROC/32;i++) {
		if(pid_bitmap[i]==0xffffffff) continue;
		asm("bsfl %1,%0" :"=r" (j) :"rm" (~pid_bitmap[i]));
		asm("btsl %1,%0 \n" : :"m"(pid_bitmap[i]),"Ir"(j));
		pid=j+i*32+1;
		proc->pid=pid;
		return pid;
	}
	return -1;
}

int __free_pid(sched_proc_t *proc) {
	int i=0;
	while((proc->pid-32)>32) { 
		proc->pid-=32;
		i++;
	}
	asm("btcl %1,%0 \n" : :"m"(pid_bitmap[i]),"Ir"(proc->pid-1));
	proc->pid=0;
}

//The algorithm bellow is based on the algorithm used in the 
//Linux 2.6 kernel of swapping the runqueues
int __swap_rqueue() {
		if(_rqueue) {
			active=_expired;
			expired=_active;
			_rqueue=0;
		} else {
			active=_active;
			expired=_expired;
			_rqueue=1;
		}
		return _rqueue;
}

int __bonus(sched_proc_t *proc) {
	//this can be way smarter
	if(proc->s_time<=0) {
		return 0;
	} else if(proc->s_time<=1000) {
		return 1;
	} else if(proc->s_time<=2000) {
		return 2;
	} else if(proc->s_time<=3000) {
		return 3;
	} else if(proc->s_time<=4000) {
		return 4;
	} else if(proc->s_time<=5000) {
		return 5;
	} else if(proc->s_time<=6000) {
		return 6;
	} else if(proc->s_time<=7000) {
		return 7;
	} else if(proc->s_time<=8000) {
		return 8;
	} else if(proc->s_time<=9000) {
		return 9;
	}
	return 10;
}

int __proc_list_init(proc_list_t *l) {
	l->proc=NULL;
	l->next=l;
	l->prev=l;
}

//flag==1 returns pointer to self on the list -- EXCLUSIVE wakeup(wait for chid)
//flag==0 returns pointter to begining of list
proc_list_t *__proc_list_add(proc_list_t *l,sched_proc_t *proc,int flag) {
	proc_list_t *tmp;
	if(!(tmp=malloc(sizeof(proc_list_t)))) {
		fprintf(stderr,"malloc() failed to creat list node: %s\n",
			strerror(errno));
		exit(-1);
	}
	tmp->proc=proc;
	tmp->next=l;
	tmp->prev=l->prev;
	tmp->next->prev=tmp;
	tmp->prev->next=tmp;
	return (flag)?tmp:l;
}

sched_proc_t * __proc_list_del(proc_list_t *l) {
	sched_proc_t *proc=l->proc;
	l->next->prev=l->prev;
	l->prev->next=l->next;
	free(l);
	return proc;
}

int __idle() {
	while(1){}
}

int __remove_proc(sched_proc_t *proc,int *code) {
	proc_list_t *cq_c;
	sched_proc_t *child;
	if(code) {
		*code=proc->exit_code;
	}
	//add all children to init->children
	for(cq_c=proc->children.next;cq_c!=&proc->children;
						cq_c=proc->children.next){
		child=__proc_list_del(cq_c);
		child->parent=init;
		__proc_list_add(&init->children,child,1);
	}
	if(munmap(proc->context.uc_stack.ss_sp,STACK_SIZE)) {
		fprintf(stderr,"munmap faild to unmap stack of pid %d: %s\n",
				proc->pid,strerror(errno));
		exit(-1);
	}
	__free_pid(proc);
	free(proc);
	return 0;
}

int __ps(sched_proc_t *proc){ 
	if(proc->pid!=0) {
	printf("%4d\t%4d\t%s\t0x%08x\t0x%08x\t%2d\t%2d\t%8llu\n",
		proc->pid,proc->parent->pid,STATE(proc),
		proc->context.uc_stack.ss_sp,
		((proc->state==SCHED_SLEEPING)?proc->wq:0),
		proc->nice,proc->prio,proc->cpu_time);
	}
}

int __ps_all(sched_proc_t *proc){ 
	//traverse `tree'!
	proc_list_t *tmp;
	__ps(proc);
	for(tmp=proc->children.next;tmp!=&proc->children;tmp=tmp->next) {
		__ps_all(tmp->proc);
	}
}