#include "axl_memory.h"

static u8 memory[AXL_HEAP_SIZE];

typedef struct mb_header mb_header;

struct mb_header
{
    u32 size;
    b8 is_free;
    mb_header* prev;
    mb_header* next;
};

static mb_header* root = NULL;
static mb_header* nomad = NULL;

#define AXL_ALIGNMENT 8

static inline u32 axl_align(u32 size)
{
    return (size + (AXL_ALIGNMENT - 1)) & ~(AXL_ALIGNMENT - 1);
}

#define MB_HEADER_SIZE axl_align(sizeof(mb_header))

void* axl_memset(void* ptr, i8 c, u32 n)
{
    if(ptr)
    {
        for(u32 i = 0; i < n; i++)
        {
            ((i8*)ptr)[i] = c; 
        }
    }

    return ptr;
}

void axl_init(void)
{
    static b8 axl_initialized = false;

    if(!axl_initialized)
    {
        axl_memset(memory, 0, AXL_HEAP_SIZE);
        root = (mb_header*)memory;
        root->is_free = true;
        root->next = NULL;
        root->prev = NULL;
        root->size = AXL_HEAP_SIZE - MB_HEADER_SIZE; 
        axl_initialized = true;
    }
}

static inline mb_header* axl_get_mb_header(void* ptr)
{
    return (mb_header*)((u8*)ptr - MB_HEADER_SIZE);
}

static mb_header* axl_find_mb(u32 req_size)
{
    mb_header* block = nomad ? nomad : root; 
    mb_header* start = block;

    while(block)
    {
        if(block->is_free && block->size >= req_size)
        {
            return block;
        }
        
        block = block->next;

        if(!block)
        {
            //nomad pointer didn't help us, start from root
            block = root;
        }

        if(block == start)
        {
            //full circle, we didn't find a free memory block
            break;
        }
    }

    return NULL;
}

static b8 axl_split_mb(mb_header* block, u32 size)
{
    if(block->size > size)
    {
        if(block->size - size >= MB_HEADER_SIZE + AXL_ALIGNMENT)
        {
            mb_header* new_block = (mb_header*)((u8*)block + MB_HEADER_SIZE + size);
            new_block->size = block->size - size - MB_HEADER_SIZE;
            new_block->is_free = true; 
            new_block->next = block->next;
            new_block->prev = block;

            if(new_block->next)
            {
                new_block->next->prev = new_block;
            }

            block->size = size; 
            block->next = new_block;

            return true;
        }
    }

    return false;
}

void* axl_malloc(u32 size)
{
    if(size == 0)
    {
        return NULL;
    }

    size = axl_align(size);

    if(size > AXL_HEAP_SIZE - MB_HEADER_SIZE)
    {
        return NULL;
    }

    mb_header* free_block = axl_find_mb(size);

    if(!free_block)
    {
        return NULL;
    }

    axl_split_mb(free_block, size);
    free_block->is_free = false;

    nomad = free_block->next;
    if(!nomad) 
    {
        nomad = root;
    }
    
    return (void*)((u8*)free_block + MB_HEADER_SIZE);
}

void* axl_memcpy(void* dst, const void* src, u32 count)
{
    for(u32 i = 0; i < count; i++)
    {
        *((u8*)dst + i) = *((u8*)src + i);
    }

    return dst;
}

void* axl_realloc(void* ptr, u32 size)
{
    if(size == 0)
    {
        axl_free(ptr);
        return NULL;
    }

    size = axl_align(size);

    if(!ptr)
    {
        return axl_malloc(size);
    }

    mb_header* old = axl_get_mb_header(ptr);
    if(axl_split_mb(old, size))
    {
        return ptr;
    }  

    void* new = axl_malloc(size);

    if(!new)
    {
        return NULL;
    }

    size = size < old->size ? size : old->size;
    axl_memcpy(new, ptr, size);

    axl_free(ptr);

    return new;
}

void axl_free(void* ptr)
{
    if(!ptr)
    {
        return;
    }

    mb_header* block = axl_get_mb_header(ptr);
    mb_header* next_block = block->next;
    mb_header* prev_block = block->prev;
    block->is_free = true; 

    mb_header* free_block = block;

    if(next_block != NULL && next_block->is_free)
    {
        block->size += next_block->size + MB_HEADER_SIZE;
        block->next = next_block->next;
        if(block->next)
        {
            block->next->prev = block;
        }
    }

    if(prev_block != NULL && prev_block->is_free)
    {
        prev_block->size += block->size + MB_HEADER_SIZE;
        prev_block->next = block->next;
        if(prev_block->next)
        {
            prev_block->next->prev = prev_block;
        }
        free_block = prev_block;
    }
    
    nomad = free_block;
}