#define MODULE "kheap"

#include "memory/KernelHeap.h"
#include "assert.h"
#include "log/Log.h"
#include "std/string.h"

#ifndef PAGE_SIZE
#define PAGE_SIZE 4096
#endif

static uint8_t page_bitmap[2048];

static int64_t kheap_free = sizeof(page_bitmap) * 8 * PAGE_SIZE;
static int64_t kheap_used = 0;

#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; }
}

void KernelHeap::clear()
{
    memset(page_bitmap, 0, sizeof(page_bitmap));
    kinfoln("page bitmap located at %p", (void*)page_bitmap);
}

uint64_t KernelHeap::request_virtual_page()
{
    for (uint64_t index = 0; index < sizeof(page_bitmap) * 8; index++)
    {
        if (bitmap_read(index)) continue;
        bitmap_set(index, true);
        // start_index = index + 1;
#ifdef KHEAP_DEBUG
        kinfoln("allocating one page for caller %p, returning %lx", __builtin_return_address(0),
                ALLOC_BASE + (index * PAGE_SIZE));
#endif
        kheap_free -= PAGE_SIZE;
        kheap_used += PAGE_SIZE;
#ifdef KHEAP_DEBUG
        dump_usage();
#endif
        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 = 0; 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);
#ifdef KHEAP_DEBUG
            kinfoln("allocating %lu pages for caller %p, returning %lx", count, __builtin_return_address(0),
                    ALLOC_BASE + (contiguous_start * PAGE_SIZE));
#endif
            kheap_free -= (count * PAGE_SIZE);
            kheap_used += (count * PAGE_SIZE);
#ifdef KHEAP_DEBUG
            dump_usage();
#endif
            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);
#ifdef KHEAP_DEBUG
    kinfoln("releasing one page for caller %p, %lx", __builtin_return_address(0), address);
#endif
    kheap_free += PAGE_SIZE;
    kheap_used -= PAGE_SIZE;
#ifdef KHEAP_DEBUG
    dump_usage();
#endif
    // 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); }
#ifdef KHEAP_DEBUG
    kinfoln("releasing %lu pages for caller %p, %lx", count, __builtin_return_address(0), address);
#endif
    kheap_free += (count * PAGE_SIZE);
    kheap_used -= (count * PAGE_SIZE);
#ifdef KHEAP_DEBUG
    dump_usage();
#endif
    // if (start_index > index) start_index = index;
}

void KernelHeap::dump_usage()
{
    kinfoln("Used: %ld KB", kheap_used / 1024);
    kinfoln("Free: %ld KB", kheap_free / 1024);
}