#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 = NULL;
static process_t* current_process = NULL;
static process_t* idle_task = NULL;

static int total_processes = 0;

void lock_scheduler() {
    disable_interrupts();
}

void unlock_scheduler() {
    enable_interrupts();
}

typedef struct registers {
    uint32_t ds;                                       // push %ds
    uint32_t edi, esi, ebp, esp, ebx, edx, ecx, eax;   // pushal
    uint32_t eip, cs, eflags, useresp, ss;             // pushed by cpu
} __attribute__((packed)) registers_t;

uint32_t switch_context(uint32_t current_esp) {
    registers_t* regs = (registers_t*)current_esp;
    current_process->esp = current_esp;

    process_t* next_process = current_process->sched_next;
    int found_ready = 0;
    while (next_process != NULL) {
        if (next_process != idle_task && next_process->state == STATE_READY) {
            found_ready = 1;
            break;
        }

        if (next_process == current_process) {
            break; 
        }

        next_process = (process_t*)next_process->sched_next;
    }

    if (!found_ready) {
        if (current_process && current_process->state == STATE_READY) {
            next_process = current_process;
        } else {
            next_process = idle_task;
        }
    }

    current_process = next_process;

    uint32_t current_cr3 = fetch_cr3();
    if (current_cr3 != current_process->cr3) {
        load_pd((void*)current_process->cr3);
        flush_tlb();
    }

    if (current_process->flags & PROCESS_FLAG_USER) {
        tss.esp0 = current_process->kstack_top;
    }
    
    return current_process->esp;
}

void init_scheduler() {
    process_t* root = kalloc(sizeof(process_t));
    root->pid = 0;
    root->esp = 0;
    root->cr3 = fetch_cr3();
    root->state = STATE_READY;
    root->sleep = 0;
    root->flags = PROCESS_FLAG_KERNEL;
    root->sched_next = NULL;

    current_process = root;
    idle_task = root;
    process_table = root;

    total_processes = 1;
}

void init_task (
    process_t* process, 
    uint32_t entry_point, 
    uint32_t cr3, 
    int user, 
    uint32_t u_esp
) {
    process->pid = get_next_pid();
    process->cr3 = cr3;
    process->state = STATE_READY;
    process->sleep = 0;
    process->flags = user ? PROCESS_FLAG_USER : PROCESS_FLAG_KERNEL;
    total_processes++;

    if (process->kstack_top != 0) {
        kfree((void*)process->kstack_top);
    }
    
    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) = 0x202;          // EFLAGS (Interrupts enabled)
        *(--esp) = 0x1B;            // User Code Segment (CS) with RPL 3 (0x18 | 3)
        *(--esp) = entry_point;     // EIP
    } else {
        // IRET values for Ring 0
        *(--esp) = 0x202;          // EFLAGS
        *(--esp) = 0x08;            // Kernel Code Segment (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

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

    process->esp = (uint32_t)esp;
}

int get_next_pid() {
    int pid = total_processes;
    total_processes++;
    return pid;
}

void add_task(process_t* task) {
    if (!task) {
        return;
    }

    if (process_table->sched_next == NULL) {
        process_table->sched_next = task;
        task->sched_next = process_table;
        return;
    }

    process_t* curr = process_table;
    while (curr->sched_next != process_table) {
        curr = curr->sched_next;
    }

    curr->sched_next = task;
    task->sched_next = process_table;
}

process_t* current_task() {
    process_t* rtn = current_process;
    return rtn;
}

void sleep(uint32_t ms) {
    if (ms == 0) return;
    current_process->sleep = ms;
    current_process->state = STATE_SLEEPING;
    yield();
}

void exit() {
    current_process->state = STATE_DEAD;
    yield(); // yield forever
}

void wake(process_t* task) {
    if (task->state == STATE_SLEEPING || task->state == STATE_BLOCKED) {
        task->state = STATE_READY;
    }
}

process_t* get_task_by_pid(uint32_t pid) {
    process_t* cur = process_table;
    while (cur) {
        if (cur->pid == pid) break;
        cur = cur->sched_next;
    }
    return cur;
}