#include <kernel/task.h>
#include <kernel/heap.h>
#include <kernel/tty.h>

struct Task* tasks;
size_t tasks_capacity;
size_t num_tasks;

size_t current_task = 0;

void setup_tasks()
{
    tasks = malloc(sizeof(*tasks));
    tasks_capacity = 1;
    num_tasks = 1;
}

void add_task(struct Task task)
{
    asm("cli");
    if (num_tasks == tasks_capacity) {
        struct Task* temp = malloc(sizeof(*tasks) * tasks_capacity * 2);

        for (size_t i = 0; i < num_tasks; i++) {
            temp[i] = tasks[i];
        }
        
        tasks_capacity *= 2;

        free(tasks);
        tasks = temp;
    }

    tasks[num_tasks] = task;
    num_tasks++;
    asm("sti");
}

size_t ss;
size_t esp;
size_t eflags;
size_t cs;
size_t eip;

void jump_kernel(
    size_t edi,
    size_t esi,
    size_t ebp,
    size_t esp,
    size_t ebx,
    size_t edx,
    size_t ecx,
    size_t eax
);
void jump_user(
    size_t edi,
    size_t esi,
    size_t ebp,
    size_t esp,
    size_t ebx,
    size_t edx,
    size_t ecx,
    size_t eax
);

static void save_state(
    size_t eax,
    size_t ebx,
    size_t ecx,
    size_t edx,

    size_t esp,
    size_t ebp,

    size_t esi,
    size_t edi,

    size_t eip,

    size_t cs,
    size_t ds,

    size_t eflags
) {
    struct CPUState state = {
        .eax = eax,
        .ebx = ebx,
        .ecx = ecx,
        .edx = edx,

        .esp = esp,
        .ebp = ebp,

        .eip = eip,

        .esi = esi,
        .edi = edi,

        .cs = cs,
        .ds = ds,

        .eflags = eflags,
    };
    
    tasks[current_task].state = state;
    
    current_task++;
    current_task %= num_tasks;
}

void next_task() {
    struct CPUState s = tasks[current_task].state;

    ss = s.ds;
    esp = s.esp;
    eflags = s.eflags;
    cs = s.cs;
    eip = s.eip;

    if ((cs & 3) == 0) {
        jump_kernel(
            s.edi, 
            s.esi, 
            s.ebp, 
            s.esp, 
            s.ebx, 
            s.edx, 
            s.ecx, 
            s.eax
        );
    } else {
        jump_user(
            s.edi, 
            s.esi, 
            s.ebp, 
            s.esp, 
            s.ebx, 
            s.edx, 
            s.ecx, 
            s.eax
        );
    }
}

void _switch_task(
    size_t edi,
    size_t esi,
    size_t ebp,
    size_t esp_current,
    size_t ebx,
    size_t edx,
    size_t ecx,
    size_t eax,

    size_t eip_passed,
    size_t cs_passed,
    size_t eflags_passed,
    size_t esp_passed,
    size_t ds_passed
) 
{
    terminal_writestring("Switching task\n");
    if (num_tasks == 0) {
        return;
    };
    
    size_t ds;
    if ((cs & 3) == 0) {
        ds = 0x10;
        esp_current = esp_current + sizeof(size_t) * 3;
    } else {
        ds = ds_passed;
        esp_current = esp_passed;
    }

    save_state(eax, ebx, ecx, edx, esp_current, ebp, esi, edi, eip_passed, cs_passed, ds, eflags_passed);

    next_task();
}