#include <stdint.h>

#include <kernel/heap.h>
#include <kernel/tty.h>
#include <debugging.h>

#define HEAP_SIZE (1 << 16)
#define NUM_PAGES (128)

uint8_t global_heap_data[HEAP_SIZE + sizeof(struct Heap_Block)];
uint8_t pages_heap_data[NUM_PAGES * 4096 + sizeof(struct Heap_Block)] __attribute__((aligned(4096)));

struct Heap_Block pages_heap_blocks[NUM_PAGES];

struct Heap_Metadata global_heap;
struct Heap_Metadata pages_heap;

void* nearest_aligned(void* ptr, size_t alignment)
{
    // TODO: Fix this solution up, it's *extremely* inneficient for numbers that aren't powers of 2
    
    size_t address = (size_t)ptr;

    if (address % alignment == 0) {
        return (void*) address;
    }

    if ((alignment & (alignment - 1)) == 0) {
        address = address & ~(alignment - 1);
        address += alignment;
        return (void*) address;
    }

    while (address % alignment != 0) {
        address++;
    }

    return (void*) address;
}

size_t sizeof_block(struct Heap_Block* block, size_t alignment)
{
    size_t start = (size_t)nearest_aligned(block->data, alignment);
    size_t end = (size_t)block->next;
    
    return end - start;
}

void heap_init() 
{
    global_heap.size = HEAP_SIZE;
    global_heap.start = (struct Heap_Block*) global_heap_data;

    struct Heap_Block* start = global_heap.start;
    start->size = HEAP_SIZE;
    start->used = false;
    start->data = start + 1;
    start->next = 0;

    struct Heap_Block* page_start = pages_heap_blocks;

    pages_heap.size = NUM_PAGES;
    pages_heap.start = (struct Heap_Block*) page_start;

    page_start->size = NUM_PAGES;
    page_start->used = false;
    page_start->data = pages_heap_data;
    page_start->next = 0;
}

void* malloc(size_t size)
{
    struct Heap_Block* current = global_heap.start;
    
    while (true) {
        if (current->used || current->size < size) { 
            current = current->next; 
            continue; 
        }

        if (current == 0) {
            terminal_writestring("WARNING: Heap Full\n");
            return 0;
        }

        struct Heap_Block* next = current->data + size;

        next->next = current->next;
        next->size = current->size - size - sizeof(struct Heap_Block);
        next->data = next + 1;

        current->next = next;
        current->size = size;
        current->used = true;

        return current->data;
    }
}

void* alloc_page(size_t num_pages, struct Heap_Block** ptr)
{
    struct Heap_Block* current = pages_heap.start;
    
    while (true) {
        if (current->used || current->size < num_pages) { 
            current = current->next; 
            continue; 
        }

        if (current == 0) {
            terminal_writestring("WARNING: Heap Full\n");
            return 0;
        }

        if ((size_t)current->data % 4096 != 0) {
            terminal_writestring("ERROR: Pages are unaligned");
        }

        size_t current_index = ((size_t)current->data - (size_t)pages_heap_data) / 4096;

        struct Heap_Block* next = pages_heap_blocks + (current_index + num_pages);
            
        next->data = current->data + 4096 * num_pages;
        next->next = current->next;
        next->size = current->size - num_pages;
            
        current->next = next;
        current->size = num_pages;
        current->used = true;

        *ptr = current;

        return current->data;
    }
}

void free(void *ptr)
{
    struct Heap_Block* current = global_heap.start;

    struct Heap_Block* first_unused;
    bool found_unused = false;

    while (current->data != ptr) { 
        if (!current->used && !found_unused) {
            first_unused = current;
            found_unused = true;
        } else {
            found_unused = false;
        }
        current = current->next;
    }

    if (!found_unused) {
        first_unused = current;
    }

    current->used = false;

    if (current->next == 0) {
        return;
    }

    struct Heap_Block* start = first_unused;
    current = start;

    while (current->used == false && current->next != 0) {
        current = current->next;
    }

    start->next = current;
    start->size = (size_t)start->next - (size_t)start->data;
}

void free_page(struct Heap_Block *ptr)
{
    struct Heap_Block* current = pages_heap.start;

    struct Heap_Block* first_unused;
    bool found_unused = false;

    while (current != ptr) { 
        if (!current->used && !found_unused) {
            first_unused = current;
            found_unused = true;
        } else {
            found_unused = false;
        }
        current = current->next;
    }

    if (!found_unused) {
        first_unused = current;
    }

    current->used = false;

    if (current->next == 0) {
        return;
    }

    struct Heap_Block* start = first_unused;
    current = start;

    while (current->used == false && current->next != 0) {
        current = current->next;
    }

    start->next = current;
    start->size = ((size_t)start->next - (size_t)start->data) / 4096;
}