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

MemoryBlock *firstBlock;

void initMemoryAllocation(uint32_t kernelEnd) {
    firstBlock = (MemoryBlock *)kernelEnd;
    firstBlock->next = NULL;
    firstBlock->last = NULL;
    firstBlock->task = NULL;
    firstBlock->size = -1;
}

#define getOffset()                                                            \
    (alignment ? alignmentBit - ((uint32_t)(current + 1) & (alignmentMask)) : 0)

void memmove(uint8_t *to, uint8_t *from, uint32_t count) {
    if (from > to) {
        for (; count > 0; count--) {
            *(to++) = *(from++);
        }
        return;
    }
    for (to += count - 1, from += count - 1; count > 0; count--) {
        *(to--) = *(from--);
    }
}

void *mallocTaskAligned(uint32_t size, uint8_t alignment, Task *task) {
    MemoryBlock *current = firstBlock;
    uint32_t alignmentBit = 1 << alignment;
    uint32_t alignmentMask = alignmentBit - 1;
    uint32_t additionalOffset = getOffset();
    while (current->task || current->size < size + additionalOffset) {
        current = current->next;
        additionalOffset = getOffset();
    }
    uint32_t neededSize = size + additionalOffset;
    if (current->size > neededSize + sizeof(MemoryBlock) * 2) {
        MemoryBlock *next = current->next;
        current->size = neededSize;
        current->next = ((void *)(current + 1)) + neededSize;
        current->next->last = current;
        current->next->next = next;
        if (next) {
            current->next->size = (void *)next - (void *)(current->next + 1);
        } else {
            current->next->size = -1;
        }
        current->next->task = NULL;
        if (next) {
            next->last = current->next;
        }
    }
    if (additionalOffset) {
        memmove((uint8_t *)current + additionalOffset, (uint8_t *)current,
                sizeof(MemoryBlock));
        current = (void *)current + additionalOffset;
        current->size = (void *)current->next - (void *)(current + 1);
        current->last->next = current;
        current->last->size = (void *)current - (void *)(current->last + 1);
    }
    current->task = task;
    return (void *)(current + 1);
}

void *mallocTask(uint32_t size, Task *task) {
    return mallocTaskAligned(size, 0, task);
}

void *malloc(uint32_t size) {
    return mallocTaskAligned(size, 0, getCurrentTask());
}

void *mallocAligned(uint32_t size, uint8_t alignment) {
    return mallocTaskAligned(size, alignment, getCurrentTask());
}

void mergeNextIfPossible(MemoryBlock *current) {
    if (current->next && !current->next->task) {
        current->next = current->next->next;
        current->size = (void *)current->next - (void *)(current + 1);
        current->next->last = current;
    }
}

void free(void *location) {
    MemoryBlock *current = ((MemoryBlock *)location) - 1;
    current->task = NULL;
    for (uint32_t i = 0; i < current->size; i++) {
        ((uint8_t *)&current[1])[i] = 0;
    }
    mergeNextIfPossible(current);
    if (current->last && !current->last->task) {
        current = current->last;
        mergeNextIfPossible(current);
    }
    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);
    }
}