#include <multitask.h>
#include <zmem/early_heap.h>
#include <cpu/gdt.h>
#include <string.h>
#include <zmem/memory.h>

volatile task_t * first_task;
volatile task_t * last_task;
volatile task_t * actual_task;

extern void idle();

int max_pid;

task_t * create_task(entrypoint * entry, uint32 gdt, uint8 stack_size){
	if(!stack_size)
		return NULL;
	task_t * task = ealloc(sizeof(task_t));
	if(!task)
		return NULL;
	
	esp_t stack= ealloc(stack_size);
	if(!stack){
		efree(task);
		return NULL;
	}
	regs_t * state = ealloc(sizeof(regs_t));
	if(!state){
			efree(task);
			efree(stack);
			return NULL;
	}
	state->ds = KERNEL_DS;
	state->es = KERNEL_DS;
	state->fs = KERNEL_DS;
	state->gs = KERNEL_DS;
	state->esp = stack;
	state->eip = entry;
	state->cs = KERNEL_CS;
	state->eflags = 0x202;
	
	task->pid = max_pid;
	task->state = TASK_READY;
	task->cpu_state = state;
	task->cr3 = get_kernel_cr3();
	task->next = NULL;
	
	return task;
}

void add_task(task_t* task){
	last_task->next = task;
	last_task = task;
	max_pid++;
}

void regs_copy(regs_t * dest,const regs_t * src){
	dest->gs	= src->gs	;
	dest->fs	= src->fs	;
	dest->es	= src->es	;
	dest->ds	= src->ds	;

	dest->edi	= src->edi	;
	dest->esi	= src->esi	;
	
	dest->ebp	= src->ebp	;
	dest->esp	= src->esp	;

	dest->eax	= src->eax	;
	dest->ebx	= src->ebx	;
	dest->ecx	= src->ecx	;
	dest->edx	= src->edx	;

	dest->int_no	= src->int_no	;
	dest->err_code	= src->err_code	;

	dest->eip	= src->eip	;
	dest->cs	= src->cs	;
	dest->eflags	= src->eflags	;
}

regs_t * task_switch(struct regs* r){
	if(!is_tasking_set)
		return;
	task_t * next = actual_task->next;
	if(next==NULL){
		next = first_task;
	}
	
	regs_copy(actual_task->cpu_state, r);
	actual_task->state = TASK_READY;
	next->state = TASK_EXECUTE;
	actual_task = next;
	return next->cpu_state;
}

void init_tasking(uint32 gdt, entrypoint * entry, uint64 sched_freq){
    asm volatile("cli");
	max_pid = 0;
    task_t * newtask = create_task(&idle, gdt, 4);
	first_task = last_task = newtask;
	max_pid++;
	actual_task = newtask;
	actual_task->state = TASK_EXECUTE;
    is_tasking_set = true;
    asm volatile("sti");
}