#include <assert.h>
#include <kprint.h>
#include <stdint.h>
#include <neobbo.h>
#include <arch/amd64/hal/smp.h>
#include <error.h>
#include <string.h>
#include <mm/kmalloc.h>
#include <mm/slab.h>
#include <scheduler/sched.h>

extern void switch_context(struct context **old, struct context *new);

struct ma_cache *thread_cache;
struct thread *idle;

#define QUANTUM_US 10000

int next_pid = 1;

void idle_task(){
    for(;;){
        kprintf("idling!\n");
    }
}

void test_task(){
    kprintf("Hello world from scheduled task!\n");
    //for(;;);
}

void best_task(){
    kprintf("Hello world I am best\n");
    //for(;;);
}

void thread_exit(){
    cpu_state *cpu = get_cpu_struct();
    struct thread *p = cpu->current_process;

    if(p == p->next){
        // If this is the only thread in the queue then set cpu->head to NULL, so scheduler knows to idle
        cpu->head = NULL;
    }else{
        // Remove process from circular linked list
        p->next->prev = p->prev;
        p->prev->next = p->next;

        if(p == cpu->head){
            //cpu->head = p->next;
        }

    }

    p->state = ZOMBIE;

    // Switch to scheduler
    yield();
}

/* Setup a process structure */
struct thread *alloc_thread(void){
    struct thread *t = ma_cache_alloc(thread_cache, 0);
    memset(t, 0, sizeof(struct thread));

    t->state = READY;

    t->kstack = kzalloc(8 * 4096);

    t->pid = next_pid++;

    uint64_t *sp = (uint64_t*)((uint64_t)t->kstack + 8 * 4096);

    // Push the exit function (the thread will return to it when it finishes)
    *--sp = (uint64_t)thread_exit;
    
    // Allocate space for context 
    sp -= sizeof(struct context)/sizeof(uint64_t);

    t->context = (struct context*)sp;
    t->context->rbp = (uint64_t)sp;

    return t;

}

struct thread *add_thread(uint64_t *entry){
    struct thread *t = alloc_thread();

    assert(t != NULL && "Thread allocation failed!");

    struct cpu_state *cpu = get_cpu_struct();
    struct thread *head = get_cpu_struct()->head;
    struct thread *base = get_cpu_struct()->base;    

    // Manage circley linked list
    if(base == NULL){
        t->next = t;
        t->prev = t;
        cpu->base = t;
        cpu->head = t;
    } else {
        t->prev = head;
        t->next = base;
        head->next = t;
        base->prev = t;
        cpu->head = t;
    }

    t->context->rip = (uint64_t)entry;

    return t;
}


[[noreturn]] void sched(){
    for(;;){
        asm("cli"); // we sti at the end of switch_context
        cpu_state *cpu = get_cpu_struct();
        struct thread *prev = cpu->current_process;

        if(prev->state == ZOMBIE){
            kprintf("we are freeing allocated thread 0x{x}\n", prev);
            ma_cache_dealloc(prev);
        }else{
            prev->state = READY;
        }

        if(cpu->head == NULL){
            cpu->current_process = idle;
        }else{
            if(prev->next){
                cpu->current_process = prev->next;
            }else{
                cpu->current_process = cpu->head;
            }
        }
        
        cpu->current_process->state = RUNNING;
        switch_context(&(get_cpu_struct()->scheduler_context), get_cpu_struct()->current_process->context);
    }
}

void scheduler_init(){
    assert(get_cpu_struct_initialized() && "CPU struct not initialized!");

    cpu_state *cpu = get_cpu_struct();

    if(cpu->current_process != NULL){
        kprintf("scheduler on CPU {d} already initialized!\n", cpu->lapic_id);
        return;
    }

    idle = alloc_thread();
    idle->context->rip = (uint64_t)idle;

    assert(idle != NULL && "Failed to allocate idle task!");

    cpu->current_process = idle;

    cpu->scheduler_stack = kzalloc(4096);
    cpu->scheduler_context = (struct context*)((uint64_t)cpu->scheduler_stack + 4096);
    cpu->scheduler_context->rbp = (uint64_t)cpu->scheduler_context;
    cpu->scheduler_context->rip = (uint64_t)sched;
    
    add_thread((uint64_t*)test_task);

    add_thread((uint64_t *)best_task);

    // Initialize scheduler -> we will now get timer interrupts to switch us into sched()
    cpu->scheduler_initialized = true;
}


void yield(){
    switch_context(&get_cpu_struct()->current_process->context, get_cpu_struct()->scheduler_context);
}