//	Darkstar Engine  (c) UnderWorld Interactive Studio (all coding by Svarog    2007)
//
//      Данный файл является частью Darkstar Engine. Любое использование кода файла
//      без согласия компании, является нарушением закона об авторском праве (в слу-
//      чае если не оговорено обратное).
//
//      Имя файла: memory_managment.cpp
//      Краткое описание: переопределение операторов new, delete, new[], delete[]
//
//   This file was coded by Zeux

#include <Darkstar/common.hpp>

#include <Darkstar/Kernel/Misc/memory_management.hpp>

#include <new>

#include <DbgHelp.h>

#pragma comment(lib, "dbghelp.lib")

#define CONFIG_DEBUG_MEMMGR 1

extern "C"
{
	NTSYSAPI USHORT NTAPI RtlCaptureStackBackTrace(ULONG FramesToSkip, ULONG FramesToCapture, PVOID *BackTrace, PULONG BackTraceHash);
};

namespace
{
	const unsigned int DEBUG_MEMMGR_STACKLIMIT = 40;

	const unsigned int DEBUG_MEMMGR_MARKER = 0xDEADBEEF;

	unsigned int g_frame_alloc_count = 0, g_frame_rtlalloc_count = 0;
	uint64_t g_alloc_count = 0, g_rtlalloc_count = 0;
	size_t g_alloc_size = 0;

	/// Memory layout is as follows:
	/// [stack size bytes for stack info] [node_t] [actual data]
	struct node_t
	{
		size_t stack_size;
		size_t size;

		node_t* next;
		node_t* prev;
	};

	node_t* g_alloc_list_begin = NULL;
	node_t* g_alloc_list_end = NULL;

	void* my_alloc(size_t size)
	{
#if CONFIG_DEBUG_MEMMGR
		{
			// Why do we need it? No idea. But stack capture does not work in Release mode without these.
			DWORD_PTR frame_pointer;
			__asm mov frame_pointer, ebp;

			void* stack_buf[DEBUG_MEMMGR_STACKLIMIT];

			size_t stack_size =  RtlCaptureStackBackTrace(0, DEBUG_MEMMGR_STACKLIMIT, stack_buf, 0);

			size_t rsize = size + sizeof(node_t) + stack_size * sizeof(void*) + 2 * sizeof(unsigned int);

			char* ptr = (char*)malloc(rsize) + sizeof(node_t) + stack_size * sizeof(void*);

			// fill node info
			node_t* node = (node_t*)ptr - 1;

			node->stack_size = stack_size;
			node->size = size;

			node->next = 0;
			node->prev = g_alloc_list_end;

			if (!g_alloc_list_begin) g_alloc_list_begin = node;
			else g_alloc_list_end->next = node;

			g_alloc_list_end = node;

			// fill stack info
			void** stack = (void**)node - stack_size;

			memcpy(stack, stack_buf, stack_size * sizeof(void*));

			++g_frame_alloc_count;
			++g_alloc_count;
			g_alloc_size += size;

			// insert markers
			ptr += sizeof(unsigned int);

			*(unsigned int*)(ptr - sizeof(unsigned int)) = DEBUG_MEMMGR_MARKER;
			*(unsigned int*)(ptr + size) = DEBUG_MEMMGR_MARKER;

			return ptr;
		}
#else
		return malloc(size);
#endif
	}

	void my_free(void* ptr)
	{
#if CONFIG_DEBUG_MEMMGR
		{
			char* rptr = (char*)ptr;

			node_t* node = (node_t*)(rptr - sizeof(unsigned int)) - 1;

			assert(*(unsigned int*)(rptr - sizeof(unsigned int)) == DEBUG_MEMMGR_MARKER && "Memory buffer underflow");
			assert(*(unsigned int*)(rptr + node->size) == DEBUG_MEMMGR_MARKER && "Memory buffer overflow");

			if (node->prev) node->prev->next = node->next;
			else g_alloc_list_begin = node->next;

			if (node->next) node->next->prev = node->prev;
			else g_alloc_list_end = node->prev;

			g_alloc_size -= node->size;

			free(rptr - sizeof(unsigned int) - sizeof(node_t) - node->stack_size * sizeof(void*));
		}
#else
		free(ptr);
#endif
	}
};

void* operator new(size_t size)
{
	void* ptr = my_alloc(size);

	if (!ptr) throw std::bad_alloc();

	return ptr;
}

void* operator new[](size_t size)
{
	void* ptr = my_alloc(size);

	if (!ptr) throw std::bad_alloc();

	return ptr;
}

void* operator new(size_t size, const std::nothrow_t&) throw()
{
	return my_alloc(size);
}

void* operator new[](size_t size, const std::nothrow_t&) throw()
{
	return my_alloc(size);
}

void operator delete(void* ptr)
{
	if (ptr) my_free(ptr);
}

void operator delete[](void* ptr)
{
	if (ptr) my_free(ptr);
}

void operator delete(void* ptr, const std::nothrow_t&) throw()
{
	if (ptr) my_free(ptr);
}

void operator delete[](void* ptr, const std::nothrow_t&) throw()
{
	if (ptr) my_free(ptr);
}

namespace ker
{
	void MemoryManager::outputLeaks(FILE* file)
	{
#if CONFIG_DEBUG_MEMMGR
		HANDLE process = GetCurrentProcess();

		bool symbols = true;

		SymSetOptions(SYMOPT_LOAD_LINES | SYMOPT_DEFERRED_LOADS | SYMOPT_UNDNAME);
		if (!SymInitialize(process, NULL, TRUE))
		{
			fprintf(file, "Error getting symbol information\n");
			symbols = false;
		}

		node_t* node = g_alloc_list_begin;

		size_t size = 0;

		fprintf(file, "Memory leaks summary:\n");

		while (node)
		{
			fprintf(file, "Memory leak: size = %d, pointer: %p\n", node->size, node + 1);

			void** stack = (void**)node - node->stack_size;

			for (size_t i = 0; i < node->stack_size; ++i)
			{
				DWORD64 w = reinterpret_cast<DWORD64>(stack[i]);

				fprintf(file, "0x%x", w);

				if (symbols)
				{
					const size_t MAXSYMBOLNAMELENGTH = 512;
					const size_t SYMBOLBUFFERSIZE = sizeof(SYMBOL_INFO) + MAXSYMBOLNAMELENGTH - 1;

					IMAGEHLP_LINE64 source_info;

					unsigned char symbol_buffer[SYMBOLBUFFERSIZE];
					SYMBOL_INFO *function_info = (SYMBOL_INFO*)symbol_buffer;

					DWORD displacement;
					DWORD64 displacement64;

					// Initialize structures passed to the symbol handler.
					memset(function_info, 0, SYMBOLBUFFERSIZE);
					function_info->SizeOfStruct = sizeof(SYMBOL_INFO);
					function_info->MaxNameLen = MAXSYMBOLNAMELENGTH;
					memset(&source_info, 0, sizeof(IMAGEHLP_LINE64));
					source_info.SizeOfStruct = sizeof(IMAGEHLP_LINE64);

					if (SymGetLineFromAddr64(process, w, &displacement, &source_info))
						fprintf(file, " %s(%d):", source_info.FileName, source_info.LineNumber);

					if (SymFromAddr(process, w, &displacement64, function_info))
						fprintf(file, " %s", function_info->Name);
				}

				fprintf(file, "\n");
			}

			fprintf(file, "\n");

			size += node->size;

			node = node->next;
		}

		SymCleanup(process);

		if (size > 0) fprintf(file, "Total memory leaked: %d\n", size);
#endif
	}

	void MemoryManager::getFrameStatistics(unsigned int& count, unsigned int& total_count)
	{
		count = g_frame_alloc_count;
		total_count = g_frame_rtlalloc_count;
	}

	void MemoryManager::getGlobalStatistics(uint64_t& count, size_t& size, uint64_t& total_count)
	{
		count = g_alloc_count;
		size = g_alloc_size;
		total_count = g_rtlalloc_count;
	}

	void MemoryManager::resetFrameStatistics()
	{
		g_frame_alloc_count = 0;
		g_frame_rtlalloc_count = 0;
	}
};

