#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <signal.h>
#include <sys/mman.h>
#include <dlfcn.h>

#include "stack.h"

void _mdbg_init();

#ifdef SOLARIS
#pragma init (_mdbg_init)
#endif
#ifdef LINUX
void _init() { _mdbg_init(); }
#endif

static size_t  _mdbg_private_len = 0;
static void   *_mdbg_private     = NULL;
static void   *_mdbg_base        = NULL;
static void   *_mdbg_curr        = NULL;

void *_malloc(size_t size);

//static unsigned long _mdbg_incr = 0;

static char mdbg_magic_head[] = "#MeM cHeCkEr: HeAdErMaGicStRiNg";
static char mdbg_magic_tail[] = "*mEm ChEcKeR: tAiLeRmAgICsTrInG";
#define mdbg_HEAD_LEN 32
#define mdbg_TAIL_LEN 32

#define mdbg_MAXSTACKS (1024 * 1024 * sizeof(void *))

typedef void *(*fcore)(size_t size);

static fcore _mdbg_sbrk = NULL;

static int _mdbg_lock = 0;

enum {mdbg_Unset, mdbg_Used, mdbg_Free};
static char *mdbg_status_str[] = {
    "Unset", "InUse", "Freed"
};

typedef struct {
    unsigned long nmalloc;
    unsigned long nrealloc;
    unsigned long nfree;
} stat_t, *pstat_t;

static stat_t _mdbg_stats;

void mdbg_priv_lock()  { 
    static const char str[]  = "mprotect lock OK\n";
    static const char fail[] = "mprotect lock FAILED\n";
    if (mprotect(_mdbg_private, _mdbg_private_len, PROT_NONE))
        write(2, fail, sizeof(fail));
}

void mdbg_priv_read()  { 
    static const char str[]  = "mprotect read unlock OK\n";
    static const char fail[] = "mprotect read unlock FAILED\n";
    if (mprotect(_mdbg_private, _mdbg_private_len, PROT_READ))
        write(2, fail, sizeof(fail));
}

void mdbg_priv_write() { 
    static const char str[]  = "mprotect write unlock OK\n";
    static const char fail[] = "mprotect write unlock FAILED\n";
    if (mprotect(_mdbg_private, _mdbg_private_len, PROT_READ | PROT_WRITE))
        write(2, fail, sizeof(fail));
}

void mdbg_lock()   {
    while (_mdbg_lock);
    _mdbg_lock = 1;
}
void mdbg_unlock() {
    _mdbg_lock = 0; 
}

typedef struct {
    size_t  size;
    size_t  status;
    void    *next;
    int             nstacks;
    void    *pstack;
    char    magic_head[mdbg_HEAD_LEN];
    char    align64[12];
} mdbg_block_head_t;

typedef struct {
    char    magic_tail[mdbg_TAIL_LEN];
    char    align64[mdbg_TAIL_LEN];
} mdbg_block_tail_t;

#define ALIGN 8 

#define mdbg_BUFF 1000
#define mdbg_ERR  2
#define mdbg_OUT  1

void _mdbg_dump_data(char *buff, size_t size) {
    int i, j;
    char hex[256], asc[256], tmp[256];
    if (buff == NULL) {
        fprintf(stderr, "   null pointer\n");
        return;
    }

    fprintf(stderr, "   block dump  :\n");

    /* 00 11 22 33 44 55 66 77 - 88 99 AA BB CC DD EE FF   azertyui - qsdfghjk */
    for (i = 0; i < size; i += 16) {
        hex[0] = asc[0] = '\0';
        for (j = 0; j < 16 && i + j < size; j++) {
            if (j == 8) {
                strcat(hex, "- ");
                strcat(asc, " - ");
            }
            sprintf(tmp, "%02x ", (unsigned char) buff[i+j]);
            strcat(hex, tmp);
            sprintf(tmp, "%c", isprint(buff[i + j]) ? buff[i + j] : '.');
            strcat(asc, tmp);
        }
        fprintf(stderr, "      %p: %-50s  %-19s\n", buff+i, hex, asc);
    }
}

void
_mdbg_hdr_dump(mdbg_block_head_t *h) {
    if (!h) return;
    fprintf(stderr, "haddr  = %x\n", h);
    fprintf(stderr, "size   = %u\n", h->size);
    if (h->status < mdbg_Unset || h->status > mdbg_Free) 
         fprintf(stderr, "invalid status\n");
    else fprintf(stderr, "status = %s\n", mdbg_status_str[h->status]);
    fprintf(stderr, "next   = %x\n", h->next);
    fprintf(stderr, "magic1 = \"%s\"\n", h->magic_head);
    fprintf(stderr, "magic2 = \"%s\"\n", (char *) h + sizeof(mdbg_block_head_t) + h->size);
    fprintf(stderr, "---------------------------\n\n");
}

int _brk(void *foo) {
    static char msg[] = "mdbg ALERT: _brk is directly called from program\n";
    write(2, msg, strlen(msg));
    if (!_mdbg_private) _mdbg_init();
    if (getenv("MDBG_DUMP")) {
        stack_live_dump();
        fprintf(stderr, "_brk(%p)\n", foo);
    }
    return 0;
}

int brk(void *foo) {
    static char msg[] = "mdbg ALERT: brk is directly called from program\n";
    write(2, msg, strlen(msg));
    if (!_mdbg_private) _mdbg_init();
    if (getenv("MDBG_DUMP")) {
        fprintf(stderr, "brk(%p)\n", foo);
        stack_live_dump();
    }
    return 0;
}

void *_sbkr(int size) {
    static char msg[] = "mdbg ALERT: _sbrk is directly called from program\n";
    write(2, msg, strlen(msg));
    if (!_mdbg_private) _mdbg_init();
    if (getenv("MDBG_DUMP")) {
        fprintf(stderr, "_sbrk(%u)\n", size);
        stack_live_dump();
    }
    return _malloc(size);
}

void *sbrk(int size) {
    static char msg[] = "mdbg ALERT: sbrk is directly called from program\n";
    write(2, msg, strlen(msg));
    if (!_mdbg_private) _mdbg_init();
    if (getenv("MDBG_DUMP")) {
        fprintf(stderr, "sbrk(%u)\n", size);
        stack_live_dump();
    }
    return _malloc(size);
}

void * 
_mdbg_morecore(size_t size) {
    return _mdbg_sbrk(size);        
}

void
mdbg_check_free(mdbg_block_head_t *hb) {
    int i, dump = 0;
    mdbg_block_tail_t *ht;

    ht = (mdbg_block_tail_t *) ((char *) hb + sizeof(mdbg_block_head_t) + hb->size);

    if (strcmp(hb->magic_head, mdbg_magic_head)) dump += 1;
    if (strcmp((char *) ht, mdbg_magic_tail))    dump += 2;

    if (dump) {
        fprintf(stderr, "trying to free corrupted block at %x:\n", hb);
        if (dump & 1) fprintf(stderr, "   magic header corrupted (\"%s\")\n", hb->magic_head);
        if (dump & 2) fprintf(stderr, "   magic tailer corrupted (\"%s\")\n", ht);
        if (hb->status < mdbg_Unset || hb->status > mdbg_Free) 
            fprintf(stderr, "   block status: %u (invalid)\n", hb->status);
        else 
            fprintf(stderr, "   block status: %s\n", mdbg_status_str[hb->status]);
        fprintf(stderr, "   block size  : %u\n", hb->size);
        fprintf(stderr, "   block next  : %x\n", hb->next);
        fprintf(stderr, "   user ptr    : %x\n", (char *) hb + sizeof(mdbg_block_head_t));
        _mdbg_dump_data((char *) hb + sizeof(mdbg_block_head_t), hb->size);
        fprintf(stderr, "   allocation time stack dump:\n");
        mdbg_priv_read();       
        stack_dump(hb->nstacks, hb->pstack);
        mdbg_priv_lock();
    }
}

void 
_mdbg_check(int n, mdbg_block_head_t *hb) {
    int i, dump = 0;
    mdbg_block_tail_t *ht;

    ht = (mdbg_block_tail_t *) ((char *) hb + sizeof(mdbg_block_head_t) + hb->size);

    if (strcmp(hb->magic_head, mdbg_magic_head)) 
        dump += 1;
    if (strcmp((char *) ht, mdbg_magic_tail)) 
        dump += 2;

    if (dump || getenv("MDBG_DUMP")) {
        fprintf(stderr, "\nBlock number %4d:\n", n);
        fprintf(stderr, "------------------\n", n);
        if   (dump & 1) fprintf(stderr, "   magic header corrupted\n");
        if   (dump & 2) fprintf(stderr, "   magic tailer corrupted\n");
        fprintf(stderr, "   block status: %s\n", mdbg_status_str[hb->status]);
        fprintf(stderr, "   block size  : %u\n", hb->size);
        fprintf(stderr, "   block next  : %x\n", hb->next);
        fprintf(stderr, "   user ptr    : %x\n", (char *) hb + sizeof(mdbg_block_head_t));
        _mdbg_dump_data((char *) hb + sizeof(mdbg_block_head_t), hb->size);
        fprintf(stderr, "   allocation time stack dump:\n");
        mdbg_priv_read();       
        stack_dump(hb->nstacks, hb->pstack);
        mdbg_priv_lock();
    }
}

void
mdbg_check_all() {
    static int count = 0;
    int i = 0;
    mdbg_block_head_t *hb;
    
    fprintf(stderr, "check memory blocks (%u)\n", ++count);
    hb = (mdbg_block_head_t *) _mdbg_base;
    do {
        _mdbg_check(i, hb);
        if ((long) hb->next < (long) hb) {
            fprintf(stderr, "block # %u is corrupted, stopping check\n", i);
            break;
        }
        hb = hb->next;
        i++;
    } while (hb);
} 

void _mdbg_stats_print() {
    fprintf(stderr, "mdbg stats:\n");
    fprintf(stderr, "mdbg stats: memory statistics:\n");
    fprintf(stderr, "mdbg stats: ------------------\n");
    fprintf(stderr, "mdbg stats: malloc:       %8d\n", _mdbg_stats.nmalloc);
    fprintf(stderr, "mdbg stats: realloc:      %8d\n", _mdbg_stats.nrealloc);
    fprintf(stderr, "mdbg stats: free:         %8d\n", _mdbg_stats.nfree);
    fprintf(stderr, "mdbg stats: still in use: %8d\n", _mdbg_stats.nmalloc - _mdbg_stats.nfree);
}   

void
_mdbg_check_all() {
    int i = 0;
    mdbg_block_head_t *hb;

    fprintf(stderr, "\n");
    fprintf(stderr, "--------------------\n");
    fprintf(stderr, " -- atexit check -- \n");
    fprintf(stderr, "--------------------\n");
    fprintf(stderr, "\n");

    hb = (mdbg_block_head_t *) _mdbg_base;
    do {
        _mdbg_check(i, hb);
        hb = hb->next;
        i++;
    } while (hb);
    _mdbg_stats_print();
} 

void
_mdbg_signals(int sig) {
    fprintf(stderr, "mdbg caught signal %u:\n", sig);
    stack_live_dump();
    fprintf(stderr, "Now preform memory check...\n");
    _mdbg_check_all();
    fprintf(stderr, "done\n");
}

size_t 
mdbg_aligned(size_t size) {
    unsigned int residue;
    if ((residue = size % ALIGN))
        size += (ALIGN - residue);
    return size;
}

void _mdbg_init() {
    long pagesize;
    pagesize = sysconf(_SC_PAGE_SIZE);

//write(2, "_mdbg_init\n", 12);

    /* fetch real "mem core function" */
    _mdbg_sbrk = (fcore) dlsym(RTLD_NEXT, "sbrk");

    /* make sure we start on pagesize multiple */
    if (((size_t) _mdbg_sbrk(0)) % pagesize) _mdbg_morecore(pagesize - ((size_t) _mdbg_sbrk(0) % pagesize));


    _mdbg_private = _mdbg_sbrk(0);
    stack_init(mdbg_MAXSTACKS);
    signal(SIGSEGV, _mdbg_signals);

    if (((size_t) _mdbg_sbrk(0)) % pagesize) _mdbg_morecore(pagesize - ((size_t) _mdbg_sbrk(0) % pagesize));
    _mdbg_base = _mdbg_sbrk(0);
    _mdbg_private_len = (char *) _mdbg_base - (char *) _mdbg_private - 1;
    _mdbg_stats.nmalloc  = 
    _mdbg_stats.nrealloc = 
    _mdbg_stats.nfree    = 0;

    mdbg_priv_lock();       
    
    atexit(_mdbg_check_all);
}

void *_mdbg_malloc(size_t size) {
    static int        count = 0;
    size_t            gsize, real;
    void              *p; 
    char                      *q, *r;
    mdbg_block_head_t *h;

    if (getenv("MDBG_ALWAYS_CHECK") && count > 0) mdbg_check_all();

    /* 
     * compute actual reservation size as : 
     * header + tailer + size (adjusted)
     */
    
    h = _mdbg_curr;

    gsize = mdbg_aligned(size);
    real = gsize + sizeof(mdbg_block_head_t) + sizeof(mdbg_block_tail_t);

    if ((p = _mdbg_morecore(real)) == (void *) -1) return NULL;
    if ((unsigned long) p % ALIGN) 
        fprintf(stderr, "ALIGNMENT ERROR: %x %% %u = %u\n", 
            (unsigned long) p, ALIGN, 
            (unsigned long) p % ALIGN);

    if (count) h->next = p;
    
    h = (mdbg_block_head_t *) (q = (char *) p);

    h->size   = gsize;
    h->status = mdbg_Used;
    h->next   = NULL;
    strcpy(h->magic_head, mdbg_magic_head);

/* Fill the allocated but not used bytes withs 'X' to check small leaks: */
    if (size < gsize) memset(q +  sizeof(mdbg_block_head_t) + size, 'X', gsize - size);
    strcpy(q + sizeof(mdbg_block_head_t) + gsize, mdbg_magic_tail);

    /* Find out a way to protect small amounts of mem HERE!!! */

    mdbg_priv_write();
    h->pstack  = stack_curr_get();
    h->nstacks = stack_trace();
    mdbg_priv_lock();

    _mdbg_curr = (void *) h;
    count++;

    if (getenv("MDBG_DBGS")) fprintf(stderr, "new block @ 0x%x\n", h);
    
    return q + sizeof(mdbg_block_head_t);
}

void _mdbg_free(void *ptr) {
    char *p;
    mdbg_block_head_t *hb;

    p = (char *) ptr;
    hb = (mdbg_block_head_t *) (p - sizeof(mdbg_block_head_t));

    mdbg_check_free(hb);
    hb->status = mdbg_Free;
}

void *_malloc(size_t size) {
    void *p;
    if (!_mdbg_private) {
        static char msg[] = "init called from malloc()\n";
        write(2, msg, strlen(msg));
        _mdbg_init();
    }
    if (getenv("MDBG_DBG")) {
        fprintf(stderr, "malloc(%u) called from:\n", size);
        stack_live_dump();
    }
    if (size == 0) fprintf(stderr, "malloc(0)\n");
    mdbg_lock();
    p = _mdbg_malloc(size); 
    if (!p) {
        fprintf(stderr, "%x = malloc(%u) FAILED\n", 0, size);
        if (getenv("MDBG_DBG")) stack_live_dump();      
    }
    _mdbg_stats.nmalloc++;
    mdbg_unlock();
    if (p && getenv("MDBG_DBG")) fprintf(stderr, "allocated %u bytes at %x\n", size, p);
    return p;       
}

void _free(void *ptr) {
    if (!_mdbg_private) return;
    if (getenv("MDBG_DBG")) {
        fprintf(stderr, "free(%x) called from:\n", ptr);
        stack_live_dump();
    }
    if (!ptr) {
        fprintf(stderr, "free(NULL):\n");
        stack_live_dump();
        return;
    } 
    mdbg_lock();
    _mdbg_free(ptr);
    _mdbg_stats.nfree++;
    mdbg_unlock();
    if (getenv("MDBG_DBG")) fprintf(stderr, "freed %x\n", ptr);
}

void *
_realloc(void *ptr, size_t newsize) {
    void *q;
    size_t size;

    if (!_mdbg_private) _mdbg_init();
    if (getenv("MDBG_DBG")) {
        fprintf(stderr, "realloc(%x, %u) called from:\n", ptr, newsize);
        stack_live_dump();
    }
    if (!ptr) {
        if (getenv("MDBG_REALLOC_AS_MALLOC")) {
            fprintf(stderr, "realloc used as malloc (null input pointer)\n");
            if (getenv("MDBG_DUMP")) stack_live_dump();
        }
        return _malloc(newsize);
    }

    mdbg_lock();

    size = (((mdbg_block_head_t *) ptr) - 1)->size;

    if (getenv("MDBG_ALWAYS_CHECK")) mdbg_check_all();
    
    _mdbg_free(ptr);
    if ((q = _mdbg_malloc(newsize))) {
        memcpy(q, ptr, newsize < size ? newsize : size);
        if (newsize < size) {
            fprintf(stderr, "realloc truncate current data:\n");
            stack_live_dump();
        }
    } else {
        fprintf(stderr, "%x = realloc(%x, %u) (was %u) FAILED\n", 0, ptr, newsize, size);
        stack_live_dump();
    }
    _mdbg_stats.nrealloc++;
    mdbg_unlock();
    fprintf(stderr, "reallocated %u (from %u) at %x (from %x)\n", newsize, size, ptr, q); 
    return q;
}

/* Redefine malloc, realloc and free : */
void *malloc(size_t size) { return _malloc(size); }
void *realloc(void *ptr, size_t newsize) { return _realloc(ptr, newsize); }
void free(void *ptr) { _free(ptr); }