#define MODULE "mem"

#include "memory/UserHeap.h"
#include "log/Log.h"
#include "misc/utils.h"
#include "std/stdlib.h"
#include "std/string.h"

#ifndef USER_HEAP_DEBUG
#undef kdbgln
#define kdbgln(...)
#endif

#ifndef PAGE_SIZE
#define PAGE_SIZE 4096
#endif

#define ALLOC_BASE 0xa00000

#define INITIAL_SIZE 0x2000
#define EXPAND_SIZE 0x1000

bool UserHeap::init()
{
    bitmap = (uint8_t*)kmalloc(INITIAL_SIZE);
    if (!bitmap) return false;
    bitmap_size = INITIAL_SIZE;
    memset(bitmap, 0, bitmap_size);
    kdbgln("new user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
    return true;
}

bool UserHeap::inherit(UserHeap& other)
{
    bitmap = (uint8_t*)kmalloc(other.bitmap_size);
    if (!bitmap) return false;
    bitmap_size = other.bitmap_size;
    memcpy(bitmap, other.bitmap, bitmap_size);
    kdbgln("child user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
    return true;
}

void UserHeap::free()
{
    kdbgln("freeing user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
    kfree(bitmap);
}

bool UserHeap::try_expand()
{
    kdbgln("attempting to expand user heap");
    void* new_bitmap = krealloc(bitmap, bitmap_size + EXPAND_SIZE);
    if (!new_bitmap)
    {
        kdbgln("expansion failed");
        return false;
    }
    bitmap = (uint8_t*)new_bitmap;
    memset(bitmap + bitmap_size, 0, EXPAND_SIZE);
    bitmap_size += EXPAND_SIZE;
    kdbgln("expanded user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
    return true;
}

bool UserHeap::try_expand_size(uint64_t size)
{
    void* new_bitmap = krealloc(bitmap, bitmap_size + size);
    if (!new_bitmap)
    {
        kdbgln("expansion failed");
        return false;
    }
    bitmap = (uint8_t*)new_bitmap;
    memset(bitmap + bitmap_size, 0, size);
    bitmap_size += size;
    kdbgln("expanded user heap, bitmap at %p, size %ld", (void*)bitmap, bitmap_size);
    return true;
}

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

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

uint64_t UserHeap::request_virtual_page()
{
    uint64_t attempts = 0;
allocate:
    for (uint64_t index = start_index; index < bitmap_size * 8; index++)
    {
        if (bitmap_read(index)) continue;
        bitmap_set(index, true);
        start_index = index + 1;
        return ALLOC_BASE + (index * PAGE_SIZE);
    }

    if (attempts == 0 && try_expand()) // We are allocating ONE PAGE, only one attempt should be necessary.
    {
        attempts++;
        goto allocate;
    }

    return 0;
}

uint64_t UserHeap::request_virtual_pages(uint64_t count)
{
    uint64_t attempts = 0;
allocate:
    uint64_t contiguous = 0;
    uint64_t contiguous_start = 0;
    for (uint64_t index = start_index; index < bitmap_size * 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);
        }
    }

    if (attempts == 0 && try_expand_size(Utilities::get_blocks_from_size(8, count) + 256))
    {
        attempts++;
        goto allocate;
    }

    return 0;
}

void UserHeap::free_virtual_page(uint64_t address)
{
    if (address < ALLOC_BASE || address >= (ALLOC_BASE + bitmap_size * 8 * PAGE_SIZE)) return;
    uint64_t index = (address - ALLOC_BASE) / PAGE_SIZE;
    bitmap_set(index, false);
    if (start_index > index) start_index = index;
}

void UserHeap::free_virtual_pages(uint64_t address, uint64_t count)
{
    if (address < ALLOC_BASE || address >= (ALLOC_BASE + bitmap_size * 8 * PAGE_SIZE)) 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;
}