#include "precompile.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdarg.h>
#include <ctype.h>
#include <assert.h>

#include "fileutil.h"
#include "toolkit.h"
#include "memutil.h"

#include <winpr/wlog.h>
#include "dbgutil.h"
#include "memtrace.h"
#undef malloc
#undef calloc
#undef realloc
#undef free
#undef strdup

#define TAG TOOLKIT_TAG(".memtrace")

/** This is a special maker which is set right before the beginning of a
 *  memory chunk and right after the end of a memory chunk.  When memory is
 *  freed, the markers are checked to see if they were overwritten by the
 *  client */
#ifndef MEMTRACE_MARK
#define MEMTRACE_MARK ((char)0x55)
#endif

 /** Used to write over freed memory, and write over malloced memory. Must be
  *  different than MEMTRACE_MARK */
#ifndef MEMTRACE_JUNK
#define MEMTRACE_JUNK ((char)0xAA)
#endif

/** How many bytes should be written before and after every allocation to be
*  checked for possible overflows when memory is freed?  Must be >= 0.
*  Should be a multiple of sizeof(wchar_t) to prevent bugs with glibc */
#ifndef MEMTRACE_GUARDSIZE
#define MEMTRACE_GUARDSIZE 4
#endif

#define MEMTRACE_HEADER_SIZE (sizeof(mem_trace_t) + MEMTRACE_GUARDSIZE)

/** The number of functions to keep track of at a time. The larger the stack
  *  size, the more effecient the function stack printing can be */
#ifndef MEMTRACE_FUNCSTACK_SIZE
#define MEMTRACE_FUNCSTACK_SIZE 1024
#endif

 /** Should we add code to keep track of function calls? */
#ifndef MEMTRACE_TRACK_FUNCTIONS
#define MEMTRACE_TRACK_FUNCTIONS 1
#endif


/** This structure is added before the memory returned to the client, and
 *  can be used to retain information. This is no longer necessary in a full
 *  memtrace report, but we leave it so that we can still do simple reports
 *  efficiently */
typedef struct {
    /** The total number of bytes in this allocation */
    size_t bytes;
} mem_trace_t;

static const char* s_program_name = NULL;

/** The total number of allocated bytes */
static intmax_t s_mem_allocated_bytes = 0;

static toolkit_mutex_t global_allocated_lock;

static toolkit_once_t once = TOOLKIT_ONCE_INIT;

/** The depth of the function calls. Does not depend on how many functions
 *  we have printed out. */
static intmax_t MemTrace_funcDepth = 0;

/** Stack of function pointers */
static void* MemTrace_funcStack[MEMTRACE_FUNCSTACK_SIZE];
static void* MemTrace_callStack[MEMTRACE_FUNCSTACK_SIZE];

/** Points to where the next function should be added in MemTrace_funcStack*/
static int s_mem_stack_index = 0;

TOOLKIT_API const char* tk_get_program_name()
{
	return s_program_name;
}

TOOLKIT_API void tk_set_program_name(const char* name)
{
	if (s_program_name != NULL) {
		char* just_progname = strrchr(name, SPLIT_SLASH);
		if (just_progname != NULL) s_program_name = toolkit_strdup(just_progname + 1);
		s_program_name = toolkit_strdup(name);
	}
}

//////////////////////////////////////////////////////////////////////////

static void init_allocated_mutex_once(void) 
{
    toolkit_mutex_init(&global_allocated_lock);
}

static void increment_allocated_bytes(size_t bytes)
{
	toolkit_once(&once, init_allocated_mutex_once);

    toolkit_mutex_lock(&global_allocated_lock);
    s_mem_allocated_bytes += bytes;
    toolkit_mutex_unlock(&global_allocated_lock);
}

static void decrement_allocated_bytes(size_t bytes)
{
    toolkit_once(&once, init_allocated_mutex_once);

	toolkit_mutex_lock(&global_allocated_lock);
	s_mem_allocated_bytes -= bytes;
	toolkit_mutex_unlock(&global_allocated_lock);
}

static void* mark_allocated_buffer(char* buf, size_t bytes)
{
    char* ret_buf = NULL;

    increment_allocated_bytes(bytes);
	((mem_trace_t*)buf)->bytes = bytes;

    ret_buf = (char*)(buf + MEMTRACE_HEADER_SIZE);

	do {
		/* Fill beginning markers */
		memset(buf + sizeof(mem_trace_t), MEMTRACE_MARK, MEMTRACE_GUARDSIZE);
		/* Fill end markers */
		memset(buf + MEMTRACE_HEADER_SIZE + bytes, MEMTRACE_MARK, MEMTRACE_GUARDSIZE);
	} while (0);

    return ret_buf;
}

void* toolkit_memtrace_zalloc(size_t bytes, const char* file, const char* function, int line)
{
	void* buf = toolkit_memtrace_malloc(bytes, file, function, line);
	if (buf) {
		memset(buf, 0, bytes);
	}
	return buf;
}

void* toolkit_memtrace_malloc(size_t bytes, const char* file, const char* function, int line)
{
    char* pcvoid = NULL;
    void* ret = NULL;

    bytes = bytes == 0 ? 1 : bytes;

    pcvoid = (char*)malloc(bytes + sizeof(mem_trace_t) + 2 * MEMTRACE_GUARDSIZE);
    if (pcvoid == NULL) {
        WLog_ERR(TAG, "malloc failed in <%s>[%d]{%s}", function, line, file);
        return NULL;
    }

    ret = mark_allocated_buffer(pcvoid, bytes);

	/* Overwrite memory to return with junk */
	memset(ret, MEMTRACE_JUNK, bytes);

    WLog_DBG(TAG, "malloc: %p+%u in <%s>[%d]{%s}", ret, (unsigned int)bytes, function, line, file);

    return ret;
}

void* toolkit_memtrace_calloc(size_t elem, size_t ele_size, const char* file, const char* function, int line)
{
	char* pcvoid = NULL;
	void* ret = NULL;
    size_t bytes = 0;

    bytes = elem * ele_size;
    pcvoid = (char*)calloc(1, bytes + sizeof(mem_trace_t) + 2 * MEMTRACE_GUARDSIZE);
	if (pcvoid == NULL) {
		WLog_ERR(TAG, "calloc failed in <%s>[%d]{%s}", function, line, file);
		return NULL;
	}

    ret = mark_allocated_buffer(pcvoid, bytes);

    WLog_DBG(TAG, "calloc: %p+%u in <%s>[%d]{%s}", ret, (unsigned int)bytes, function, line, file);

    return ret;
}

void* toolkit_memtrace_czalloc(size_t elem, size_t ele_size, const char* file, const char* function, int line)
{
	void* buf = toolkit_memtrace_calloc(elem, ele_size, file, function, line);
	if (buf) {
		memset(buf, 0, elem * ele_size);
	}
	return buf;
}

static int check_and_junk_buffer(char* pcvoid, size_t bytes)
{
    int ret = 0;
    int i;

	do {
		int err = 0;
		char* mark = pcvoid + sizeof(mem_trace_t);
		for (i = 0; i < MEMTRACE_GUARDSIZE; ++i) {
			if (*mark != MEMTRACE_MARK) {
				err = 1;
				break;
			}
			*mark = MEMTRACE_JUNK;
			mark++;
		}

		if (err) {
            ret = 1;
            break;
		}

		err = 0;
		mark += bytes;

		for (i = 0; i < MEMTRACE_GUARDSIZE; ++i) {
			if (*mark != MEMTRACE_MARK) {
				err = 1;
				break;
			}
			*mark = MEMTRACE_JUNK;
			mark++;
		}
		if (err) {
            ret = 2;
            break;
		}

	} while (0);

    return ret;
}

void  toolkit_memtrace_free(void* space, const char* file, const char* function, int line)
{
    char* pcvoid = NULL;
    size_t bytes = 0;
    int i = 0;
    int err = 0;

    if (space == NULL)
        return;

    pcvoid = (char*)space - MEMTRACE_HEADER_SIZE;
    bytes = ((mem_trace_t*)pcvoid)->bytes;

    decrement_allocated_bytes(bytes);

    err = check_and_junk_buffer(pcvoid, bytes);
    if (err == 1) {
		WLog_ERR(TAG, "A memory underflow was detected, address: %p in <%s>[%d]{%s}", space, function, line, file);
		abort();
    }
    else if (err == 2) {
		WLog_ERR(TAG, "A memory overflow was detected. address: %p in <%s>[%d]{%s}", space, function, line, file);
		abort();
    }

	/* Overwrite memory to return with junk */
	memset(space, MEMTRACE_JUNK, bytes);

    WLog_DBG(TAG, "free: %p in <%s>[%d]{%s}", space, function, line, file);

    /* Actually free the memory */
    free(pcvoid);

	WLog_WARN(TAG, "%u bytes still not free for now!", s_mem_allocated_bytes);

    return;
}

void* toolkit_memtrace_realloc(void* space, size_t bytes, const char* file, const char* function, int line)
{
    void* ret = NULL;
    char* pcvoid = NULL;
    size_t old_bytes = 0;
    int i;
    int err;

    if (bytes == 0) {
        toolkit_memtrace_free(space, file, function, line);
        return NULL;
    }

    if (space == NULL) {
        return toolkit_memtrace_malloc(bytes, file, function, line);
    }

	pcvoid = (char*)space - MEMTRACE_HEADER_SIZE;
	old_bytes = ((mem_trace_t*)pcvoid)->bytes;

	err = check_and_junk_buffer(pcvoid, old_bytes);
	if (err == 1) {
		WLog_ERR(TAG, "A memory underflow was detected, address: %p in <%s>[%d]{%s}", space, function, line, file);
		abort();
	}
	else if (err == 2) {
		WLog_ERR(TAG, "A memory overflow was detected. address: %p in <%s>[%d]{%s}", space, function, line, file);
		abort();
	}

    pcvoid = (char*)realloc(pcvoid, bytes + sizeof(mem_trace_t) + 2 * MEMTRACE_GUARDSIZE);

    if (pcvoid == NULL) {
        WLog_ERR(TAG, "realloc() ran out of memory. %u in <%s>[%d]{%s}", (unsigned int)bytes, function, line, file);
        return NULL;
    }

    ret = mark_allocated_buffer(pcvoid, bytes);
    decrement_allocated_bytes(old_bytes);

    /* If we allocated more memory, then we should fill in junk */
    if (old_bytes < bytes) {
        const size_t count = bytes - old_bytes;
        char* mark = pcvoid + MEMTRACE_HEADER_SIZE + old_bytes;
        memset(mark, MEMTRACE_JUNK, count);
    }

    WLog_DBG(TAG, "realloc: %p+%u from address: %p+%u in <%s>[%d]{%s}",
        ret, (unsigned int)bytes, space, (unsigned int)old_bytes, function, line, file);

    WLog_DBG(TAG, "free: %p in <%s>[%d]{%s}", space, function, line, file);
    return ret;
}

char* toolkit_memtrace_strdup(const char* string, const char* file, const char* function, int line)
{
	char* string_copy = NULL;

	/* Don't actually handle NULL string, since regular strdup() does not */
	TOOLKIT_ASSERT(string != NULL);

    //if (strlen(string) == 0) {
    //    abort(); /**debug: memory leak.*/
    //}

	string_copy = toolkit_memtrace_malloc(strlen(string) + 1, file, function, line);
	if (string_copy != NULL) {
		strcpy(string_copy, string);
	}
	return string_copy;
}