#ifndef MEMDEBUG_H
#define MEMDEBUG_H

/* OPTIONS:
 * - MEMDEBUG_IMPLEMENTATION - include the implementation
 *   (should be inlcuded in only one source file)
 * - MEMDEBUG_(MALLOC|REALLOC|CALLOC|FREE)_SYMBOL - change
 *   the function used for the operation
 * - MEMDEBUG_OUT_OF_BOUNDS - enable the out-of-bounds
 *   check
 * - MEMDEBUG_VOID - the main function is main(void)
 *   instead of main(int argc, char **argv)
 * - MEMDEBUG_OUTPUT_LOG - output the log in a file
 *   (format is "memdebug-<timestamp>.log")
 * - MEMDEBUG_OUTPUT_DIR - set the directory for the log,
 *   by default it is the current directory (automatically
 *   enables the MEMDEBUG_OUTPUT_LOG
 */

void *__memdebug_malloc(size_t size, char *file, int line);
void *__memdebug_calloc(size_t nmemb, size_t size, char *file, int line);
void *__memdebug_realloc(void *ptr, size_t size, char *file, int line);
void __memdebug_free(void *ptr, char *file, int line);

#ifdef MEMDEBUG_IMPLEMENTATION

// Default memory allocation functions
#ifndef MEMDEBUG_MALLOC_SYMBOL
#define MEMDEBUG_MALLOC_SYMBOL malloc
#endif

#ifndef MEMDEBUG_REALLOC_SYMBOL
#define MEMDEBUG_REALLOC_SYMBOL realloc
#endif

#ifndef MEMDEBUG_CALLOC_SYMBOL
#define MEMDEBUG_CALLOC_SYMBOL calloc
#endif

#ifndef MEMDEBUG_FREE_SYMBOL
#define MEMDEBUG_FREE_SYMBOL free
#endif

// Log output
#ifdef MEMDEBUG_OUTPUT_DIR
#define MEMDEBUG_OUTPUT_LOG
#else
#define MEMDEBUG_OUTPUT_DIR "."
#endif

#ifdef MEMDEBUG_OUTPUT_LOG
#define MEMDEBUG_OUTPUT_FMT "memdebug-%lu.log"
#endif

// Out-of-bounds check
#define MEMDEBUG_MAGIC_SUFFIX 0xDEADB00BCAFEF00D // 64 bit

typedef long long int memdebug_suffix;
#define MEMDEBUG_OUT_OF_BOUNDS_EXTRA_SIZE \
    (sizeof(size_t) + sizeof(memdebug_suffix))


// Implementation

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <time.h>

static FILE *_memdebug_fp = NULL;

#define MEMDEBUG_LOG(...) \
    fprintf(_memdebug_fp, __VA_ARGS__)
#define MEMDEBUG_LOG_FUNC(func, ret, file, line)                \
    do {                                                        \
        MEMDEBUG_LOG("(%s:%4d) %-8s ", file, line, #func);      \
        if(ret == NULL)                                         \
            MEMDEBUG_LOG("FAILED %d (%s) ",                     \
                         errno, strerror(errno));               \
    } while(0)

static inline void *memdebug_add_out_of_bounds_check(void *addr, size_t size)
{
    uintptr_t size_addr = (uintptr_t)addr;
    uintptr_t suffix_addr = size_addr +
        sizeof(size) + size;

    memcpy((void *)size_addr, &size, sizeof(size));

    memdebug_suffix suffix = MEMDEBUG_MAGIC_SUFFIX;
    memcpy((void *)suffix_addr, &suffix, sizeof(suffix));

    return (void *)((uintptr_t)addr + sizeof(size));
}

void *__memdebug_malloc(size_t size, char *file, int line)
{
#ifndef MEMDEBUG_OUT_OF_BOUNDS
    void *addr = MEMDEBUG_MALLOC_SYMBOL(size);
#else
    void *addr = MEMDEBUG_MALLOC_SYMBOL(
        size + MEMDEBUG_OUT_OF_BOUNDS_EXTRA_SIZE);
    if(addr) addr = memdebug_add_out_of_bounds_check(addr, size);
#endif

    MEMDEBUG_LOG_FUNC(malloc, addr, file, line);

    MEMDEBUG_LOG("size: %zu, ret: %p", size, addr);
    MEMDEBUG_LOG("\n");

    fflush(_memdebug_fp);
    return addr;
}

void *__memdebug_realloc(void *ptr, size_t size, char *file, int line)
{
#ifndef MEMDEBUG_OUT_OF_BOUNDS
    void *addr = MEMDEBUG_REALLOC_SYMBOL(ptr, size);
#else
    void *addr = MEMDEBUG_REALLOC_SYMBOL(
        ptr, size + MEMDEBUG_OUT_OF_BOUNDS_EXTRA_SIZE);
    if(addr) addr = memdebug_add_out_of_bounds_check(addr, size);
#endif

    MEMDEBUG_LOG_FUNC(realloc, addr, file, line);

    MEMDEBUG_LOG("ptr: %p, size: %zu, ret: %p", ptr, size, addr);
    MEMDEBUG_LOG("\n");

    fflush(_memdebug_fp);
    return addr;
}

void *__memdebug_calloc(size_t nmemb, size_t size, char *file, int line)
{
#ifndef MEMDEBUG_OUT_OF_BOUNDS
    void *addr = MEMDEBUG_CALLOC_SYMBOL(nmemb, size);
#else
    void *addr = MEMDEBUG_MALLOC_SYMBOL(
        nmemb * size + MEMDEBUG_OUT_OF_BOUNDS_EXTRA_SIZE);
    if(addr) addr = memdebug_add_out_of_bounds_check(addr, nmemb * size);
    memset(addr, 0, nmemb * size);
#endif

    MEMDEBUG_LOG_FUNC(calloc, addr, file, line);

    MEMDEBUG_LOG("nmemb: %zu, size: %zu, ret: %p", nmemb, size, addr);
    MEMDEBUG_LOG("\n");

    fflush(_memdebug_fp);
    return addr;
}

void __memdebug_free(void *ptr, char *file, int line)
{
    MEMDEBUG_LOG_FUNC(free, (void *)1, file, line);

    MEMDEBUG_LOG("ptr: %p", ptr);

#ifdef MEMDEBUG_OUT_OF_BOUNDS
    if(ptr != NULL) {
        size_t size = *(size_t *)((uintptr_t)ptr - sizeof(size_t));
        memdebug_suffix suffix = 0;
        memcpy(&suffix, (void *)((uintptr_t)ptr + size), sizeof(suffix));

        MEMDEBUG_LOG(", ");
        MEMDEBUG_LOG("out-of-bounds-check: ");

        if(suffix == (memdebug_suffix)MEMDEBUG_MAGIC_SUFFIX)
            MEMDEBUG_LOG("SUCCESS");
        else MEMDEBUG_LOG("FAILED");

        ptr = (void *)((uintptr_t)ptr + sizeof(size_t));
    }
#endif

    MEMDEBUG_LOG("\n");

    fflush(_memdebug_fp);
    MEMDEBUG_FREE_SYMBOL(ptr);
}

#ifdef MEMDEBUG_MAIN_VOID
int real_main(void);
#define CALL_MAIN real_main()
#else
int real_main(int argc, char **argv);
#define CALL_MAIN real_main(argc, argv)
#endif

int main(int argc, char **argv)
{
#ifdef MEMDEBUG_MAIN_VOID
    (void)argc;
    (void)argv;
#endif

#ifdef MEMDEBUG_OUTPUT_LOG
    size_t filename_sz = 64 + sizeof(MEMDEBUG_OUTPUT_DIR) + sizeof(MEMDEBUG_OUTPUT_FMT);
    char *filename = malloc(filename_sz);
    if(!filename) return -ENOMEM;

    memset(filename, 0, filename_sz);
    snprintf(filename, filename_sz, MEMDEBUG_OUTPUT_DIR"/"MEMDEBUG_OUTPUT_FMT, time(NULL));

    _memdebug_fp = fopen(filename, "w");
    if(!_memdebug_fp) {
        perror(filename);
        free(filename);
        return errno;
    }

    int ret = CALL_MAIN;

    fclose(_memdebug_fp);
    free(filename);
    return ret;
#else
    _memdebug_fp = stdout;
    return CALL_MAIN;
#endif
}

#define main real_main

#endif

#define __MEMDEBUG_CALL(f, ...) f(__VA_ARGS__, __FILE__, __LINE__)

#define malloc(size)        __MEMDEBUG_CALL(__memdebug_malloc,         size)
#define calloc(nmemb, size) __MEMDEBUG_CALL(__memdebug_calloc,  nmemb, size)
#define realloc(ptr, size)  __MEMDEBUG_CALL(__memdebug_realloc, ptr,   size)
#define free(ptr)           __MEMDEBUG_CALL(__memdebug_free,    ptr)

#endif