#include "Log.h"
#include "Pledge.h"
#include "binfmt/BinaryFormat.h"
#include "fs/VFS.h"
#include "memory/MemoryManager.h"
#include "sys/Syscall.h"
#include "thread/Scheduler.h"
#include "thread/ThreadImage.h"
#include <bits/modes.h>
#include <bits/open-flags.h>
#include <luna/CString.h>
#include <luna/ScopeGuard.h>
#include <luna/Vector.h>

static Result<Vector<String>> copy_string_vector_from_userspace(u64 address)
{
    Vector<String> result;

    const u64* user_vector = (const u64*)address;

    u64 string_addr;
    while (true)
    {
        if (!MemoryManager::copy_from_user_typed(user_vector, &string_addr)) return err(EFAULT);
        if (!string_addr) break;

        auto string = TRY(MemoryManager::strdup_from_user(string_addr));
        TRY(result.try_append(move(string)));
        user_vector++;
    }

    return result;
}

static u64 calculate_userspace_stack_size(const Vector<String>& v)
{
    u64 total { 0 };

    for (const auto& str : v)
    {
        // The string's byte count + a terminating NUL byte.
        total += str.length() + 1;
        // The pointer to said string in the userspace array.
        total += sizeof(char*);
    }

    // The NULL pointer at the end of the userspace array.
    total += sizeof(char*);

    return total;
}

static constexpr usize MAX_ARGV_STACK_SIZE = 2 * ARCH_PAGE_SIZE;

Result<u64> sys_execve(Registers* regs, SyscallArgs args)
{
    auto path = TRY(MemoryManager::strdup_from_user(args[0]));
    auto argv = TRY(copy_string_vector_from_userspace(args[1]));
    Vector<String> envp;
    if (args[2]) envp = TRY(copy_string_vector_from_userspace(args[2]));

    if ((calculate_userspace_stack_size(argv) + calculate_userspace_stack_size(envp)) > MAX_ARGV_STACK_SIZE)
        return err(E2BIG);

    auto current = Scheduler::current();

    TRY(check_pledge(current, Promise::p_exec));

    auto inode = TRY(VFS::resolve_path(path.chars(), current->auth, current->current_directory));

    if (!VFS::can_execute(inode, current->auth)) return err(EACCES);

#ifdef EXEC_DEBUG
    kdbgln("exec: attempting to replace current image with %s", path.chars());
#endif

    bool is_setuid = VFS::is_setuid(inode);
    bool is_setgid = VFS::is_setgid(inode);
    bool is_secure_environment = is_setgid || is_setuid;
    if (is_secure_environment && current->execpromises >= 0) return err(EACCES);

    auto loader = TRY(BinaryFormat::create_loader(inode));

#ifdef EXEC_DEBUG
    kdbgln("exec: created loader for binary format %s", loader->format().chars());
#endif

    auto guard = make_scope_guard([current] { MMU::switch_page_directory(current->self_directory()); });

    auto image = TRY(ThreadImage::try_load_from_binary(loader));

    argv = TRY(loader->cmdline(path, move(argv)));

    u64 user_argv = TRY(image->push_string_vector_on_stack(argv));
    usize user_argc = argv.size();

    u64 user_envp = TRY(image->push_string_vector_on_stack(envp));
    usize user_envc = envp.size();

    // From now on, nothing should fail.

#ifdef EXEC_DEBUG
    kdbgln("exec: image load ok, will now replace existing process image");
#endif

    guard.deactivate();

    for (int i = 0; i < FD_MAX; i++)
    {
        auto& descriptor = current->fd_table[i];
        if (!descriptor.has_value()) continue;
        if (descriptor->flags & O_CLOEXEC) { descriptor = {}; }
    }

    if (is_setuid) current->auth.euid = current->auth.suid = inode->metadata().uid;
    if (is_setgid) current->auth.egid = current->auth.sgid = inode->metadata().gid;

    current->name = path.chars();

    image->apply(current);

    MMU::switch_page_directory(current->self_directory());

    current->set_arguments(user_argc, user_argv, user_envc, user_envp);

    current->promises = current->execpromises;
    current->execpromises = -1;

    memcpy(regs, &current->regs, sizeof(*regs));

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

    current->has_called_exec = true;

    kinfoln("exec: thread %d was replaced with %s", current->id, path.chars());

    return 0;
}

Result<u64> sys_fork(Registers* regs, SyscallArgs)
{
    auto current = Scheduler::current();

    TRY(check_pledge(current, Promise::p_proc));

    auto guard = make_scope_guard([current] { MMU::switch_page_directory(current->self_directory()); });

    memcpy(&current->regs, regs, sizeof(*regs));

    auto current_directory_path = TRY(current->current_directory_path.clone());

    auto image = TRY(ThreadImage::clone_from_thread(current));

    auto thread = TRY(new_thread());

    thread->state = ThreadState::Runnable;
    thread->is_kernel = false;
    thread->fp_data.save();
    thread->name = current->name;
    thread->auth = current->auth;
    thread->current_directory = current->current_directory;
    thread->current_directory_path = move(current_directory_path);
    thread->umask = current->umask;
    thread->parent = current;
    thread->promises = current->promises;
    thread->execpromises = current->execpromises;

    for (int i = 0; i < FD_MAX; i++) { thread->fd_table[i] = current->fd_table[i]; }

    image->apply(thread);

    memcpy(&thread->regs, regs, sizeof(*regs));

    for (int i = 0; i < NSIG; i++) thread->signal_handlers[i] = current->signal_handlers[i];
    thread->signal_mask = current->signal_mask;

    thread->set_return(0);

    Scheduler::add_thread(thread);

#ifdef FORK_DEBUG
    kdbgln("fork: thread %d forked into child %d", current->id, thread->id);
#endif

    return thread->id;
}