#include "Log.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include "sys/Syscall.h"
#include "thread/Scheduler.h"
#include <bits/mmap-flags.h>
#include <luna/Alignment.h>

constexpr uintptr_t USERSPACE_HEAP_BASE = 0x3000000;

Result<u64> sys_allocate_memory(Registers*, SyscallArgs args)
{
    usize size = (usize)args[0];
    int flags = (int)args[1];

    if (size != ARCH_PAGE_SIZE) return err(EINVAL);
    if (flags < 0) return err(EINVAL);

    if (size == 0) return 0;

    Thread* current = Scheduler::current();
    if (!current->heap_bitmap.initialized())
    {
        void* bitmap_location = (void*)TRY(MemoryManager::alloc_for_kernel(1, MMU::ReadWrite));
        current->heap_bitmap.initialize(bitmap_location, ARCH_PAGE_SIZE);
        current->heap_bitmap.clear(false);
    }

    u64 index = TRY(Result<u64>::from_option(current->heap_bitmap.find_and_toggle(false), ENOMEM));

    u64 address = USERSPACE_HEAP_BASE + (index * ARCH_PAGE_SIZE);

    int mmu_flags = MMU::User | MMU::NoExecute;
    if (flags & PROT_WRITE) mmu_flags |= MMU::ReadWrite;
    if (flags & PROT_EXEC) mmu_flags &= ~MMU::NoExecute;

    kdbgln("allocate_memory: allocating memory at %#lx", address);

    return MemoryManager::alloc_at(address, 1, mmu_flags);
}

Result<u64> sys_deallocate_memory(Registers*, SyscallArgs args)
{
    u64 address = (u64)args[0];

    Thread* current = Scheduler::current();
    if (!current->heap_bitmap.initialized()) return err(EFAULT);

    u64 index = (address - USERSPACE_HEAP_BASE) / ARCH_PAGE_SIZE;

    if (!current->heap_bitmap.get(index)) return err(EFAULT);

    current->heap_bitmap.set(index, false);

    kdbgln("deallocate_memory: deallocating memory at %#lx", address);

    TRY(MemoryManager::unmap_owned(address, 1));

    return { 0 };
}