#include "memory/RangeAllocator.h"
#include "assert.h"
#include "bootboot.h"
#include "memory/Memory.h"
#include "std/string.h"

extern BOOTBOOT bootboot;

RangeAllocator physical_allocator;

void RangeAllocator::init_from_mmap()
{
    uint64_t total_mem = Memory::get_system();

    void* biggest_chunk = nullptr;
    uint64_t biggest_chunk_size = 0;

    MMapEnt* ptr = &bootboot.mmap;
    uint64_t mmap_entries = (bootboot.size - 128) / 16;
    for (uint64_t i = 0; i < mmap_entries; i++)
    {
        if (MMapEnt_Size(ptr) > biggest_chunk_size)
        {
            biggest_chunk = (void*)MMapEnt_Ptr(ptr);
            biggest_chunk_size = MMapEnt_Size(ptr);
        }
        ptr++;
    }

    bitmap_addr = (char*)biggest_chunk;
    ASSERT((total_mem / 4096 / 8) < biggest_chunk_size);
    bitmap_size = total_mem / 4096 / 8 + 1;
    memset(bitmap_addr, 0, bitmap_size);

    free_mem = total_mem;

    lock_pages(bitmap_addr, bitmap_size / 4096 + 1);

    ptr = &bootboot.mmap;
    for (uint64_t i = 0; i < mmap_entries; i++)
    {
        if (!MMapEnt_IsFree(ptr)) { reserve_pages((void*)MMapEnt_Ptr(ptr), MMapEnt_Size(ptr) / 4096); }
        ptr++;
    }
}

bool RangeAllocator::bitmap_read(uint64_t index)
{
    return (bitmap_addr[index / 8] & (0b10000000 >> (index % 8))) > 0;
}

void RangeAllocator::bitmap_set(uint64_t index, bool value)
{
    uint64_t byteIndex = index / 8;
    uint8_t bitIndexer = 0b10000000 >> (index % 8);
    bitmap_addr[byteIndex] &= ~bitIndexer;
    if (value) { bitmap_addr[byteIndex] |= bitIndexer; }
}

void* RangeAllocator::request_page()
{
    for (uint64_t index = start_index; index < (bitmap_size * 8); index++)
    {
        if (bitmap_read(index)) continue;
        bitmap_set(index, true);
        start_index = index + 1;
        free_mem -= 4096;
        used_mem += 4096;
        return (void*)(index * 4096);
    }

    return 0;
}

void RangeAllocator::free_page(void* address)
{
    uint64_t index = (uint64_t)address / 4096;
    if (!bitmap_read(index)) return;
    bitmap_set(index, false);
    used_mem -= 4096;
    free_mem += 4096;
    if (start_index > index) start_index = index;
}

void RangeAllocator::free_pages(void* address, uint64_t count)
{
    for (uint64_t index = 0; index < count; index++) { free_page((void*)((uint64_t)address + index)); }
}

void RangeAllocator::lock_page(void* address)
{
    uint64_t index = (uint64_t)address / 4096;
    if (bitmap_read(index)) return;
    bitmap_set(index, true);
    used_mem += 4096;
    free_mem -= 4096;
}

void RangeAllocator::reserve_page(void* address)
{
    uint64_t index = (uint64_t)address / 4096;
    if (bitmap_read(index)) return;
    bitmap_set(index, true);
    reserved_mem += 4096;
    free_mem -= 4096;
}

void RangeAllocator::lock_pages(void* address, uint64_t count)
{
    for (uint64_t index = 0; index < count; index++) { lock_page((void*)((uint64_t)address + index)); }
}

void RangeAllocator::reserve_pages(void* address, uint64_t count)
{
    for (uint64_t index = 0; index < count; index++) { reserve_page((void*)((uint64_t)address + index)); }
}

uint64_t RangeAllocator::get_free()
{
    return free_mem;
}

uint64_t RangeAllocator::get_used()
{
    return used_mem;
}

uint64_t RangeAllocator::get_reserved()
{
    return reserved_mem;
}