#include "thread/Scheduler.h"
#include "ELF.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include "thread/ThreadImage.h"
#include <luna/Alignment.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.is_kernel = true;
        g_idle.parent_id = 0;
        g_idle.name = "[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, const char* name)
    {
        // If anything fails, make sure to clean up.
        auto guard = make_scope_guard([&] { delete thread; });

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

        guard.deactivate();

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

        thread->stack = thread_stack;

        thread->parent_id = 0;

        thread->name = name;

        thread->is_kernel = true;

        thread->auth = Credentials { .uid = 0, .euid = 0, .suid = 0, .gid = 0, .egid = 0, .sgid = 0 };

        g_threads.append(thread);

        thread->state = ThreadState::Runnable;

        kinfoln("Created kernel thread: id %lu with ip %#lx and sp %#lx", thread->id, thread->ip(), thread->sp());

        return {};
    }

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

        return new_kernel_thread_impl(thread, name);
    }

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

        return new_kernel_thread_impl(thread, name);
    }

    Result<void> new_kernel_thread(void (*func)(void*), void* arg, const char* name)
    {
        Thread* const 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, name);
    }

    Result<Thread*> new_userspace_thread(SharedPtr<VFS::Inode> inode, const char* name)
    {
        Thread* const thread = TRY(make<Thread>());

        thread->state = ThreadState::None;
        thread->is_kernel = false;
        thread->id = 1;
        thread->name = name;
        thread->parent_id = 0;
        thread->auth = Credentials { .uid = 0, .euid = 0, .suid = 0, .gid = 0, .egid = 0, .sgid = 0 };

        Vector<String> args;
        auto name_string = TRY(String::from_cstring(name));
        TRY(args.try_append(move(name_string)));

        Vector<String> env;

        auto guard = make_scope_guard([&] { delete thread; });

        auto image = TRY(ThreadImage::try_load_from_elf(inode));
        u64 argv = TRY(image->push_string_vector_on_stack(args));
        u64 envp = TRY(image->push_string_vector_on_stack(env));

        guard.deactivate();

        image->apply(thread);
        thread->set_arguments(args.size(), argv, env.size(), envp);

        kinfoln("Created userspace thread: id %lu with ip %#.16lx and sp %#.16lx (ksp %#lx)", thread->id, thread->ip(),
                thread->sp(), thread->kernel_stack.top());

        g_threads.append(thread);

        return thread;
    }

    void add_thread(Thread* thread)
    {
        g_threads.append(thread);
    }

    void reap_thread(Thread* thread)
    {
        CPU::disable_interrupts();

        kinfoln("reap: reaping thread with id %zu", thread->id);

        if (thread->is_kernel)
        {
            auto stack = thread->stack;
            MemoryManager::unmap_owned_and_free_vm(stack.bottom(), stack.bytes() / ARCH_PAGE_SIZE).release_value();
        }
        else
        {
            auto stack = thread->kernel_stack;
            MemoryManager::unmap_owned_and_free_vm(stack.bottom(), stack.bytes() / ARCH_PAGE_SIZE).release_value();
        }

        if (!thread->is_kernel) MMU::delete_userspace_page_directory(thread->directory);

        delete thread;

        CPU::enable_interrupts();
    }

    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 (!old_thread->is_kernel) old_thread->fp_data.save();
            if (!new_thread->is_kernel)
            {
                MMU::switch_page_directory(new_thread->directory);
                CPU::switch_kernel_stack(new_thread->kernel_stack.top());
                new_thread->fp_data.restore();
            }
        }

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

        for (auto* const thread : g_threads)
        {
            if (thread->state == ThreadState::Sleeping)
            {
                if (--thread->sleep_ticks_left == 0) thread->state = ThreadState::Runnable;
            }
        }

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

    LinkedList<Thread> check_for_dying_threads()
    {
        LinkedList<Thread> result;

        g_threads.delayed_for_each([&](Thread* thread) {
            if (thread->state == ThreadState::Dying)
            {
                g_threads.remove(thread);
                result.append(thread);
            }
        });

        return result;
    }

    Option<Thread*> find_by_pid(pid_t pid)
    {
        for (auto* const thread : g_threads)
        {
            if (thread->id == (u64)pid && thread->state != ThreadState::Dying) return thread;
        }

        return {};
    }

    bool has_children(Thread* thread)
    {
        bool result { false };

        for_each_child(thread, [&](Thread*) {
            result = true;
            return false;
        });

        return result;
    }

    Option<Thread*> find_exited_child(Thread* thread)
    {
        Option<Thread*> result;

        for_each_child(thread, [&](Thread* child) {
            if (!result.has_value() && child->state == ThreadState::Exited)
            {
                result = child;
                return false;
            }
            return true;
        });

        return result;
    }
}

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

[[noreturn]] void kernel_exit()
{
    g_current->state = ThreadState::Dying;
    kernel_yield();
    unreachable();
}