#include "thread/Scheduler.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "binfmt/ELF.h"
#include "memory/MemoryManager.h"
#include "thread/ThreadImage.h"
#include "thread/Timer.h"
#include <luna/Alignment.h>
#include <luna/ScopeGuard.h>
#include <luna/Stack.h>

static Thread g_idle;
static Process g_idle_process;
static Thread* g_current = nullptr;
static Process* g_init = nullptr;
static Thread* g_reap = nullptr;
static Thread* g_oom = nullptr;

static const usize TICKS_PER_TIMESLICE = 20;

namespace Scheduler
{
    void init()
    {
        g_idle.tid = 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.process = &g_idle_process;
        g_idle.cmdline = "[idle]";
        g_idle.active_directory = nullptr;

        g_idle_process.id = 0;
        g_idle_process.parent = nullptr;
        g_idle_process.thread_count = 1;
        g_idle_process.is_kernel = true;

        g_idle.ticks_left = 1;

        // Map some stack for the idle task
        u64 idle_stack_vm = mark_critical(MemoryManager::alloc_for_kernel(1, MMU::NoExecute | MMU::ReadWrite),
                                          "Failed to allocate stack memory for the CPU idle thread");

        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 %d with ip %#lx and sp %#lx", g_idle_process.id, g_idle.ip(), g_idle.sp());

        g_current = &g_idle;
    }

    Thread* current()
    {
        return g_current;
    }

    Thread* idle()
    {
        return &g_idle;
    }

    Process* init_process()
    {
        return g_init;
    }

    void set_reap_thread(Thread* thread)
    {
        g_reap = thread;
    }

    void signal_reap_thread()
    {
        if (g_reap) g_reap->wake_up();
    }

    void set_oom_thread(Thread* thread)
    {
        g_oom = thread;
        g_oom->unrestricted_task = true;
    }

    void signal_oom_thread()
    {
        if (g_oom) g_oom->wake_up();
    }

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

        Process* process = TRY(make<Process>());

        auto guard2 = make_scope_guard([&] { delete process; });

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

        guard.deactivate();
        guard2.deactivate();

        const Stack thread_stack { thread_stack_vm, ARCH_PAGE_SIZE * 4 };
        thread->set_sp(thread_stack.top());
        thread->stack = thread_stack;
        thread->cmdline = name;
        thread->is_kernel = true;
        thread->active_directory = MMU::kernel_page_directory();

        thread->process = process;

        process->id = thread->tid;
        process->parent = nullptr;
        process->thread_count = 1;
        process->virtual_clock.set_resolution(1'000'000);
        process->profiling_clock.set_resolution(1'000'000);
        process->cmdline = name;
        process->is_kernel = true;

        g_threads.append(thread);
        g_processes.append(process);

        thread->state = ThreadState::Runnable;

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

        return thread;
    }

    Result<Thread*> 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<Thread*> 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<Thread*> 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*> create_init_process(SharedPtr<VFS::Inode> inode, const char* name)
    {
        check(!g_init);

        Thread* const thread = TRY(make<Thread>());
        Process* const process = TRY(make<Process>());

        thread->state = ThreadState::None;
        thread->tid = 1;
        thread->cmdline = name;
        thread->process = process;

        process->id = 1;
        process->pgid = 1;
        process->thread_count = 1;
        process->cmdline = name;

        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;
            delete process;
        });

        // Contrary to other programs, which use BinaryFormat::create_loader(), init must be a native executable.
        auto loader = TRY(ELFLoader::create(inode, nullptr, 0));

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

        const u64 kernel_stack_base = TRY(MemoryManager::alloc_for_kernel(4, MMU::ReadWrite | MMU::NoExecute));
        Stack kernel_stack { kernel_stack_base, 4 * ARCH_PAGE_SIZE };

        guard.deactivate();

        thread->kernel_stack = kernel_stack;

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

        for (int i = 0; i < NSIG; i++)
        {
            thread->signal_handlers[i] = { .sa_handler = SIG_DFL, .sa_mask = 0, .sa_flags = 0 };
        }

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

        g_threads.append(thread);
        g_processes.append(process);
        g_init = process;

        return thread;
    }

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

    void add_process(Process* process)
    {
        g_processes.append(process);
    }

    void reap_process(Process* process)
    {
        CPU::disable_interrupts();

        // FIXME: Shouldn't all this be done when the timers' destructors are called?
        process->real_timer.disarm();
        process->virtual_timer.disarm();
        process->profiling_timer.disarm();
        for (int i = 0; i < MAX_POSIX_TIMERS; i++)
        {
            if (process->posix_timers[i].has_value()) process->posix_timers[i]->disarm();
        }

        delete process;

        CPU::enable_interrupts();
    }

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

#ifdef REAP_DEBUG
        kdbgln("reap: reaping thread with id %d", thread->tid);
#endif

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

        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 (old_thread->state != ThreadState::Idle && MMU::get_page_directory() != MMU::kernel_page_directory())
                old_thread->active_directory = MMU::get_page_directory();
            if (new_thread->active_directory) MMU::switch_page_directory(new_thread->active_directory);
            if (!new_thread->is_kernel)
            {
                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 if (new_thread->unrestricted_task)
        {
            check(new_thread->is_kernel);
            new_thread->ticks_left = -1;
        }
        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);
        if (!is_in_kernel(regs)) new_thread->process_pending_signals(regs);
    }

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

        if (is_in_kernel(regs))
        {
            g_current->process->kernel_ticks_self++;
            g_current->kernel_ticks_self++;
        }
        else
        {
            g_current->process->virtual_clock.tick();
            g_current->process->user_ticks_self++;
            g_current->user_ticks_self++;
        }

        g_current->process->profiling_clock.tick();

        g_current->ticks_left--;

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

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

    LinkedList<Process> check_for_dead_processes()
    {
        LinkedList<Process> result;

        g_processes.delayed_for_each([&](Process* p) {
            if (p->thread_count == PROCESS_SHOULD_REAP)
            {
                g_processes.remove(p);
                result.append(p);
            }
        });

        return result;
    }

    Option<Process*> find_by_pid(pid_t pid)
    {
        for (auto* const process : g_processes)
        {
            if (process->id == pid) return process;
        }

        return {};
    }

    Option<Thread*> find_by_tid(pid_t tid)
    {
        for (auto* const thread : g_threads)
        {
            if (thread->tid == tid) return thread;
        }

        return {};
    }

    bool has_children(Process* process)
    {
        bool result { false };

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

        return result;
    }

    Option<Process*> find_exited_child(Process* process)
    {
        Option<Process*> result;

        for_each_child(process, [&](Process* child) {
            if (!result.has_value() && child->dead())
            {
                result = child;
                return false;
            }
            return true;
        });

        return result;
    }

    void dump_state()
    {
        CPU::disable_interrupts();

        kdbgln("--- BEGIN SCHEDULER DUMP ---");
        kdbgln("Current thread at %p, tid = %d", g_current, g_current->tid);
        kdbgln("Current process at %p, pid = %d", g_current->process, g_current->process->id);

        for (const auto* thread : g_threads)
        {
            kdbgln("Thread %p (belongs to pid %4d) %c [%-20s] %4d, state = %d, ip = %p", thread, thread->process->id,
                   thread->is_kernel ? 'k' : 'u', thread->cmdline.chars(), thread->tid, (int)thread->state,
                   (void*)thread->ip());
        }

        for (const auto* process : g_processes)
        {
            kdbgln("Process %p (%zu threads) %4d, parent = (%-18p,%d), cwd = %s, ticks: (k:%04zu,u:%04zu), "
                   "status = %d",
                   process, process->thread_count.load(), process->id, process->parent,
                   process->parent ? process->parent->id : 0,
                   process->current_directory_path.is_empty() ? "/" : process->current_directory_path.chars(),
                   process->kernel_ticks_self.load(), process->user_ticks_self.load(), process->status);
        }

        kdbgln("--- END SCHEDULER DUMP ---");

        CPU::enable_interrupts();
    }
}

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

void kernel_wait(pid_t pid)
{
    g_current->child_being_waited_for = pid;
    g_current->state = ThreadState::Waiting;
    kernel_yield();
}

void kernel_wait_for_event()
{
    g_current->state = ThreadState::Waiting;
    kernel_yield();
}

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