#include "../include/paging.h"
#include <stdint.h>
#include "../include/stdio.h"

// Структура заголовка блока памяти
typedef struct Block {
    uint32_t size;      // Размер блока (в байтах, включая заголовок)
    uint32_t is_free;   // 1 = свободен, 0 = занят
    struct Block* next; // Указатель на следующий блок
} Block;

// Структура для управления кучей
typedef struct {
    void* start;    // Начало кучи
    uint32_t size;  // Текущий размер кучи в байтах
    Block* first;   // Первый блок в куче
} Heap;

Heap kernel_heap;

void heap_init() {
    kernel_heap.start = (void*)0xD0000000; // Начало кучи в ядре
    kernel_heap.size = 0x1000;            // Начальный размер — 1 страница (4 КБ)
    printf("before map_kernel_page\n");
    // Отображаем первую страницу
    map_kernel_page(kernel_heap.start, 0x2); // writable, supervisor
    printf("after map_kernel_page\n");
    // Создаём первый блок, который занимает всю страницу
    Block* initial_block = (Block*)kernel_heap.start;
    initial_block->size = 0x1000 - sizeof(Block); // Размер без учёта заголовка
    initial_block->is_free = 1;                   // Свободен
    initial_block->next = (Block*)0;              // Пока нет следующего

    kernel_heap.first = initial_block;
}
/*
void* heap_alloc(uint32_t size) {
    // Выравниваем размер до 4 байт
    size = (size + 3) & ~3;

    Block* current = kernel_heap.first;
    Block* prev = (Block*)0;

    // Ищем подходящий свободный блок
    while (current) {
        if (current->is_free && current->size >= size) {
            if (current->size >= size + sizeof(Block) + 4) {
                Block* new_block = (Block*)((uint32_t)current + sizeof(Block) + size);
                new_block->size = current->size - size - sizeof(Block);
                new_block->is_free = 1;
                new_block->next = current->next;

                current->size = size;
                current->next = new_block;
            }
            current->is_free = 0;
            return (void*)((uint32_t)current + sizeof(Block));
        }
        prev = current;
        current = current->next;
    }

    // Нет подходящего блока — выделяем нужное количество страниц
    uint32_t total_size = size + sizeof(Block); // Размер с заголовком
    uint32_t pages_needed = (total_size + 0xFFF) / 0x1000; // Округляем вверх до страниц
    uint32_t alloc_size = pages_needed * 0x1000; // Сколько байт выделяем

    void* new_page = (void*)((uint32_t)kernel_heap.start + kernel_heap.size);
    for (uint32_t i = 0; i < pages_needed; i++) {
        map_kernel_page((void*)((uint32_t)new_page + i * 0x1000), 0x2);
    }
    kernel_heap.size += alloc_size;

    Block* new_block = (Block*)new_page;
    new_block->size = alloc_size - sizeof(Block);
    new_block->is_free = 1;
    new_block->next = (Block*)0;

    if (prev) prev->next = new_block;
    if (!kernel_heap.first) kernel_heap.first = new_block;

    // Теперь current указывает на новый блок, продолжаем цикл
    current = new_block;
    if (current->size >= size) {
        if (current->size >= size + sizeof(Block) + 4) {
            Block* split_block = (Block*)((uint32_t)current + sizeof(Block) + size);
            split_block->size = current->size - size - sizeof(Block);
            split_block->is_free = 1;
            split_block->next = current->next;

            current->size = size;
            current->next = split_block;
        }
        current->is_free = 0;
        return (void*)((uint32_t)current + sizeof(Block));
    }

    // Если что-то пошло не так, возвращаем NULL
    return (void*)0;
}
*/
//TODO: fix it yo
void *heap_alloc(uint32_t size)
{
    return alloc_page();
}

void heap_free(void* ptr) {
    if (!ptr || (uint32_t)ptr < (uint32_t)kernel_heap.start) return;

    // Находим блок по указателю (указатель указывает после заголовка)
    Block* block = (Block*)((uint32_t)ptr - sizeof(Block));
    block->is_free = 1;

    // Слияние с предыдущим блоком, если он свободен
    Block* current = kernel_heap.first;
    Block* prev = (Block*)0;
    while (current && current != block) {
        prev = current;
        current = current->next;
    }
    if (prev && prev->is_free) {
        prev->size += sizeof(Block) + block->size;
        prev->next = block->next;
        block = prev;
    }

    // Слияние с следующим блоком, если он свободен
    if (block->next && block->next->is_free) {
        block->size += sizeof(Block) + block->next->size;
        block->next = block->next->next;
    }
}