#include "ELF.h"
#include "Log.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 Result<u64> copy_string_vector_to_userspace(const Vector<String>& vec, ThreadImage& image)
{
    Vector<u64> user_vec;
    for (const auto& item : vec)
    {
        // Copy each individual string and retrieve a userspace pointer to said copy
        u64 addr = TRY(image.push_mem_on_stack((const u8*)item.chars(), item.length() + 1));
        TRY(user_vec.try_append(addr));
    }

    TRY(user_vec.try_append((u64) nullptr));
    // Copy the actual vector of userspace pointers to the stack
    return TRY(image.push_mem_on_stack((u8*)user_vec.data(), user_vec.size() * sizeof(u64)));
}

Result<u64> sys_exec(Registers* regs, SyscallArgs args)
{
    auto path = TRY(MemoryManager::strdup_from_user(args[0]));
    auto argv = TRY(copy_string_vector_from_userspace(args[1]));

    // FIXME: Make sure argv is not too big.

    auto inode = TRY(VFS::resolve_path(path.chars()));

    // Not executable
    if ((inode->mode() & S_IXUSR) != S_IXUSR) return err(EACCES);

    kinfoln("exec: attempting to replace current image with %s", path.chars());

    auto current = Scheduler::current();

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

    auto image = TRY(ThreadImage::try_load_from_elf(inode));

    kdbgln("exec: copying argv to image memory (argc = %zu)", argv.size());

    u64 user_argv = TRY(copy_string_vector_to_userspace(argv, *image));
    usize user_argc = argv.size();

    // From now on, nothing should fail.

    kinfoln("exec: image load ok, will now replace existing process image");

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

    MMU::delete_userspace_page_directory(current->directory);

    current->name = path.chars();

    image->apply(current);

    MMU::switch_page_directory(current->directory);

    current->set_arguments(user_argc, user_argv, 0, 0);

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

    kinfoln("exec: done");

    return 0;
}

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

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

    kinfoln("fork: trying to duplicate process %lu", current->id);

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

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

    auto thread = TRY(new_thread());

    thread->state = ThreadState::Runnable;
    thread->is_kernel = false;
    thread->parent_id = current->id;
    thread->fp_data.save();
    thread->name = current->name;

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

    thread->set_return(0);

    Scheduler::add_thread(thread);

    return thread->id;
}