#ifndef TK_MEMORY_H
#define TK_MEMORY_H

#define tk_malloc(size) tk_internal_malloc(size, __FILE__, __LINE__)
#define tk_free(size) tk_internal_free(size, __FILE__, __LINE__)
#define tk_realloc(ptr, size) tk_internal_realloc(ptr, size, __FILE__, __LINE__)

#include <stddef.h>

#ifdef __cplusplus
extern "C" {
#endif

void* tk_internal_malloc(size_t size, const char* file, int line);
void tk_internal_free(void* ptr, const char* file, int line);
void* tk_internal_realloc(void* ptr, const size_t size, const char* file, const int line);

#ifdef __cplusplus
}
#endif

#endif /* TK_MEMORY_H */

#ifdef TK_MEMORY_IMPLEMENTATION

struct tk_memory_chunk {
    size_t size;
    void* addr;
    const char* file;
    int line;

    struct tk_memory_chunk* next;
};

struct tk_memory_chunk* tk_internal_memory_chunks = NULL;

#include <stdlib.h>
#include <string.h>

inline void* tk_internal_malloc(const size_t size, const char* file, const int line) {
    void* ptr = malloc(size);

    if (ptr == NULL) {
        return NULL;
    }

#ifdef DEBUG
    memset(ptr, 0xFE, size);

#ifdef TK_MEMORY_TRACK
    struct tk_memory_chunk* chunk = malloc(sizeof(struct tk_memory_chunk));
    if (chunk != NULL) {
        chunk->size = size;
        chunk->addr = ptr;
        chunk->file = file;
        chunk->line = line;
        chunk->next = NULL;

        if (tk_internal_memory_chunks == NULL) {
            tk_internal_memory_chunks = chunk;
        } else {
            struct tk_memory_chunk* current = tk_internal_memory_chunks;
            while (current->next != NULL) {
                current = current->next;
            }
            current->next = chunk;
        }
    }
#endif
#endif

    return ptr;
}

inline void tk_internal_free(void* ptr, const char* file, const int line) {
    free(ptr);
#ifdef DEBUG
#ifdef TK_MEMORY_TRACK
    struct tk_memory_chunk* prev = NULL;
    struct tk_memory_chunk* current = tk_internal_memory_chunks;
    while (current != NULL) {
        if (current->addr == ptr) {
            if (prev != NULL) {
                prev->next = current->next;
            }
            if (prev == NULL && current == tk_internal_memory_chunks) {
                tk_internal_memory_chunks = current->next;
            } else if (prev == NULL && current == tk_internal_memory_chunks && current->next == NULL) {
                tk_internal_memory_chunks = NULL;
            }

            free(current);

            return;
        }
        prev = current;
        current = current->next;
    }
#endif
#endif
}

inline void* tk_internal_realloc(void* ptr, const size_t size, const char* file, const int line) {
    void* new_ptr = realloc(ptr, size);
#ifdef DEBUG
#ifdef TK_MEMORY_TRACK
    if (new_ptr != NULL) {
        struct tk_memory_chunk* current = tk_internal_memory_chunks;
        while (current != NULL) {
            if (current->addr == ptr) {
                current->addr = new_ptr;
                current->size = size;
                current->file = file;
                current->line = line;
                break;
            }
            current = current->next;
        }
    }
#endif
#endif

    return new_ptr;
}

void tk_check_memory_leaks() {
#ifdef TK_MEMORY_TRACK
    struct tk_memory_chunk* current = tk_internal_memory_chunks;
    while (current != NULL) {
        printf("Memory leak detected at %s:%d\n", current->file, current->line);
        current = current->next;
    }
#endif
}

#endif /* TK_MEMORY_IMPLEMENTATION */