#include "MemoryManager.h"
#include "Types.h"
#include <stdlib.h>
#include <string.h>

/** M_std **/

Result
std_alloc(MemoryManager mm, void **ptr, usize size) {
    if (NULL == ptr)
        return R_NullArgument;

    *ptr = malloc(size);
    return (*ptr == NULL) ? R_Err : R_Ok;
}

Result
std_realloc(MemoryManager mm, void **ptr, usize size) {
    if (NULL == ptr)
        return R_NullArgument;

    *ptr = realloc(*ptr, size);
    return (*ptr == NULL) ? R_Err : R_Ok;
}

Result
std_free(MemoryManager mm, void **ptr) {
    if (NULL == ptr)
        return R_NullArgument;

    free(*ptr);
    *ptr = NULL;
    return R_Ok;
}

Result
std_destroy(MemoryManager *mm) {
    return R_Unavailable;
}

MemoryManager
M_std() {
    return (MemoryManager){
        .alloc = &std_alloc,
        .realloc = &std_realloc,
        .free = &std_free,
        .destroy = &std_destroy,
    };
}

/** M_std_tracked **/

struct MemoryTracking {
    usize alloc;
    void **ptrs;
};

static inline Result
MemoryTracking_store(struct MemoryTracking *mt, void *ptr) {
    if (NULL == mt || NULL == ptr)
        return R_NullArgument;

    if (NULL == mt->ptrs || 0 == mt->alloc)
        return R_Uninitialized;

    usize index = 0;
    while (index < mt->alloc && NULL != mt->ptrs[index]) {
        ++index;
    }

    if (index == mt->alloc) {
        mt->ptrs = realloc(mt->ptrs, sizeof(void*) * (mt->alloc << 1));
        memset(mt->ptrs + (sizeof(void*) * mt->alloc), 0, sizeof(void*) * mt->alloc);

        if (NULL == mt->ptrs) {
            mt->alloc = 0;
            return R_OutOfMemory;
        }

        mt->alloc <<= 1;
    }

    mt->ptrs[index] = ptr;
    return R_Ok;
}

static inline Result
MemoryTracking_remove(struct MemoryTracking *mt, void *ptr) {
    if (NULL == mt || NULL == ptr)
        return R_NullArgument;

    if (NULL == mt->ptrs || 0 == mt->alloc)
        return R_Uninitialized;

    usize index = 0;
    while (index < mt->alloc && ptr != mt->ptrs[index]) {
        ++index;
    }

    if (index == mt->alloc)
        return R_NotFound;

    mt->ptrs[index] = NULL;
    return R_Ok;
}

Result
stdt_alloc(MemoryManager mm, void **ptr, usize size) {
    if (NULL == ptr)
        return R_NullArgument;

    if (NULL == mm.priv)
        return R_Uninitialized;

    *ptr = malloc(size);
    if (NULL == *ptr)
        return R_Err;

    return MemoryTracking_store(mm.priv, *ptr);
}

Result
stdt_realloc(MemoryManager mm, void **ptr, usize new_size) {
    if (NULL == ptr)
        return R_NullArgument;

    if (NULL == mm.priv)
        return R_Uninitialized;

    void *ptr_orig = *ptr;
    *ptr = realloc(*ptr, new_size);

    if (NULL == *ptr)
        return R_Err;

    if (ptr_orig != *ptr) {
        RETURN_NOTOK(MemoryTracking_remove(mm.priv, ptr_orig));
        RETURN_NOTOK(MemoryTracking_store(mm.priv, *ptr));
    }
    return R_Ok;
}

Result
stdt_free(MemoryManager mm, void **ptr) {
    if (NULL == ptr)
        return R_NullArgument;

    if (NULL == mm.priv)
        return R_Uninitialized;

    RETURN_NOTOK(MemoryTracking_remove(mm.priv, *ptr));

    free(*ptr);
    *ptr = NULL;

    return R_Ok;
}

Result
stdt_destroy(MemoryManager *mm) {
    if (NULL == mm)
        return R_NullArgument;

    if (NULL == mm->priv)
        return R_Uninitialized;

    struct MemoryTracking *mt = mm->priv;

    for (usize i = 0; i < mt->alloc; ++i) {
        if (NULL != mt->ptrs[i])
            free(mt->ptrs[i]);
    }

    free(mt->ptrs);
    free(mt);
    mm->priv = NULL;

    return R_Ok;
}

MemoryManager
M_std_tracked() {
    struct MemoryTracking *mt = malloc(sizeof(struct MemoryTracking));
    mt->alloc = 8;
    mt->ptrs = malloc(sizeof(void*) * mt->alloc);
    memset(mt->ptrs, 0, sizeof(void*) * mt->alloc);

    return (MemoryManager) {
        .alloc = &stdt_alloc,
        .realloc = &stdt_realloc,
        .free = &stdt_free,
        .destroy = &stdt_destroy,
        .priv = mt,
    };
}

/* Helper functions */

char*
M_strdup(MemoryManager mm, const char* str) {
    if (NULL == str)
        return NULL;

    char* newstr = NULL;
    if (R_Ok != mm.alloc(mm, (void**)&newstr, strlen(str) + 1) ||
        NULL == newstr)
    {
        return NULL;
    }

    strcpy(newstr, str);
    return newstr;
}