#define MODULE "kheap"

#include "memory/KernelHeap.h"
#include "assert.h"

#ifndef PAGE_SIZE
#define PAGE_SIZE 4096
#endif

static uint8_t page_bitmap[2048];

#define ALLOC_BASE 0xfffffffff8000000
#define ALLOC_END 0xfffffffffc000000

static uint64_t start_index = 0;

static bool bitmap_read(uint64_t index)
{
    return (page_bitmap[index / 8] & (0b10000000 >> (index % 8))) > 0;
}

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

uint64_t KernelHeap::request_virtual_page()
{
    for (uint64_t index = start_index; index < sizeof(page_bitmap) * 8; index++)
    {
        if (bitmap_read(index)) continue;
        bitmap_set(index, true);
        start_index = index + 1;
        return ALLOC_BASE + (index * PAGE_SIZE);
    }

    return 0;
}

uint64_t KernelHeap::request_virtual_pages(uint64_t count)
{
    uint64_t contiguous = 0;
    uint64_t contiguous_start = 0;
    for (uint64_t index = start_index; index < sizeof(page_bitmap) * 8; index++)
    {
        if (bitmap_read(index))
        {
            contiguous = 0;
            continue;
        }
        if (contiguous == 0)
        {
            contiguous_start = index;
            contiguous++;
        }
        else
            contiguous++;
        if (contiguous == count)
        {
            for (uint64_t i = 0; i < count; i++) bitmap_set(contiguous_start + i, true);
            return ALLOC_BASE + (contiguous_start * PAGE_SIZE);
        }
    }

    return 0;
}

void KernelHeap::free_virtual_page(uint64_t address)
{
    if (address < ALLOC_BASE || address >= ALLOC_END) return;
    uint64_t index = (address - ALLOC_BASE) / PAGE_SIZE;
    bitmap_set(index, false);
    if (start_index > index) start_index = index;
}

void KernelHeap::free_virtual_pages(uint64_t address, uint64_t count)
{
    if (address < ALLOC_BASE || address >= ALLOC_END) return;
    uint64_t index = (address - ALLOC_BASE) / PAGE_SIZE;
    for (uint64_t i = 0; i < count; i++) { bitmap_set(index + i, false); }
    if (start_index > index) start_index = index;
}