#define MODULE "sched"

#include "thread/Scheduler.h"
#include "assert.h"
#include "interrupts/Interrupts.h"
#include "log/Log.h"
#include "memory/KernelMemoryManager.h"
#include "misc/hang.h"
#include "std/string.h"
#include "thread/PIT.h"
#include "thread/Task.h"

static Task tasks[32];
static uint64_t task_num = 0;

static Task idle_task;

static uint64_t free_tid = 0;

static Task* sched_current_task;
static Task* base_task;
static Task* end_task;

static void idle_task_function()
{
    Interrupts::enable();
    while (1) halt();
}

static uint64_t frequency;

void Scheduler::init()
{
    memset(&idle_task, 0, sizeof(Task));
    idle_task.id = free_tid++;
    idle_task.regs.rip = (uint64_t)&idle_task_function;
    idle_task.regs.rsp = (uint64_t)KernelMemoryManager::get_page();
    idle_task.regs.cs = 0x08;
    idle_task.regs.ds = 0x10;
    asm volatile("pushfq; movq (%%rsp), %%rax; movq %%rax, %0; popfq;" : "=m"(idle_task.regs.rflags)::"%rax");
    idle_task.regs.rflags |= 0x200;
    idle_task.task_sleep = 1000;

    base_task = &tasks[task_num++];
    memset(base_task, 0, sizeof(Task));
    end_task = base_task;
    sched_current_task = base_task;
    sched_current_task->id = free_tid++;
    sched_current_task->task_time = 20; // gets 20 ms of cpu time before next switch
    sched_current_task->next_task = nullptr;
    // the other registers will be saved next task switch

    frequency = 1000 / PIT::frequency();
}

void Scheduler::add_kernel_task(void (*task)(void))
{
    if (task_num == 32) return; // FIXME: allow for dynamically allocated linked list instead of a fixed array of Tasks
    Task* new_task = &tasks[task_num++];
    memset(new_task, 0, sizeof(Task));
    new_task->id = free_tid++;
    new_task->regs.rip = (uint64_t)task;
    new_task->regs.rsp = (uint64_t)KernelMemoryManager::get_pages(2); // 8 KB is enough for everyone, right?
    new_task->regs.cs = 0x08;
    new_task->regs.ds = 0x10;
    asm volatile("pushfq; movq (%%rsp), %%rax; movq %%rax, %0; popfq;" : "=m"(new_task->regs.rflags)::"%rax");
    new_task->regs.rflags |= 0x200; // enable interrupts
    new_task->task_sleep = 0;
    new_task->task_time = 0;
    end_task->next_task = new_task;
    end_task = new_task;
    kinfoln("Adding task: starts at %lx, tid %ld, stack at %lx, total tasks: %ld", new_task->regs.rip, new_task->id,
            new_task->regs.rsp, task_num);
}

static void sched_decrement_sleep_times()
{
    Task* task = base_task;
    while (task)
    {
        if (task->task_sleep > 0)
        {
            task->task_sleep -= frequency;
            if (task->task_sleep < 0) task->task_sleep = 0;
        }
        task = task->next_task;
    }
}

void Scheduler::task_tick(Context* context)
{
    ASSERT(Interrupts::is_in_handler());
    sched_decrement_sleep_times();
    sched_current_task->task_time -= frequency;
    if (sched_current_task->task_time < 0)
    {
        sched_current_task->task_time = 0;
        task_yield(context);
    }
}

void Scheduler::task_yield(Context* context)
{
    ASSERT(Interrupts::is_in_handler());
    get_context_to_task(*sched_current_task, context);
    bool was_idle = false;
    if (sched_current_task->id == 0) // idle task
    {
        sched_current_task = base_task;
        was_idle = true;
    }
    Task* original_task = sched_current_task;
    do {
        sched_current_task = sched_current_task->next_task;
        if (!sched_current_task) { sched_current_task = base_task; }
        if (sched_current_task->task_sleep == 0)
        {
            sched_current_task->task_time = 20;
            set_context_from_task(*sched_current_task, context);
            return;
        }
    } while (sched_current_task != original_task);
    if (original_task->task_sleep > 0)
    {
        sched_current_task = &idle_task;
        sched_current_task->task_time = frequency;
        if (!was_idle) { set_context_from_task(*sched_current_task, context); }
        return;
    }
    original_task->task_time = 20; // grant 30 more ms, there is no other task available
    return;
}

void Scheduler::yield()
{
    asm volatile("int $48");
}

void Scheduler::sleep(unsigned long ms)
{
    current_task()->task_sleep = ms;
    yield();
}

Task* Scheduler::current_task()
{
    return sched_current_task;
}