#include "axl_memory.h"

static u8 memory[AXL_HEAP_SIZE];

#if AXL_HEAP_SIZE <= 0xFF - 1
typedef u8 mb_handle;
#define AXL_INVALID_MB_HANDLE 0xFF
#elif AXL_HEAP_SIZE <= 0xFFFF - 1
typedef u16 mb_handle;
#define AXL_INVALID_MB_HANDLE 0xFFFF
#elif AXL_HEAP_SIZE <= 0xFFFFFFFF - 1
typedef u32 mb_handle;
#define AXL_INVALID_MB_HANDLE 0xFFFFFFFF
#else
#error UNSUPPORTED HEAP SIZE
#endif

typedef struct mb_header mb_header;

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

static mb_handle root_handle = 0;
static mb_handle nomad_handle = 0;

#define AXL_ALIGNMENT 8

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

static mb_header* axl_id_to_mb(mb_handle id)
{
    if(id == AXL_INVALID_MB_HANDLE)
    {
        return NULL;
    }

    return (mb_header*)&memory[id];
}

static mb_handle axl_mb_to_id(mb_header* ptr)
{
    return (mb_handle)((u8*)ptr - memory);
}

#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);
        mb_header* root = axl_id_to_mb(root_handle);
        root->is_free = true;
        root->next = AXL_INVALID_MB_HANDLE;
        root->prev = AXL_INVALID_MB_HANDLE;
        root->size = AXL_HEAP_SIZE - MB_HEADER_SIZE;
        axl_initialized = true;
    }
}

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

static mb_header* axl_find_mb(u32 req_size)
{
    mb_header* block = nomad_handle ? axl_id_to_mb(nomad_handle) : axl_id_to_mb(root_handle); 
    mb_header* start = block;

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

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

        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 = axl_mb_to_id(block);

            if(new_block->next != AXL_INVALID_MB_HANDLE)
            {
                axl_id_to_mb(new_block->next)->prev = axl_mb_to_id(new_block);
            }

            block->size = size; 
            block->next = axl_mb_to_id(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_handle = free_block->next;
    if(nomad_handle == AXL_INVALID_MB_HANDLE) 
    {
        nomad_handle = 0;
    }
    
    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_ptr_to_mb(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;
}

i32 axl_memcmp(const void* s1, const void* s2, u32 n)
{
    const u8* a = (u8*)s1;
    const u8* b = (u8*)s2;

    for(u32 i = 0; i < n; i++)
    {
        if(a[i] != b[i])
        {
            return a[i] - b[i]; //checked: will do negative values too
        }
    }

    return 0;
}

void* axl_memchr(const void* ptr, u8 value, u32 n)
{
    if(ptr)
    {
        for(u32 i = 0; i < n; i++)
        {
            if(*((u8*)ptr + i) == value)
            {
                return (u8*)ptr + i;
            }
        }
    }

    return NULL;
}

void* axl_memmove(void* dst, const void* src, u32 n)
{
    if(!dst || !src || !n)
    {
        return dst;
    }

    const u8* s = (const u8*)src;
    u8* d = (u8*)dst;

    if(d > s)
    {
        s += n - 1;
        d += n - 1;

        while(n--)
        {
            *d-- = *s--;
        }
    }
    else 
    {
        while(n--)
        {
            *d++ = *s++;
        }
    }

    return dst;
}

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

    mb_header* block = axl_ptr_to_mb(ptr);
    mb_header* next_block = axl_id_to_mb(block->next);
    mb_header* prev_block = axl_id_to_mb(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 != AXL_INVALID_MB_HANDLE)
        {
            axl_id_to_mb(block->next)->prev = axl_mb_to_id(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 != AXL_INVALID_MB_HANDLE)
        {
            axl_id_to_mb(prev_block->next)->prev = axl_mb_to_id(prev_block);
        }
        free_block = prev_block;
    }
    
    nomad_handle = axl_mb_to_id(free_block);
}

#if !defined(__STDC_HOSTED__) || __STDC_HOSTED__ == 0
void* memset(void* s, int c, size_t n)
{
    return axl_memset(s, (i8)c, n);
}

void* memcpy(void* dst, const void* src, unsigned long long count)
{
    return axl_memcpy(dst, src, (u32)count);
}

void* memmove(void* dst, const void* src, size_t n)
{
    return axl_memmove(dst, src, n);
}
#endif