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

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

extern "C" void clock(){
    if(show_ticks)
        print(".");
    if(scheduling)
        schedule();
}
void schedule(){
    // 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
    u64* stack;
    asm("mov %%rbp, %%rax;mov (%%rax), %%rbx; add $16, %%rbx; mov %%rbx,%0": "=m"(stack)::"rax","rbx");

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

    // 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];
    t->id=nproc;
    t->pid=nproc;
    t->size=size;

    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=0;    // Always 0 in long-mode

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

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