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

#include <memory/mm.h>

#include <stddef.h>

#include "common.h"
#include "process.h"
#include "scheduler.h"
#include "gdt.h"

static process_t* process_table;
  
static int current_process = 0;
static int total_processes = 0;

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

static void lock_scheduler() {
    disable_interrupts();
}

static void unlock_scheduler() {
    enable_interrupts();
}

// called from isr and yield, interrupts are already disabled and re enabled in the isr/yield func
uint32_t switch_context(uint32_t current_esp) {
    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 (process_table[current_process].flags & PROCESS_FLAG_USER ) {
        tss.esp0 = process_table[current_process].kstack_top;
    }

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

    return process_table[current_process].esp;
}

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;
    process->flags = PROCESS_FLAG_KERNEL;
    total_processes++;

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

    current_process = 0;
}

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

    lock_scheduler();

    // default to a new task, but if we find a dead one, re use it
    int reuse = 0;
    process_t* process = &process_table[total_processes];
    uint32_t pid = total_processes;
    for (uint32_t i = 0; i < total_processes; i++) {
        if (process_table[i].state == STATE_DEAD) {
            process = &process_table[i];
            pid = i;
            reuse = 1;
        }
    }

    process->pid = pid;
    process->cr3 = cr3;
    process->state = STATE_READY;
    process->sleep = 0;
    process->flags = user ? PROCESS_FLAG_USER : PROCESS_FLAG_KERNEL;

    void* kstack_bottom = kalloc(PAGE_SIZE);
    uint32_t kstack_top = (uint32_t)kstack_bottom + PAGE_SIZE;
    process->kstack_top = kstack_top;

    uint32_t* esp = (uint32_t*)kstack_top;
    if (user) {
        *(--esp) = 0x23;            // User Data Segment (SS) with RPL 3 (0x20 | 3)
        *(--esp) = u_esp;      // User Stack Pointer (ESP) - location mapped in user space
        *(--esp) = 0x0202;          // EFLAGS (Interrupts enabled)
        *(--esp) = 0x1B;            // User Code Segment (CS) with RPL 3 (0x18 | 3)
        *(--esp) = entry_point;     // EIP
        *(--esp) = 0x23;
    } else {
        // IRET values for Ring 0
        *(--esp) = 0x23;            // User Data Segment (SS) with RPL 3 (0x20 | 3)
        *(--esp) = kstack_top;      // User Stack Pointer (ESP) - location mapped in user space
        *(--esp) = 0x0202;          // EFLAGS
        *(--esp) = 0x08;            // Kernel Code Segment (CS)
        *(--esp) = entry_point;     // EIP
        *(--esp) = 0x10;
    }

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

    if (user) {
        *(--esp) = 0x23;            // User Data Segment
    } else {
        *(--esp) = 0x10;            // Kernel Data Segment
    }

    process->esp = (uint32_t)esp;
    if (!reuse) {
        total_processes++;
    }
    unlock_scheduler();
    return 1;
}

process_t* current_task() {
    lock_scheduler();
    process_t* rtn = &process_table[current_process];
    unlock_scheduler();
    return rtn;
}

void sleep(uint32_t ms) {
    if (ms == 0) return;
    lock_scheduler();
    process_table[current_process].sleep = ms;
    process_table[current_process].state = STATE_SLEEPING;
    yield();
    unlock_scheduler();
}

void exit() {
    lock_scheduler();
    process_table[current_process].state = STATE_DEAD;
    yield(); // yield forever
    unlock_scheduler();
}