#include "thread/Scheduler.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include <luna/ScopeGuard.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 = MemoryManager::alloc_for_kernel(1, MMU::NoExecute | MMU::ReadWrite)
                                .expect_value("Error while setting up the idle task, cannot continue");

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

        g_idle.stack = idle_stack;

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

    Thread* idle()
    {
        return &g_idle;
    }

    Result<void> new_kernel_thread_impl(Thread* thread)
    {
        // If anything fails, make sure to clean up.
        auto guard = make_scope_guard([&] { delete thread; });

        u64 thread_stack_vm = TRY(MemoryManager::alloc_for_kernel(4, MMU::NoExecute | MMU::ReadWrite));

        guard.deactivate();

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

        thread->stack = thread_stack;

        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 new_kernel_thread_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 new_kernel_thread_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 new_kernel_thread_impl(thread);
    }

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

        bool has_found_thread = false;

        do {
            auto maybe_next = g_threads.next(g_current);
            if (!maybe_next.has_value()) g_current = g_threads.expect_first();
            else
                g_current = maybe_next.value();

            if (g_current->state == ThreadState::Runnable)
            {
                has_found_thread = true;
                break;
            }
        } while (g_current != old);

        if (!has_found_thread) g_current = &g_idle;

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

        g_threads.for_each([](Thread* thread) {
            if (thread->state == ThreadState::Sleeping)
            {
                if (--thread->sleep_ticks_left == 0) thread->state = ThreadState::Runnable;
            }
        });

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

void kernel_sleep(u64 ms)
{
    g_current->sleep_ticks_left = ms;
    g_current->state = ThreadState::Sleeping;
    kernel_yield();
}