#include <stdio.h>
//#include <pthread.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <malloc.h>
#include "util.h"
#include "mem.h"

/* 本文中不使用自定义 malloc,calloc,free 等函数*/
#undef malloc
#undef calloc
#undef realloc
#undef strdup
#undef strndup
#undef free
#undef vasprintf
#undef asprintf
#define SOME_PRIME 563
#define FENCE_MAGIC 0xdeadbeef

static struct mm_region {
    struct mm_region *next;
    char file[40];              // 分配所在的文件
    char func[40];              // 分配所在的函数
    unsigned int lineno;        // 分配所在的行
    size_t len;                 // 内存分配的大小
    unsigned int fence;         // 内存起始边界,用于头越界检测
    unsigned char data[0];      // 用户内存分配首地址， malloc 等函数返回以此为首地址的 len 长度的一块内存
} *regions[SOME_PRIME];

#define HASH(a) \
    (((unsigned long)(a)) % SOME_PRIME)

//static pthread_mutex_t mmlock = PTHREAD_MUTEX_INITIALIZER;
#define mm_log(...) \
    do { \
    fprintf(stderr, __VA_ARGS__); \
    } while (0)

static void *__mem_alloc_region(size_t size, const char *file, int lineno, const char *func)
{
    struct mm_region *reg;
    void *ptr = NULL;
    unsigned int *fence;
    int hash;
    if (!(reg = (struct mm_region *)malloc(size + sizeof(*reg) + sizeof(*fence) ))){
        mm_log("Memory Allocation Failure - '%d' bytes in \
            function %s " "at line %d of %s\n", (int) size, func, lineno, file);
    }
    strncpy(reg->file, file, sizeof(reg->file));
    strncpy(reg->func, func, sizeof(reg->func));
    reg->lineno = lineno;
    reg->len = size;
    ptr = reg->data;
    hash = HASH(ptr);
    // 内存起始标志
    reg->fence = FENCE_MAGIC;
    // 内存结束标志
    fence = (unsigned int *)((unsigned char *)ptr + reg->len);
    *fence =FENCE_MAGIC;

    //pthread_mutex_lock(&mmlock);
    reg->next = regions[hash]; /* 一个 hash 可能对应多个值*/
    regions[hash] = reg;
    //pthread_mutex_unlock(&mmlock);

    return ptr;
}
static size_t __mem_sizeof_region(void *ptr)
{
    int hash = HASH(ptr);
    struct mm_region *reg;
    size_t len = 0;
    //pthread_mutex_lock(&mmlock);
    for (reg = regions[hash]; reg; reg = reg->next) {
        if (reg->data == ptr) {
            len = reg->len;
            break;
        }
    }
    //pthread_mutex_unlock(&mmlock);
    return len;
}
static void __mem_free_region(void *ptr, const char *file, int lineno, const char *func)
{
    int hash = HASH(ptr);
    struct mm_region *reg, *prev = NULL;
    unsigned int *fence;
    //pthread_mutex_lock(&mmlock);
    for (reg = regions[hash]; reg; reg = reg->next) {
        if (reg->data == ptr) {
            if (prev)
                prev->next = reg->next;
            else
                regions[hash] = reg->next;
            break;
        }
        prev = reg;
    }
    //pthread_mutex_unlock(&mmlock);
    if (reg) {
        // 头越界检测
        if (reg->fence != FENCE_MAGIC) {
            mm_log("WARNING: Head fence violation at %p, in %s of %s, line %d\n", 
                reg->data, reg->func, reg->file, reg->lineno); 
        }
        // 尾越界检测
        fence = (unsigned int *)(reg->data + reg->len);
        if ( *fence != FENCE_MAGIC) {
            mm_log("WARNING: Tail fence violation at %p, in %s of %s, "
                "line %d\n", reg->data, reg->func, reg->file, reg->lineno); 
        }
        free(reg);
    } else {
        mm_log("WARNING: Freeing unused memory at %p, in %s of %s,line %d\n",
            ptr, func, file, lineno);
    }
}

void *__mem_calloc(size_t nmemb, size_t size, const char *file, int lineno, const char *func)
{
    void *ptr;
    if ((ptr = __mem_alloc_region(size * nmemb, file, lineno, func)))
        memset(ptr, 0, size * nmemb);
    return ptr;
}

void *__mem_malloc(size_t size, const char *file, int lineno, const char *func) 
{
    return __mem_alloc_region(size, file, lineno, func);
}

void __mem_free(void *ptr, const char *file, int lineno, const char *func)
{
    __mem_free_region(ptr, file, lineno, func);
}

void *__mem_realloc(void *ptr, size_t size, const char *file, int lineno, const char *func)
{
    void *tmp; size_t len = 0;
    if (ptr && !(len = __mem_sizeof_region(ptr))) { 
        mm_log("WARNING: Realloc of unalloced memory at %p, in %s of %s,line %d\n", ptr, func, file, lineno);
        return NULL;
    }
    if (!(tmp = __mem_alloc_region(size, file, lineno, func)))
        return NULL;
    if (len > size) len = size;
    if (ptr) {
        memcpy(tmp, ptr, len);
        __mem_free_region(ptr, file, lineno, func);
    }
    return tmp;
}

char *__mem_strdup(const char *s, const char *file, int lineno, const char *func) 
{
    size_t len;
    void *ptr;
    if (!s)
        return NULL;

    len = strlen(s) + 1;
    if ((ptr = __mem_alloc_region(len, file, lineno, func)))
        strcpy((char *)ptr, s);

    return (char *)ptr;
}

char *__mem_strndup(const char *s, size_t n, const char *file, int lineno, const char *func)
{
    size_t len;
    void *ptr;
    if (!s)
        return NULL;

    len = strlen(s) + 1;
    if (len > n)
        len = n;

    if ((ptr = __mem_alloc_region(len, file, lineno, func)))
        strcpy((char *)ptr, s);

    return (char *)ptr;
}

int __mem_asprintf(const char *file, int lineno, const char *func, char **strp, const char *fmt, ...)
{
    int size;
    va_list ap, ap2;
    char s;
    *strp = NULL;
    va_start(ap, fmt);
    va_copy(ap2, ap);
    size = vsnprintf(&s, 1, fmt, ap2);
    va_end(ap2);
    if (!(*strp = (char *)__mem_alloc_region(size + 1, file, lineno, func))) {
        va_end(ap);
        return -1;
    }

    vsnprintf(*strp, size + 1, fmt, ap); 
    va_end(ap);

    return size;
}

int __mem_vasprintf(char **strp, const char *fmt, va_list ap, const char *file, int lineno, const char *func)
{
    int size;
    va_list ap2;
    char s;
    *strp = NULL;
    va_copy(ap2, ap);
    size = vsnprintf(&s, 1, fmt, ap2);
    va_end(ap2);
    if (!(*strp = (char *)__mem_alloc_region(size + 1, file, lineno, func))) {
        va_end(ap);
        return -1;
    }
    vsnprintf(*strp, size + 1, fmt, ap);

    return size;
}

int show_memory(void)
{
    char *fn = NULL;
    struct mm_region *reg;
    unsigned int x;
    unsigned int len = 0;
    unsigned int count = 0;
    unsigned int *fence;
    mm_log("\nLEAK DETAIL:\n");
    //pthread_mutex_lock(&mmlock);
    for (x = 0; x < SOME_PRIME; x++) {
        for (reg = regions[x]; reg; reg = reg->next) { 
            if (!fn || !strcasecmp(fn, reg->file) || !strcasecmp(fn,"anomolies"))
            {
                // 头越界检测
                if (reg->fence != FENCE_MAGIC) {
                    mm_log("WARNING: Head fence violation at %p, "
                        "in %s of %s, line %d\n", reg->data,
                        reg->func, reg->file, reg->lineno);
                }
                // 尾越界检测
                fence = (unsigned int *)(reg->data + reg->len);
                if ( *fence != FENCE_MAGIC) {
                    mm_log("WARNING: Tail fence violation at %p, in %s of %s, line %d\n", 
                        reg->data, reg->func, reg->file, reg->lineno);
                }
            }
            if (!fn || !strcasecmp(fn, reg->file)) {
                mm_log("%10d bytes allocated in %20s at line %5d of %s\n", (int)reg->len, reg->func, reg->lineno, reg->file); 
                len += reg->len;
                count++;
            }
        }
    }
    //pthread_mutex_unlock(&mmlock);
    mm_log("%d bytes allocated in %d allocations\n", len, count);

    return 0;
}

int show_memory_summary(void)
{
    char *fn = NULL;
    int x;
    struct mm_region *reg;
    unsigned int len = 0;
    int count = 0;
    struct file_summary {
        char fn[80];
        int len;
        int count;
        struct file_summary *next;
    } *list = NULL, *cur;

    mm_log("\nLEAK SUMMARY:\n");
    //pthread_mutex_lock(&mmlock);
    for (x = 0; x < SOME_PRIME; x++) {
        for (reg = regions[x]; reg; reg = reg->next) {
            if (fn && strcasecmp(fn, reg->file))
                continue;
            for (cur = list; cur; cur = cur->next) {
                if ((!fn && !strcmp(cur->fn, reg->file)) || 
                    (fn && !strcmp(cur->fn, reg->func)))
                    break;
            }
            if (!cur) {
                cur = (struct file_summary *)alloca(sizeof(*cur));
                memset(cur, 0, sizeof(*cur));
                strncpy(cur->fn, fn ? reg->func : reg->file, sizeof(cur->fn));
                cur->next = list;
                list = cur;
            }
            cur->len += reg->len;
            cur->count++;
        }
    }
    //pthread_mutex_unlock(&mmlock);
    /* Dump the whole list */
    for (cur = list; cur; cur = cur->next) {
        len += cur->len;
        count += cur->count;
        if (fn) {
            mm_log("%10d bytes in %d allocations in function '%s' of '%s'\n", 
                cur->len, cur->count, cur->fn, fn);
        } 
        else {
            mm_log("%10d bytes in %d allocations in file '%s'\n",cur->len, cur->count, cur->fn);
        }
    }
    mm_log("%d bytes allocated in %d allocations\n", len, count);

    return 0;
}