#define MODULE "exec"

#include "interrupts/Interrupts.h"
#include "log/Log.h"
#include "memory/Memory.h"
#include "memory/MemoryManager.h"
#include "memory/PMM.h"
#include "memory/VMM.h"
#include "std/ensure.h"
#include "std/errno.h"
#include "std/stdlib.h"
#include "std/string.h"
#include "sys/Syscall.h"
#include "sys/UserMemory.h"
#include "sys/elf/ELFLoader.h"
#include "thread/Scheduler.h"

void sys_fork(Context* context)
{
    kinfoln("fork(): attempting fork");

    Task* parent = Scheduler::current_task();

    Task* child = Scheduler::create_user_task();
    if (!child)
    {
        context->rax = -ENOMEM;
        return;
    }

    if (!child->allocator.inherit(parent->allocator))
    {
        child->state = child->Exited;
        child->exit_status = -127; // so the reaper reaps it on next reaping
        context->rax = -ENOMEM;
        return;
    }

    child->save_context(context);
    child->save_floating();

    for (int i = 0; i < TASK_MAX_FDS; i++) { child->files[i] = parent->files[i]; }

    child->address_space = parent->address_space.clone();

    child->ppid = parent->id;

    child->uid = parent->uid;
    child->euid = parent->euid;
    child->gid = parent->gid;
    child->egid = parent->egid;

    child->regs.rax = 0;
    context->rax = child->id;

    strlcpy(child->name, parent->name, sizeof(child->name));

    child->state = child->Running;

    kinfoln("fork(): forked parent %ld into child %ld", parent->id, child->id);

    return;
}

void push_on_user_stack(uint64_t* rsp, char* value,
                        size_t size) // FIXME: Handle segments of stack that extend beyond one page.
{
    (*rsp) -= size;
    char* kvalue = (char*)VMM::get_physical(*rsp);
    ensure(kvalue != (char*)UINT64_MAX);
    memcpy(kvalue, value, size);
}

void sys_execv(Context* context, const char* pathname, char** argv)
{
    char* kpathname = strdup_from_user(pathname);
    if (!kpathname)
    {
        context->rax = -EFAULT;
        return;
    }

    kinfoln("exec(): executing %s", kpathname);

    VFS::Node* program = VFS::resolve_path(kpathname);
    if (!program)
    {
        kfree(kpathname);
        context->rax = -ENOENT;
        return;
    }

    if (program->type == VFS_DIRECTORY)
    {
        kfree(kpathname);
        context->rax = -EISDIR;
        return;
    }

    if (!VFS::can_execute(program, Scheduler::current_task()->euid, Scheduler::current_task()->egid))
    {
        kfree(kpathname);
        context->rax = -EACCES;
        return;
    }

    long memusage;
    if ((memusage = ELFLoader::check_elf_image(program)) < 0)
    {
        kfree(kpathname);
        context->rax = -ENOEXEC;
        return;
    }

    if ((uint64_t)memusage > PMM::get_free())
    {
        kfree(kpathname);
        context->rax = -ENOMEM;
        return;
    }

    uint64_t kargc = 0;
    char** kargv = nullptr;
    char* arg;

    auto free_kernel_argv_copy = [&]() {
        for (uint64_t i = 0; i < kargc; i++)
        {
            if (kargv[i]) kfree(kargv[i]);
        }
        if (kargv) kfree(kargv);
    };

    // FIXME: This code is a bit messy. Should probably be refactored and moved into a separate function.

    do {
        if (!copy_from_user(argv, &arg, sizeof(char*)))
        {
            free_kernel_argv_copy();
            context->rax = -EFAULT;
            return;
        }
        kargv = (char**)krealloc(kargv, (kargc + 1) * sizeof(char*)); // we need a vector class for the kernel.
        if (!kargv)
        {
            free_kernel_argv_copy();
            context->rax = -ENOMEM;
            return;
        }
        if (arg)
        {
            char* kcopy = strdup_from_user(arg);
            if (!kcopy) // FIXME: This could also be EFAULT.
            {
                free_kernel_argv_copy();
                context->rax = -ENOMEM;
                return;
            }
            kargv[kargc] = kcopy;
        }
        else
        {
            kargv[kargc] = nullptr;
            break;
        }
        kargc++;
        argv++;
    } while (arg != nullptr);

    kinfoln("Copied %lu arguments from user process", kargc);

    size_t stack_size = 0;
    for (uint64_t i = 0; i <= kargc; i++)
    {
        stack_size += sizeof(char*);
        if (kargv[i])
        {
            stack_size += strlen(kargv[i]) + 1; // count the null byte
        }
    }

    if (stack_size >
        ((TASK_PAGES_IN_STACK / 2) *
         PAGE_SIZE)) // FIXME: Maybe we should allocate a larger stack in this case, but still set a larger upper limit.
    {
        free_kernel_argv_copy();
        context->rax = -E2BIG;
        return;
    }

    char** user_argv = (char**)kcalloc(kargc + 1, sizeof(char*));
    if (!user_argv)
    {
        free_kernel_argv_copy();
        context->rax = -ENOMEM;
        return;
    }

    Interrupts::disable();
    ensure(!Interrupts::are_enabled()); // This part is pretty sensitive.

    Task* task = Scheduler::current_task();
    ensure(task);

    // At this point, pretty much nothing can fail.

    task->allocator.free();
    task->allocator
        .init(); // If we had enough space for the old bitmap, we should have enough space for the new bitmap.

    task->address_space.clear();
    task->allocated_stack = (uint64_t)MemoryManager::get_pages_at(
        0x100000, TASK_PAGES_IN_STACK,
        MAP_USER | MAP_READ_WRITE | MAP_AS_OWNED_BY_TASK); // If we had enough space for the old stack, there should be
                                                           // enough space for the new stack.

    ELFImage* image = ELFLoader::load_elf_from_vfs(program);
    ensure(image); // If check_elf_image succeeded, load_elf_from_vfs MUST succeed, unless something has gone terribly
                   // wrong.

    if (VFS::is_setuid(program)) task->euid = program->uid;
    if (VFS::is_setgid(program)) task->egid = program->gid;

    strlcpy(task->name, kpathname, sizeof(task->name));

    Scheduler::reset_task(task, image);

    for (int i = 0; i < TASK_MAX_FDS; i++)
    {
        Descriptor& file = task->files[i];
        if (file.close_on_exec()) { file.close(); }
    }

    for (uint64_t i = 0; i <= kargc; i++)
    {
        if (kargv[i])
        {
            push_on_user_stack(&task->regs.rsp, kargv[i], strlen(kargv[i]) + 1);
            user_argv[i] = (char*)task->regs.rsp;
        }
        else
            user_argv[i] = nullptr;
    }
    push_on_user_stack(&task->regs.rsp, (char*)user_argv, (kargc + 1) * sizeof(char*));
    task->regs.rdi = kargc;          // argc
    task->regs.rsi = task->regs.rsp; // argv

    task->regs.rsp &= (UINT64_MAX ^ 15); // align it

    free_kernel_argv_copy();

    kfree(user_argv);

    kfree(kpathname);

    task->restore_context(context);

    return;
}