#include <kstdio.h>
#include <kstring.h>
#include <kstdlib.h>
#include <sched.h>
#include <mm/mm.h>
#include <spinlock.h>
#include <devices/timer.h>





static spinlock_t s;
static task_t *task_list = NULL;
static int next_task_id = 1;









void sched_init(void) {
    spinlock_init(&s);
}

void sched_create_task(void (*__start)()) {
    task_t *new_task = NULL;


    spinlock_acquire(&s);

    new_task = kmalloc(sizeof(task_t));

    if (!new_task) {
        printf("SCHED: Failed to create new task! (out of memory)\n");

        spinlock_release(&s);
        return;
    }

    memset(new_task, 0, sizeof(task_t));

    memset(&new_task->ctx, 0, sizeof(task_context_t));

    new_task->task_id = next_task_id++;
    new_task->next = NULL;

    new_task->stack_size = TASK_STACK_SIZE;
    new_task->stack = kmalloc(new_task->stack_size);

    if (!new_task->stack) {
        printf("SCHED: Failed to allocate new task stack! (out of memory)\n");

        kfree(new_task, sizeof(task_t));
        spinlock_release(&s);
        return;
    }

    memset(new_task->stack, 0, new_task->stack_size);

    uintptr_t sp = (uintptr_t)new_task->stack + new_task->stack_size;
    sp = ALIGN_DOWN(sp, 16);

    new_task->ctx.sp = sp;
    new_task->ctx.elr_el1 = (uint64_t)(uintptr_t)__start;
    new_task->ctx.spsr_el1 = __spsr_read();

    new_task->task_state |= TASK_STATE_READY;

    if (!task_list) {
        task_list = new_task;
        task_list->next = new_task;
    } else {
        task_t *curr = task_list;

        while(curr->next != task_list) {
            curr = curr->next;
        }

        curr->next = new_task;
        new_task->next = task_list;
    }

    printf("SCHED: Created new task ID: %d, SP: 0x%lx, ELR: 0x%lx, SPSR: 0x%lx\n", new_task->task_id, new_task->ctx.sp, new_task->ctx.elr_el1, new_task->ctx.spsr_el1);
    spinlock_release(&s);
}

task_t *sched_get_current_task(void) {
    if (!task_list) {
        return NULL;
    }

    return task_list;
}

void sched_context_switch(task_context_t *old, task_context_t *new) {
    if (!old || !new) {
        return;
    }

    __context_switch(old, new);
}

void sched_context_restore(task_context_t *ctx) {
    __context_restore(ctx);
}

void sched_yield(void) {
    schedule();
}

void sched_start_task(void) {
    if (!task_list) {
        printf("SCHED: No tasks registered!\n");
        return;
    }

    // Set timer
    timer_set(TASK_QUANTUM_MS);

    // Start task
    sched_context_restore(&task_list->ctx);
}

void schedule(void) {
    task_t *prev = task_list;
    task_t *next = task_list->next;




    // Disable interrupts
    __daif_set();

    // Locate next task to be run
    while(!(next->task_state & TASK_STATE_READY) && next != task_list) {
        next = next->next;
    }

    if (next == prev) {
        // No other task to be run (keep running current task)
        timer_set(TASK_QUANTUM_MS);

        // Re-enable interrupts and return
        __daif_clr();
        return;
    }

    // Set previous task to ready state
    prev->task_state &= ~(TASK_STATE_RUNNING);

    // Update next dask
    task_list = next;

    // Set task to running state
    task_list->task_state |= TASK_STATE_RUNNING;

    timer_set(TASK_QUANTUM_MS);
    sched_context_switch(&prev->ctx, &next->ctx);
}