#include "scheduler.h"
#include "api/error.h"
#include "hal/smp.h"
#include "lib/printf.h"
#include "panic/panic.h"
#include "sys/lock.h"
#include "sys/objdir.h"
#include "proc/init.h"
#include "proc/process.h"
#include "proc/ipc.h"
#include <stdatomic.h>

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

atomic_flag s_ready_queue_lock;
static struct pcb_header *s_ready_queue;

atomic_flag s_mail_queue_lock;
static struct pcb_header *s_mail_queue;

static uintmax_t s_this_tick;

static void
s_enqueue_pcb(struct pcb_header *pcb, struct pcb_header **list, atomic_flag *lock)
{
    take_mutex(lock);
    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;
    *list = pcb;
    release_mutex(lock);
}

static void
s_scheduler_swap_pcb(processor_t *proc, pcb_t *current_pcb, pcb_t *target_pcb)
{
    jove_obj_t *target_rootdir = pcb_getrootdir(target_pcb);
    jove_owned_obj_t *target_pcb_obj = pcb_getself(target_rootdir);
    jove_obj_t target_pagedir = objdir_peeklocal(target_rootdir, PCB_ROOTDIR_PAGEDIR_I);

    processor_setpcb(proc, target_pcb_obj);
    
    _scheduler_switch_pcb_arch(target_pcb);
    scheduler_switch(
            pcb_getksp_ref(current_pcb), 
            pcb_getksp(target_pcb), 
            (void*)obj_getvalue((jove_owned_obj_t*)&target_pagedir));
}

int scheduler_make_ready(jove_owned_obj_t *pcb_obj)
{
    pcb_t *pcb = obj_deref(pcb_obj);
    struct pcb_header *pcb_header = (struct pcb_header*)pcb;

    pcb_header->state = PCB_STATE_READY;
    s_enqueue_pcb(pcb_header, &s_ready_queue, &s_ready_queue_lock);

    return 0;
}

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) {
        if(target_pid == 0) {
            //We need to pre-emptively tick the scheduler.
            scheduler_tick();
            return 0;
        }
        return -EBADPROC;
    }
    
    scheduler_savestate();
    s_scheduler_swap_pcb(proc, current_pcb, target_pcb);
    scheduler_restorestate();
    return 0;
}

int scheduler_stop(jove_owned_obj_t *pcb_obj, int stop_code)
{
    struct pcb_header *pcb = obj_deref(pcb_obj);
    struct pcb_header *current_pcb = obj_deref(process_current());

    atomic_flag *state_lock = pcb->state == PCB_STATE_READY ? &s_ready_queue_lock : &s_running_list_lock;

    take_mutex(state_lock);
    if(pcb->sched_prev) pcb->sched_prev->sched_next = pcb->sched_next;
    if(pcb->sched_next) pcb->sched_next->sched_prev = pcb->sched_prev;

    if(pcb->state == PCB_STATE_READY) {
        if(s_ready_queue == pcb) s_ready_queue = pcb->sched_next;
    }else{
        kprintf("SCHEDULER STOP IS NOT READY FOR SMP\n");
        if(s_running_list == pcb) s_running_list = pcb->sched_next;

        if(current_pcb == pcb) {
            release_mutex(state_lock);
            return scheduler_yield(0);
        }
    }
    release_mutex(state_lock);
    return 0;
}

void 
scheduler_block_for_mail(struct pcb_header *pcb, uintmax_t timeout, pid_t yield_target)
{
    ipc_klogf("Blocking PID %i for %i ticks for %i\n", 
            pcb_getpid((pcb_t*)pcb), timeout, yield_target);
    pcb->state = PCB_STATE_BLOCKED_MAIL;

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

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

        //Verify that the root directory of the PCB is valid.
        jove_owned_obj_t *pcb_rootdir_obj = obj_take(pcb_getrootdir((pcb_t*)pcb));
        jove_objdir_t *pcb_rootdir = obj_deref(pcb_rootdir_obj);
        
        jove_owned_obj_t *pcb_selfobj = (jove_owned_obj_t*)&pcb_rootdir->entries[PCB_ROOTDIR_SELF_I];
        if(obj_deref(pcb_selfobj) != pcb) {
            obj_release(pcb_rootdir_obj);
            continue;
        }

        jove_owned_obj_t *pcb_pdirobj = (jove_owned_obj_t*)&pcb_rootdir->entries[PCB_ROOTDIR_PAGEDIR_I];
        if(obj_gettype(pcb_pdirobj) != JOVE_OBJ_PAGEDIR) {
            obj_release(pcb_rootdir_obj);
            continue;
        }

        //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_lock);
        obj_release(pcb_rootdir_obj);
    }
    release_mutex(&s_ready_queue_lock);
}

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

        struct port *pcb_port = pcb_getport((pcb_t*)pcb);
        if(port_poll_send(pcb_port)) {
            //Send buffer is full. Try to transfer to target recv buffer.
            //First check if the target recv buffer has space.
            struct port *recv_port = port_get_send_target(pcb_port);
            if(port_poll_recv(recv_port)) continue;
 
            //Now check if the recv buffer is accepting from this port.
            struct port *recv_target = port_get_recv_target(recv_port);
            if(recv_target != 0 && pcb_port != recv_target) continue;

            //This never hits, but if things change in impl we may have to abort here.
            if(port_recv(recv_port, pcb_port) < 0) continue;
            s_enqueue_pcb(pcb, &s_ready_queue, &s_ready_queue_lock);
        }
        if(port_poll_recv(pcb_port)) {
            //Recv buffer is full, so unblock.
            s_enqueue_pcb(pcb, &s_ready_queue, &s_ready_queue_lock);
        }

        if(s_this_tick >= pcb->block_until) {
            ipc_klogf("Unblocking %i due to timeout\n", pcb_getpid((pcb_t*)pcb));

            s_enqueue_pcb(pcb, &s_ready_queue, &s_ready_queue_lock);
        }
loop_end:
        pcb = next_pcb;
    }
    release_mutex(&s_mail_queue_lock);
}

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

    s_schedule_mail_queue();

    s_schedule_ready_queue();

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

    take_mutex(&s_running_list_lock);
    struct pcb_header *next_pcb = s_running_list;
    if(next_pcb == NULL) {
        kpanic("Kernel ran out of running processes!\n");
    }

    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;
    release_mutex(&s_running_list_lock);

    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)
{
    release_mutex(&s_running_list_lock);
    release_mutex(&s_ready_queue_lock);
    release_mutex(&s_mail_queue_lock);

    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);

    jove_owned_obj_t syscall_stack = init_pop_alloc();
    jove_owned_obj_t interrupt_stack = init_pop_alloc();

    pcb_set_syscall_stack((pcb_t*)pcb_header, obj_getvalue(&syscall_stack));
    pcb_set_interrupt_stack((pcb_t*)pcb_header, obj_getvalue(&interrupt_stack));

    _scheduler_switch_pcb_arch((pcb_t*)pcb_header);

    pcb_map_syscall_stack((pcb_t*)pcb_header);
    pcb_map_interrupt_stack((pcb_t*)pcb_header);

    s_running_list = pcb_header;
    __asm__ volatile("\
            movq %0, %%rsp; \
            xorq %%rbp, %%rbp; \
            jmpq %1":: "r"(pcb_header->ksp), "r"(process_main));
}