#include "heap.h"
#include "vmm.h"
#include "pmm.h"
#include "../drivers/serial.h"

static heap_chunk_t* heap_start_chunk = NULL;

void kheap_init(void) {
    kprint_serial("HEAP: Initializing Kernel Heap...\n");

    // 1. Выделяем одну физ. страницу под старт кучи
    void* phys_page = pmm_alloc_page();
    
    // 2. Мапим её на виртуальный адрес HEAP_START
    // Нам нужны флаги Present (1) + Write (2) = 3
    vmm_map_page((uint32_t)phys_page, HEAP_START, 3);

    // 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;

    kprint_serial("HEAP: Kernel Heap initialized successfully at 0x00400000!\n");
}

void* kmalloc(size_t size) {
    if (size == 0) return NULL;

    // Выравниваем размер по 4 байта для стабильности на x86
    size = (size + 3) & ~3;

    heap_chunk_t* current = heap_start_chunk;

    // Ищем подходящий свободный блок (First Fit)
    while (current != NULL) {
        if (current->is_free && current->size >= size) {
            
            // Если блок сильно больше, чем надо, дробим его
            // sizeof(heap_chunk_t) + 4 байта минимально на остаток
            if (current->size >= size + sizeof(heap_chunk_t) + 4) {
                heap_chunk_t* new_chunk = (heap_chunk_t*)((uint32_t)current + sizeof(heap_chunk_t) + size);
                new_chunk->size = current->size - size - sizeof(heap_chunk_t);
                new_chunk->is_free = 1;
                new_chunk->next = current->next;

                current->size = size;
                current->next = new_chunk;
            }

            current->is_free = 0;
            
            // Возвращаем указатель на память СРАЗУ ЗА заголовком
            return (void*)((uint32_t)current + sizeof(heap_chunk_t));
        }
        current = current->next;
    }

    // По-хорошему, если памяти не хватило, нужно выделить еще страницу из PMM 
    // и расширить кучу. Но для первого теста выкинем панику.
    kprint_serial("HEAP ERROR: OUT OF KERNEL HEAP MEMORY!\n");
    return NULL;
}

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;

    // Склеиваем последовательные свободные блоки (Coalescing), чтобы не было фрагментации
    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;
            // Не шагаем вперед, проверяем текущий блок еще раз с новым next
            continue; 
        }
        current = current->next;
    }
}