#include "scheduler.h"
#include "api/error.h"
#include "hal/smp.h"
#include "sys/objdir.h"
#include "proc/init.h"
#include "proc/process.h"

//TODO: Better scheduling algorithm than round robin.
static struct pcb_header *s_running_list;

static struct pcb_header *s_ready_queue;
static struct pcb_header *s_mail_queue;

static uintmax_t s_this_tick;

static inline void
s_scheduler_swap_pcb(processor_t *proc, pcb_t *current_pcb, pcb_t *target_pcb)
{
    jove_owned_obj_t *target_rootdir = obj_take(pcb_getrootdir(target_pcb));
    jove_owned_obj_t *target_pcb_obj = pcb_getself(target_rootdir);
    obj_release(target_rootdir);

    processor_setpcb(proc, target_pcb_obj);
    
    _scheduler_switch_pcb_arch(target_pcb);
    scheduler_switch(pcb_getksp_ref(current_pcb), pcb_getksp(target_pcb));
}

int scheduler_yield(pid_t target_pid)
{
    processor_t *proc = processor_current();
    pcb_t *current_pcb = obj_deref(processor_getpcb(proc));
    pcb_t *target_pcb = (pcb_t*)((struct pcb_header*)current_pcb)->sched_next;

    if(target_pid != 0) {
        pcb_t *target_pcb = pcb_frompid(target_pid);
    }
    if(target_pcb == 0) {
        return -EBADPROC;
    }
    
    scheduler_savestate();
    s_scheduler_swap_pcb(proc, current_pcb, target_pcb);
    scheduler_restorestate();
    return 0;
}

void
s_enqueue_pcb(struct pcb_header *pcb, struct pcb_header *list)
{
    if(pcb->sched_prev) pcb->sched_prev->sched_next = pcb->sched_next;
    if(pcb->sched_prev) pcb->sched_next->sched_prev = pcb->sched_prev;

    list->sched_prev = pcb;
    pcb->sched_next = list;
    pcb->sched_prev = NULL;
}

void 
scheduler_block_for_mail(struct pcb_header *pcb, uintmax_t timeout, pid_t yield_target)
{
    pcb->state = PCB_STATE_BLOCKED_MAIL;

    s_enqueue_pcb(pcb, s_mail_queue);
    s_mail_queue = pcb;
    
    pcb->block_until = s_this_tick + timeout;
    scheduler_yield(yield_target);
}

void
s_schedule_ready_queue(void)
{
    for(struct pcb_header *pcb = s_ready_queue; pcb; pcb = s_ready_queue) {
        s_ready_queue = pcb->sched_next;

        //Add this process to the ready list after the last ready process.
        pcb->state = PCB_STATE_READY;
        s_enqueue_pcb(pcb, s_running_list);
        s_running_list = pcb;
    }
}

void
s_schedule_mail_queue(void)
{
    struct pcb_header *pcb = s_mail_queue;
    while(pcb) {
        struct pcb_header *next_pcb = pcb->sched_next;

        if(pcb->send.obj) {
            struct pcb_header *send_target = (struct pcb_header*)pcb_frompid(pcb->send.pid);
            if(send_target->recv.pid == 0 || send_target->recv.pid == pcb->pid) {
                objdir_swapobj(send_target->recv.obj, pcb->send.obj);
                send_target->recv.pid = pcb->pid;
            
                s_enqueue_pcb(pcb, s_running_list);
                s_running_list = pcb;
                goto loop_end;
            }
        }
        if(pcb->recv.obj) {
            s_enqueue_pcb(pcb, s_running_list);
            s_running_list = pcb;
            goto loop_end;
        }

        if(s_this_tick >= pcb->block_until) {
            s_enqueue_pcb(pcb, s_running_list);
            s_running_list = pcb;
        }
loop_end:
        pcb = next_pcb;
    }
}

void
scheduler_tick()
{
    size_t last_tick = s_this_tick++;

    s_schedule_ready_queue();
    s_schedule_mail_queue();

    processor_t *proc = processor_current();
    struct pcb_header *current_pcb = obj_deref(processor_getpcb(proc));
    struct pcb_header *next_pcb = s_running_list;
    struct pcb_header *tail = next_pcb;

    while(tail->sched_next) {
        tail = tail->sched_next;
        if(tail == next_pcb) break;
    }

    tail->sched_next = next_pcb;
    next_pcb->sched_prev = tail;
    s_running_list = next_pcb->sched_next;

    s_scheduler_swap_pcb(proc, (pcb_t*)current_pcb, (pcb_t*)next_pcb);
}

void scheduler_init(jove_owned_obj_t *pcb_obj, void *process_main)
{
    struct pcb_header *pcb_header = obj_deref(pcb_obj);
    pcb_header->state = PCB_STATE_RUNNING;

    processor_t *proc = processor_current();
    processor_setpcb(proc, pcb_obj);

    s_running_list = pcb_header;

    __asm__("movq %0, %%rsp; \
             xorq %%rbp, %%rbp; \
             jmp %1"::
             "m"(pcb_header->ksp), "r"(process_main));
}