#include "scheduler.hpp"
#include "boucane.hpp"
#include "core/apic.hpp"
#include "libs/string.hpp"

PROC procs[MAX_TASK];
u32 nproc=0;
char show_ticks=0;
char scheduling=0;
u32 active_process=0;

extern "C" void clock(){
    u64* stack;
    asm("mov %%rbp, %0": "=r"(stack)::"rax");
    stack=&stack[2];

    if(show_ticks)
        print(".");
    if(scheduling)
        schedule(stack);
}

void tdump(PROC*t){
    printk("ss:%x rsp:%x eflags:%x cs:%x rip:%x\n",
        t->registers.ss,
        t->registers.rsp,
        t->registers.eflags,
        t->registers.cs,
        t->registers.rip
    );
}

void schedule(u64* stack){
    // First get a pointer to the first process saved register.
    // Since this is called by clock(), %rbp contains a pointer
    // to the clock() %rbp value and then we access to the registers SAVE_REGS in int.S

    // Save current task
    PROC *t=&procs[active_process];
    t->registers.ds=stack[0];
    t->registers.fs=stack[1];
    t->registers.es=stack[2];
    t->registers.gs=stack[3];
    t->registers.rdi=stack[4];
    t->registers.rsi=stack[5];
    t->registers.rbp=stack[6];
    t->registers.rdx=stack[7];
    t->registers.rcx=stack[8];
    t->registers.rbx=stack[9];
    t->registers.rax=stack[10];
    t->registers.r15=stack[11];
    t->registers.r14=stack[12];
    t->registers.r13=stack[13];
    t->registers.r12=stack[14];
    t->registers.r11=stack[15];
    t->registers.r10=stack[16];
    t->registers.r9=stack[17];
    t->registers.r8=stack[18];
    t->registers.rip=stack[19];
    t->registers.cs=stack[20];
    t->registers.eflags=stack[21];
    t->registers.rsp=stack[22];
    t->registers.ss=stack[23];

    // Goto next task
    active_process++;
    if(active_process>=nproc)
        active_process=0;

    t=&procs[active_process];
    kvar_tss.rsp0=t->registers.rsp0;


    // Clock acknownledgement
    apic_ack();

    asm volatile(
        "mov %0, %%rdi       \n\t"  
        "jmp switch          \n\t"
        :: "a" (t)
    );
}


void create_task(void* task, u32 size){
    if(nproc>=MAX_TASK){
        printk("Could not create more tasks.");
        return;
    }
    PROC *t=&procs[nproc];
    memset(t, 0, sizeof(PROC));
    t->id=nproc;
    t->pid=nproc;
    t->size=size;
    t->registers.eflags=0x246;

    u32 npages=size%4096 ? size/4096 + 1 : size/4096;
    // Note that paging_create_task() allocate 2 more pages (one for the user stack and
    // the other for the kernel stack)
    t->pml4=paging_create_task(npages);
    t->registers.rsp=TASK_VMA+npages*4096+4096;   // User stack
    t->registers.rsp0=TASK_VMA+npages*4096+4096*2;  // Kernel stack on the last page

    t->registers.rip=TASK_VMA;
    t->registers.cs=0x1B; // 0x18 and 0x3 privilege
    t->registers.ds=0x23; // 0x20 and 0x3 privilege
    t->registers.ss=0x23;

    // Load task using
    lpml4(t->pml4);
    memcpy(task, TASK_VMA, size);
    lpml4(kpml4);
    
    nproc++;
}

void scheduler_start(){
    cli();
    scheduling=1;
    active_process=0;
    PROC *t=&procs[active_process];
    kvar_tss.rsp0=t->registers.rsp0;
    asm(
        "mov %0, %%rdi   \n\t"    
        "jmp switch      \n\t"
        :: "a" (t)
    );
}