#include <_stdio.h>
#include <alloc.h>
#include <stdint.h>

MemoryBlock *firstBlock;

void initMemoryAllocation(uint32_t kernelEnd) {
    firstBlock = (MemoryBlock *)kernelEnd;
    firstBlock->next = 0x00;
    firstBlock->last = 0x00;
    firstBlock->state = FREE;
    firstBlock->size = -1;
}

void *mallocTask(uint32_t size, Task *task) {
    MemoryBlock *current = firstBlock;
    while (!(current->next == 0x00 ||
             current->state == FREE && current->size >= size)) {
        current = current->next;
    }
    MemoryBlock *next;
    if (current->next == 0x00) {
        next = ((void *)(&current[1]) + size);
        next->last = current;
        next->state = FREE;
        next->next = 0x00;
        next->size = -1;
    } else {
        next = current->next;
    }
    if (current->size > size + sizeof(MemoryBlock) * 2 &&
        current->next != 0x00) {
        // segment the current block if it is big enough
        current->next = (MemoryBlock *)(((void *)(current + 1)) + size);
        current->next->size = current->size - sizeof(MemoryBlock) - size;
        current->size = size;
        current->next->state = FREE;
        current->next->next = next;
        next->last = current->next;
        next = current->next;
    } else {
        current->next = next;
    }
    current->size = ((void *)next) - (void *)(current + 1);
    current->state = IN_USE;
    current->task = task;
    return (void *)(current + 1);
}

void *malloc(uint32_t size) {
    return mallocTask(size, (Task *)getCurrentTask());
}

void *mallocAligned(uint32_t size, uint8_t alignment) {
    uint64_t buffer = (uint64_t)malloc(size + (1 << alignment));
    buffer >>= alignment - 1;
    buffer <<= alignment - 1;
    buffer += 1 << alignment;
    return (void *)buffer;
}

void mergeNextIfPossible(MemoryBlock *current) {
    if (current->next != 0x00 && current->next->state == FREE) {
        current->size += sizeof(MemoryBlock) + current->next->size;
        current->next = current->next->next;
        current->next->last = current;
    }
}

void free(void *location) {
    MemoryBlock *current = ((MemoryBlock *)location) - 1;
    current->state = FREE;
    for (uint32_t i = 0; i < current->size; i++) {
        ((uint8_t *)&current[1])[i] = 0;
    }
    mergeNextIfPossible(current);
    if (current->last != 0x00 && current->last->state == FREE) {
        current = current->last;
    }
    mergeNextIfPossible(current); // call this function twice in case a used
                                  // block between 2 free blocks is freed
    mergeNextIfPossible(current);
}

void freeTaskAllocations(Task *task) {
    MemoryBlock *current = firstBlock;
    while (current->next != 0x00) {
        if (current->task == task) {
            free(current + 1);
        }
        current = current->next;
    }
    if (current->task == task) {
        free(current + 1);
    }
}