#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include "jl_mem.h"

#ifdef jl_malloc_exit
#undef jl_malloc_exit
#undef jl_calloc_exit
#undef jl_realloc_exit
#endif

// ALIGN_SIZE should be larger than sizeof(malloc_info), and properly aligned.
// On x86_64, memory is generally aligned on 16 byte boundaries.
#define ALIGN_SIZE 32

void *jl_malloc_exit(size_t size)
{
    void *data = aligned_alloc(ALIGN_SIZE, size);
    if(data == NULL) {
        printf("Out of memory.\n");
        exit(EXIT_FAILURE);
    }
    return data;
}

typedef struct malloc_info malloc_info;

struct malloc_info
{
    char  *file;
    int    line;
    size_t size;
    malloc_info *prev, *next;
};

static malloc_info *mi_head = NULL;
static pthread_mutex_t lock;

static size_t alloc_count = 0;
static size_t free_count = 0;
static size_t alloc_bytes = 0;
static size_t free_bytes = 0;

void jl_mem_leak_check_init()
{
    if(pthread_mutex_init(&lock, NULL) != 0) {
        printf("Failed to initialize mutex.\n");
        exit(EXIT_FAILURE);
    }
    atexit(jl_mem_print_leaks);
}

void jl_free_chk(void *ptr)
{
    if(ptr == NULL) {
        return;
    }

    pthread_mutex_lock(&lock);

    malloc_info *mi = ptr - ALIGN_SIZE;

    if(mi->prev == NULL) {
        mi_head = mi->next;
    }

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

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

    free_bytes += mi->size;
    free_count++;

    free(mi);

    pthread_mutex_unlock(&lock);
}

void *jl_malloc_chk(size_t size, char *file, int line)
{
    pthread_mutex_lock(&lock);

    malloc_info *mi = aligned_alloc(ALIGN_SIZE, size + ALIGN_SIZE);
    if(mi == NULL) {
        printf("\nMALLOC FAILED:\n");
        printf("    File: %s\n", file);
        printf("    Line: %d\n", line);
        printf("    Size: %u bytes\n", (unsigned int)size);
        exit(EXIT_FAILURE);
    }

    mi->file = file;
    mi->line = line;
    mi->size = size;
    alloc_bytes += size;
    alloc_count++;

    mi->prev = NULL;
    mi->next = mi_head;

    if(mi_head) {
        mi_head->prev = mi;
    }
    mi_head = mi;

    pthread_mutex_unlock(&lock);

    return (void*)mi + ALIGN_SIZE;
}

void jl_mem_print_leaks(void)
{
    int count = 0;
    malloc_info *mi = mi_head;
    while(mi) {
        printf("MEMORY LEAK DETECTED\n");
        printf("    File: %s\n", mi->file);
        printf("    Line: %d\n", (int)mi->line);
        printf("    Size: %d bytes\n", (int)mi->size);
        mi = mi->next;
        ++count;
        if(count > 10) {
            printf("ADDITIONAL LEAKS NOT PRINTED.\n");
            break;
        }
    }

    printf("\n");
    printf("MEMORY SUMMARY\n");
    printf("    Num allocs: %lu\n", (unsigned long)alloc_count);
    printf("    Num frees:  %lu\n", (unsigned long)free_count);
    printf("\n");
    printf("    Bytes alloced: %lu\n", (unsigned long)alloc_bytes);
    printf("    Bytes freed:   %lu\n", (unsigned long)free_bytes);
    printf("    Bytes leaked:  %lu\n",
           (unsigned long)(alloc_bytes - free_bytes));
    printf("\n");
}