#include "memory/UserVM.h"
#include "Log.h"
#include "arch/MMU.h"
#include "memory/Heap.h"
#include <luna/CString.h>
#include <luna/ScopeGuard.h>

static constexpr u64 VM_BASE = 0x10000000;

static constexpr usize INITIAL_VM_SIZE = 80;
static constexpr usize MAX_VM_SIZE = 1024 * 1024 * 16;

Result<OwnedPtr<UserVM>> UserVM::try_create()
{
    void* const base = TRY(kmalloc(INITIAL_VM_SIZE));

    auto guard = make_scope_guard([&] { kfree(base); });

    OwnedPtr<UserVM> ptr = TRY(make_owned<UserVM>(base, INITIAL_VM_SIZE));

    guard.deactivate();

    return move(ptr);
}

Result<OwnedPtr<UserVM>> UserVM::clone()
{
    void* const base = TRY(kmalloc(m_bitmap.size_in_bytes()));

    auto guard = make_scope_guard([&] { kfree(base); });

    OwnedPtr<UserVM> ptr = TRY(make_owned<UserVM>(base, m_bitmap.size_in_bytes()));

    memcpy(ptr->m_bitmap.location(), m_bitmap.location(), m_bitmap.size_in_bytes());

    guard.deactivate();

    return move(ptr);
}

UserVM::UserVM(void* base, usize size)
{
    m_bitmap.initialize(base, size);
    m_bitmap.clear(false);
}

Result<bool> UserVM::try_expand(usize size)
{
    if (m_bitmap.size_in_bytes() == MAX_VM_SIZE) { return false; }

    const usize old_size = m_bitmap.size_in_bytes();
    usize new_size = old_size + size;

    if (new_size > MAX_VM_SIZE) new_size = MAX_VM_SIZE;

    m_bitmap.resize(new_size);
    m_bitmap.clear_region(old_size * 8, (new_size - old_size) * 8, false);

    return true;
}

Result<u64> UserVM::alloc_one_page()
{
    u64 index;
    bool ok = m_bitmap.find_and_toggle(false).try_set_value(index);
    if (!ok)
    {
        bool success = TRY(try_expand());
        if (!success) return err(ENOMEM);
        index = TRY(Result<u64>::from_option(m_bitmap.find_and_toggle(false), ENOMEM));
    }

    return VM_BASE + index * ARCH_PAGE_SIZE;
}

Result<u64> UserVM::alloc_several_pages(usize count)
{
    u64 index;
    bool ok = m_bitmap.find_and_toggle_region(false, count).try_set_value(index);
    if (!ok)
    {
        bool success = TRY(try_expand((count / 8) + INITIAL_VM_SIZE));
        if (!success) return err(ENOMEM);
        index = TRY(Result<u64>::from_option(m_bitmap.find_and_toggle_region(false, count), ENOMEM));
    }

    return VM_BASE + index * ARCH_PAGE_SIZE;
}

Result<bool> UserVM::free_one_page(u64 address)
{
    if (address < VM_BASE) return err(EINVAL);
    const u64 index = (address - VM_BASE) / ARCH_PAGE_SIZE;
    if (index > (MAX_VM_SIZE * 8)) return err(EINVAL);

    // NOTE: POSIX says munmap() should silently do nothing if the address is not mapped, instead of throwing an error
    // like EFAULT.
    if (!m_bitmap.get(index)) return false;

    m_bitmap.set(index, false);

    return true;
}

Result<bool> UserVM::free_several_pages(u64 address, usize count)
{
    if (address < VM_BASE) return err(EINVAL);
    const u64 index = (address - VM_BASE) / ARCH_PAGE_SIZE;
    if ((index + count) > (MAX_VM_SIZE * 8)) return err(EINVAL);

    // NOTE: Same as above.
    if (!TRY(m_bitmap.try_match_region(index, count, true))) return false;

    m_bitmap.clear_region(index, count, false);

    return true;
}

UserVM::~UserVM()
{
    m_bitmap.deallocate();
}