/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   Kernel heap
 * @license GPL-3.0-only
 */

#include <mem/kheap.h>
#include <mem/paging.h>
#include <mem/pmm.h>
#include <stddef.h>
#include <kernel.h>
#include <sched/process.h>
#include <config.h>

extern uint64_t kernel_phys_base;
extern uint64_t kernel_virt_base;

uintptr_t kheap_start;

static struct heap_block_t* head = NULL;
static uintptr_t end;

// Kernel root table (level 4)
extern uint64_t *kernel_pml4;

/*
 * kheap_init - Kernel heap initialization
 *
 * This function physically allocates and maps enough pages
 * of memory for KHEAP_SIZE, which is defined in config.h.
 *
 * It then creates one big heap block, which will be the
 * base for a linked list.
 */
void kheap_init()
{
    kheap_start = ALIGN_UP(kernel_virt_base + KERNEL_SIZE, PAGE_SIZE);
    
    size_t heap_pages = ALIGN_UP(KHEAP_SIZE, PAGE_SIZE) / PAGE_SIZE;
    DEBUG("Mapping %d kernel heap pages at 0x%p", heap_pages, kheap_start);

    uintptr_t current_addr = kheap_start;

    // Map/alloc enough pages for heap (KHEAP_SIZE)
    for (size_t i=0; i<heap_pages; i++) {
        uintptr_t phys = pmm_alloc();
        if (phys == 0) {
            panic(NULL, "Not enough memory available to initialize kernel heap.");
        }

        paging_map_page(kernel_pml4, current_addr, phys, PTE_PRESENT | PTE_WRITABLE);
        current_addr += PAGE_SIZE;
    }

    end = current_addr;

    // Give linked list head its properties
    head = (struct heap_block_t*)kheap_start;
    head->size = (end-kheap_start) - sizeof(struct heap_block_t);
    head->free = true;
    head->next = NULL;
    DEBUG("Kernel heap initialized, head=0x%p, size=%u bytes", head, head->size);
}

/*
 * kmalloc - Kernel memory allocation
 * @size: number of bytes to allocate
 *
 * Looks for a big enough free block and marks it
 * as taken. Each block of memory is preceded by
 * the linked list header.
 *
 * Return:
 * <ptr> - Pointer to at least <size> bytes of usable memory
 * NULL  - No more memory, or no valid size given
 */
void* kmalloc(size_t size)
{
    // No size, no memory allocated!
    if (!size) return NULL;
    size = ALIGN(size);

    struct heap_block_t* curr = head;

    while (curr) {
        // Is block free and big enough for us?
        if (curr->free && curr->size >= size) {
            // We split the block if it is big enough
            if (curr->size >= size + sizeof(struct heap_block_t) + 16) {
                struct heap_block_t* split = (struct heap_block_t*)((uintptr_t)curr + sizeof(struct heap_block_t) + size);
                
                split->size = curr->size - size - sizeof(struct heap_block_t);
                split->free = true;
                split->next = curr->next;

                curr->next = split;
                curr->size = size;
            }

            // Found a good block, we return it
            curr->free = false;
            return (void*)((uintptr_t)curr + sizeof(struct heap_block_t));
        }
        // Continue browsing the list if nothing good was found yet
        curr = curr->next;
    }

    // No growing. If we're here it means the initial pool
    // wasn't sufficient. Too bad.
    DEBUG("Kernel heap is OUT OF MEMORY!");
    // if we were terrorists maybe we should panic
    // or just wait for others to free stuff?
    return NULL;
}

/*
 * kfree - Kernel memory freeing
 * @ptr: pointer to memory region to free
 *
 * Marks the memory block beginning at <ptr>
 * as free. Also merges adjacent free blocks
 * to lessen fragmentation.
 */
void kfree(void* ptr)
{
    // Nothing to free
    if (!ptr) return;

    // Set it free!
    struct heap_block_t* block = (struct heap_block_t*)((uintptr_t)ptr - sizeof(struct heap_block_t));
    block->free = true;

    // merge adjacent free blocks (coalescing)
    struct heap_block_t* curr = head;
    while (curr && curr->next) {
        if (curr->free && curr->next->free) {
            curr->size += sizeof(*curr) + curr->next->size;
            curr->next = curr->next->next;
            continue;
        }
        curr = curr->next;
    }
}

/*
 * kalloc_stack - Stack memory allocation
 *
 * Allocates a memory region of at least PROCESS_STACK_SIZE,
 * to be used as a stack for a process. The pointer returned
 * points to the end of the region, as the stack grows downwards.
 *
 * Return:
 * <ptr> - Pointer to a region after at least PROCESS_STACK_SIZE bytes of usable memory
 * NULL  - No more memory
 */
void* kalloc_stack()
{
    uint8_t* ptr = kmalloc(PROCESS_STACK_SIZE); // As it's out of kmalloc, stack is already mapped into kernel space
    return ptr ? ptr+PROCESS_STACK_SIZE : NULL;
}