#include <lib/memory/alloc.h>
#include <stdint.h>
#include <lib/textMode/stdio.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) {
        // this is currently the last block -> initialize the new last block
        next = (MemoryBlock*) ((void*) (current + 1) + size);
        next->last = current;
        next->state = FREE;
        next->next = 0x00;
        next->size = -1;
    } else {
        // there is a next block
        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;
    } else {
        current->next = next;
        if (current->next->next == 0x00) {
            current->size = size;
        }
    }
    current->state = IN_USE;
    current->task = task;
    return (void*) (current + 1);
}

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

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

void free(void* location) {
    MemoryBlock* current = ((MemoryBlock*) location) - 1;
    current->state = FREE;
    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 printMemoryStack() {
    MemoryBlock* current = firstBlock;
    while (current->next != 0x00) {
        printf("AT: %x size: %x next: %x last: %x, state: %x\n", current, current->size, current->next, current->last, current->state);
        current = current->next;
    }
    printf("AT: %x size: %x next: %x last: %x, state: %x\n", current, current->size, current->next, current->last, current->state);
}

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