#include "thread/Scheduler.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/KernelVM.h"
#include "memory/MemoryManager.h"
#include <luna/Stack.h>

static Thread g_idle;
static Thread* g_current = nullptr;

static const usize TICKS_PER_TIMESLICE = 20;

namespace Scheduler
{
    void init()
    {
        g_idle.id = 0;
        g_idle.init_regs_kernel();
        g_idle.set_ip((u64)CPU::idle_loop);
        g_idle.state = ThreadState::Idle;

        g_idle.ticks_left = 1;

        // Map some stack for the idle task
        u64 idle_stack_vm = KernelVM::alloc_one_page().release_value();
        MemoryManager::alloc_at(idle_stack_vm, 1, MMU::NoExecute | MMU::ReadWrite).release_value();

        Stack idle_stack{idle_stack_vm, ARCH_PAGE_SIZE};
        g_idle.set_sp(idle_stack.top());

        kinfoln("CREATED IDLE THREAD: id %lu with ip %lx and sp %lx", g_idle.id, g_idle.ip(), g_idle.sp());

        g_current = &g_idle;
    }

    Thread* current()
    {
        return g_current;
    }

    Result<void> kernel_thread_alloc_stack_and_append_impl(Thread* thread)
    {
        // FIXME: We will leak the thread if VM allocation or alloc_at fail.
        u64 thread_stack_vm = TRY(KernelVM::alloc_several_pages(4));
        TRY(MemoryManager::alloc_at(thread_stack_vm, 4, MMU::NoExecute | MMU::ReadWrite));

        Stack thread_stack{thread_stack_vm, ARCH_PAGE_SIZE * 4};
        thread->set_sp(thread_stack.top());

        g_threads.append(thread);

        kinfoln("CREATED THREAD: id %lu with ip %lx and sp %lx", thread->id, thread->ip(), thread->sp());

        return {};
    }

    Result<void> new_kernel_thread(u64 address)
    {
        Thread* thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip(address);

        return kernel_thread_alloc_stack_and_append_impl(thread);
    }

    Result<void> new_kernel_thread(void (*func)(void))
    {
        Thread* thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip((u64)func);

        return kernel_thread_alloc_stack_and_append_impl(thread);
    }

    Result<void> new_kernel_thread(void (*func)(void*), void* arg)
    {
        Thread* thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip((u64)func);
        thread->set_arguments((u64)arg, 0, 0, 0);

        return kernel_thread_alloc_stack_and_append_impl(thread);
    }

    Thread* pick_task()
    {
        Thread* old = g_current;
        if (old->is_idle())
        {
            auto maybe_first = g_threads.last();
            if (maybe_first.has_error()) // No threads!!
                return &g_idle;
            g_current = old = maybe_first.value();
        }

        do {
            auto maybe_next = g_threads.next(g_current);
            if (maybe_next.has_error()) g_current = g_threads.first().value();
            else
                g_current = maybe_next.value();

            if (true) // FIXME: Check if the current task is runnable.
                break;
        } while (g_current != old);

        return g_current;
    }

    void generic_switch_context(Thread* old_thread, Thread* new_thread, Registers* regs)
    {
        if (old_thread != new_thread) switch_context(old_thread, new_thread, regs);

        if (new_thread->is_idle())
        {
            new_thread->ticks_left = 1; // The idle task only runs for 1 tick so we can check for new runnable tasks
                                        // as fast as possible.
        }
        else
            new_thread->ticks_left = TICKS_PER_TIMESLICE;
    }

    void switch_task(Registers* regs)
    {
        Thread* old_thread = g_current;
        Thread* new_thread = pick_task();
        generic_switch_context(old_thread, new_thread, regs);
    }

    void invoke(Registers* regs)
    {
        CPU::disable_interrupts();

        g_current->ticks++;

        if (is_in_kernel(regs)) g_current->ticks_in_kernel++;
        else
            g_current->ticks_in_user++;

        g_current->ticks_left--;

        if (!g_current->ticks_left) switch_task(regs);
    }
}