#include "task.h"
#include "alloc.h"
#include "debug.h"
#include "interrupts.h"
#include "kpanic.h"
#include "kprintf.h"
#include "memory.h"
#include "string.h"

#include <stdbool.h>
#include <stdint.h>

uint32_t esp_backup;

struct task *create_task(uint8_t uid)
{
	static uint32_t pid = 1;
	switch_pd(kernel_pd, (uint32_t *)((uint32_t)kernel_pd - HEAP_END));
	struct task *new_task = vmalloc(sizeof(struct task));
	if (!new_task) {
		switch_pd(current_task->heap, current_task->cr3);
		return NULL;
	}
	switch_pd(current_task->heap, current_task->cr3);
	new_task->status = RUN;
	new_task->uid = uid;
	new_task->esp = new_task->esp0 + STACK_SIZE;
	new_task->pid = pid++;

	// Allocate new pd
	new_task->heap = alloc_pages(4096, (void **)&new_task->cr3);
	if (!new_task->heap) {
		vfree(new_task);
		return NULL;
	}
	new_task->heap[768] = ((uint32_t)boot_page_table1 - HEAP_END) | 0x03;
	// memcpy(new_task->heap, current_task->heap, 4096);
	current_pd = new_task->heap;
	// switch_pd(new_task->heap, new_task->cr3);

	// Allocate new stack on the newly allocated pd
	new_task->esp0 = alloc_pages(STACK_SIZE, NULL);
	current_pd = kernel_pd;
	// switch_pd(current_task->heap, current_task->cr3);
	if (!new_task->esp0) {
		vfree(new_task);
		free_pages(new_task->heap, 4096);
		return NULL;
	}

	new_task->next = current_task->next;
	new_task->prev = current_task;
	current_task->next = new_task;
	if (current_task->prev == current_task)
		current_task->prev = new_task;

	new_task->signals.pending = SIG_IGN;

	return new_task;
}

int8_t create_kernel_task(void)
{
	struct task *new_task = vmalloc(sizeof(struct task));
	if (!new_task)
		return -1;
	new_task->status = RUN;
	new_task->uid = 0;
	new_task->pid = 0;
	new_task->heap = kernel_pd;
	new_task->cr3 = (uint32_t *)((uint32_t)kernel_pd - HEAP_END);
	new_task->prev = new_task;
	new_task->next = new_task;
	new_task->signals.pending = SIG_IGN;
	current_task = new_task;
	return 0;
}

void exec_fn(void (*fn)(void))
{
	struct task *new_task = create_task(OWNER_KERNEL);
	if (!new_task)
		kpanic("failed to create new task");
	new_task->eip = (uint32_t *)fn;
}

void zombify_task(struct task *task)
{
	cli(); // Technically useless
	free_pages(task->heap, 4096);
	free_pages(task->esp0, STACK_SIZE);
	task->esp0 = NULL;
	task->heap = NULL;
	toris();
}

void remove_task(struct task *task)
{
	cli();
	struct task *left = task->prev;
	struct task *right = task->next;

	if (task->child) {
		remove_task(task->child);
		task->child = NULL;
	}
	if (task->heap)
		free_pages(task->heap, 4096);
	if (task->esp0)
		free_pages(task->esp0, STACK_SIZE);
	left->next = right;
	right->prev = left;
	vfree(task);
	toris();
}

struct task *copy_task(const struct task *task)
{
	struct task *new_task = create_task(task->uid);
	if (!new_task)
		return NULL;
	memcpy(new_task->esp0, task->esp0, STACK_SIZE);
	new_task->esp = new_task->esp0 + (task->esp - task->esp0);
	new_task->status = task->status;
	return new_task;
}

void exit_task(void)
{
	cli();
	if (current_task->daddy && current_task->daddy->status != WAIT)
		current_task->status = ZOMBIE;
	else {
		if (current_task->daddy->status == WAIT)
			current_task->daddy->status = RUN;
		current_task->status = STOPPED;
	}
	toris();
	scheduler();
}