#include <lib/tasks.h>
#include <lib/memory.h>

#include <memory/mm.h>

#include <stddef.h>

#include "process.h"
#include "scheduler.h"

static process_t* process_table;
static spinlock_t process_table_lock = {0};

static int current_process = 0;
static int total_processes = 0;

extern uint32_t fetch_cr3();

#define IDLE_TASK_PID 1
static void idle_task() {
  while (1) asm volatile("hlt");
}

uint32_t schedule_next_task(uint32_t current_esp) {
    spin_lock(&process_table_lock);
    process_table[current_process].esp = current_esp;

    for (int i = 0; i < total_processes; i++) {
        if (process_table[i].state == STATE_SLEEPING) {
            if (process_table[i].sleep > 0) {
                process_table[i].sleep--;
            }

            if (process_table[i].sleep == 0) {
                process_table[i].state = STATE_READY;
            }
        }
    }

    int next_process = current_process;
    int found_ready = 0;
    while (1) {
        next_process = (next_process + 1) % total_processes;
        if (next_process != IDLE_TASK_PID && process_table[next_process].state == STATE_READY) {
            found_ready = 1;
            break;
        }

        if (next_process == current_process) {
            break; 
        }
    }

    if (!found_ready) {
        if (process_table[current_process].state == STATE_READY) {
            next_process = current_process;
        } else {
            next_process = IDLE_TASK_PID;
        }
    }

    current_process = next_process;

    if (fetch_cr3() != process_table[current_process].cr3) {
        extern void load_pd(uint32_t table);
        load_pd(process_table[current_process].cr3);
    }

    uint32_t rtn = process_table[current_process].esp;
    spin_unlock(&process_table_lock);

    return rtn;
}

void init_scheduler() {
    process_table = kalloc(16 * sizeof(process_t));
    total_processes = 0;

    process_t* process = &process_table[total_processes];
    process->pid = total_processes;
    process->esp = 0;
    process->cr3 = fetch_cr3();
    process->state = STATE_READY;
    process->sleep = 0;
    total_processes++;

    create_task((uint32_t)idle_task, fetch_cr3());

    current_process = 0;
}

int create_task(uint32_t entry_point, uint32_t cr3) {
    if (total_processes >= 16) {
        return 0;
    }

    spin_lock(&process_table_lock);
    process_t* process = &process_table[total_processes];
    process->pid = total_processes;
    process->cr3 = cr3;
    process->state = STATE_READY;
    process->sleep = 0;

    void* stack_bottom = kalloc(4096);
    uint32_t* esp = (uint32_t*)((uint32_t)stack_bottom + 4096);

    *(--esp) = 0x0202;        // EFLAGS   
    *(--esp) = 0x08;          // CS
    *(--esp) = entry_point;   // EIP

    *(--esp) = 0;             // EAX
    *(--esp) = 0;             // ECX
    *(--esp) = 0;             // EDX
    *(--esp) = 0;             // EBX
    *(--esp) = 0;             // ESP
    *(--esp) = 0;             // EBP
    *(--esp) = 0;             // ESI
    *(--esp) = 0;             // EDI

    process->esp = (uint32_t)esp;
    spin_unlock(&process_table_lock);

    total_processes++;
    return 1;
}

void sleep(uint32_t ticks) {
    if (ticks == 0) return;

    spin_lock(&process_table_lock);
    process_table[current_process].sleep = ticks;
    process_table[current_process].state = STATE_SLEEPING;
    spin_unlock(&process_table_lock);

    yield();
}