#include "heap.h"
#include "vmm.h"
#include "pmm.h"

static heap_chunk_t* heap_start_chunk = NULL;

static uint32_t current_heap_top = HEAP_START + PAGE_SIZE;

int heap_expand(size_t size) {
    size_t total_needed = size + sizeof(heap_chunk_t);
    size_t pages_needed = (total_needed + PAGE_SIZE - 1) / PAGE_SIZE;

    heap_chunk_t* new_chunk = (heap_chunk_t*)current_heap_top;

    for (size_t i = 0; i < pages_needed; i++) {
        void* phys = pmm_alloc_page();
        if (!phys) {
            return 0;
        }

        vmm_map_page((uint32_t)phys, current_heap_top, 3);

        if (i < pages_needed - 1) {
            current_heap_top += PAGE_SIZE;
        }
    }

    new_chunk->size = (pages_needed * PAGE_SIZE) - sizeof(heap_chunk_t);
    new_chunk->is_free = 1;
    new_chunk->next = NULL;

    heap_chunk_t* curr = heap_start_chunk;
    while (curr->next != NULL) {
        curr = curr->next;
    }
    curr->next = new_chunk;

    current_heap_top += PAGE_SIZE;
    return 1;
}

void kheap_init(void) {
    void* phys_page = pmm_alloc_page();
    if (phys_page == NULL) {
        while(1);
    }

    vmm_map_page((uint32_t)phys_page, HEAP_START, 3);

    heap_start_chunk = (heap_chunk_t*)HEAP_START;
    heap_start_chunk->size = PAGE_SIZE - sizeof(heap_chunk_t);
    heap_start_chunk->is_free = 1;
    heap_start_chunk->next = NULL;
}

void* kmalloc(size_t size) {
    if (size == 0) return NULL;
    size = (size + 3) & ~3;

    heap_chunk_t* current = heap_start_chunk;

    while (current != NULL) {
        if (current->is_free && current->size >= size) {
            current->is_free = 0;
            return (void*)((uint32_t)current + sizeof(heap_chunk_t));
        }
        current = current->next;
    }

    if (!heap_expand(size)) {
        return NULL;
    }
    return kmalloc(size);
}

void kfree(void* ptr) {
    if (ptr == NULL) return;

    heap_chunk_t* chunk = (heap_chunk_t*)((uint32_t)ptr - sizeof(heap_chunk_t));
    chunk->is_free = 1;

    heap_chunk_t* current = heap_start_chunk;
    while (current != NULL) {
        if (current->is_free && current->next && current->next->is_free) {
            current->size += sizeof(heap_chunk_t) + current->next->size;
            current->next = current->next->next;
            continue;
        }
        current = current->next;
    }
}
