#include <windows.h>
#include <dbghelp.h>
#include <vector>
#include <string>
#include <sstream>
#include <iterator>
#include <algorithm>
#include <iostream>

struct ThreadParams
{
	HANDLE hExceptionThread;
	LPEXCEPTION_POINTERS pException;
};

void ContextToStackFrame(const CONTEXT& ctx, STACKFRAME64& sf)
{
	sf.AddrReturn.Mode = sf.AddrFrame.Mode = sf.AddrPC.Mode = sf.AddrStack.Mode = sf.AddrBStore.Mode = AddrModeFlat;
	sf.Virtual = TRUE;

#ifndef _WIN64

	sf.AddrPC.Offset = sf.AddrReturn.Offset = ctx.Eip;
	sf.AddrStack.Offset = ctx.Esp;
	sf.AddrFrame.Offset = ctx.Ebp;
	sf.AddrBStore = sf.AddrFrame;

#elif (defined(_IA64_) || defined(_M_IA64_))

	sf.AddrPC.Offset = sf.AddrReturn.Offset = ctx.StIIP;
	sf.AddrStack.Offset = ctx.IntSp;
	sf.AddrBStore.Offset = ctx.RsBSP;
	sf.AddrFrame = sf.AddrBStore;
		
#elif (defined(_AMD64_) || defined(_M_AMD64))

	sf.AddrPC.Offset = sf.AddrReturn.Offset = ctx.Rip;
	sf.AddrFrame.Offset = ctx.Rbp;
	sf.AddrStack.Offset = ctx.Rsp;
	sf.AddrBStore = sf.AddrFrame;

#elif
#error "Unknown Platform for context structure"
#endif
}

bool AddressToSymbol(HANDLE hProc, DWORD64 address, std::string& name)
{
	bool success = false;
	// initialize the output buffer with the address
	std::ostringstream outputBuffer;
	outputBuffer << reinterpret_cast<PVOID>(address) << " : ";

	// setup the variables and structures required
	BYTE arr[sizeof(SYMBOL_INFO) + (MAX_SYM_NAME - 1)] = {0};
	PSYMBOL_INFO inf = reinterpret_cast<PSYMBOL_INFO>(arr);
	inf->SizeOfStruct = sizeof(SYMBOL_INFO);
	inf->MaxNameLen = MAX_SYM_NAME;
	
	DWORD64 disp = 0;
	IMAGEHLP_MODULE64 modInf = {0};
	// the IMAGEHLP_MODULE64 structure has changed sizes 
	// during the various releases of dbghelp.dll, 
	// the one shipped with 2000/XP (v5.1)
	// is incredibly old and is of smaller size compared to the one in the current
	// MS headers. This, combined with a bug in SymGetModuleInfo64
	// where the comparison on the SizeOfStruct member field is inverted,
	// causes the function to fail if SizeOfStruct is greater than 0x248 
	// (the sizeof(IMAGEHLP_MODULE64) when it was built) 
	// but carry on successfully if it is less than that size.
	// I'm using the newer Windows 7 headers where the sizeof
	// is greater than that so I need this hack to test it,
	// however you may not.
	modInf.SizeOfStruct = 0x248;
	// modInf.SizeOfStruct = sizeof(IMAGEHLP_MODULE64);

	if(SymGetModuleInfo64(hProc, address, &modInf))
	{
		outputBuffer << modInf.ModuleName;
		if(SymFromAddr(hProc, address, &disp, inf))
		{
			outputBuffer << "!" << inf->Name;
			if(disp)
			{
				outputBuffer << "+0x" << std::hex << disp;
			}
		}
		else
		{
			if(modInf.BaseOfImage)
			{
				outputBuffer << "+0x" << std::hex << address - modInf.BaseOfImage;
			}
		}
		name = outputBuffer.str();
		success = true;
	}
	return success;
}

DWORD WINAPI TraceFunc(LPVOID lpParams)
{
	ThreadParams* tp = reinterpret_cast<ThreadParams*>(lpParams);
	//MessageBoxA(NULL, "", "", NULL);
    HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, GetCurrentProcessId());
	if(hProc)
	{
		//DWORD opts = SymGetOptions();
		SymSetOptions(SYMOPT_UNDNAME | SYMOPT_CASE_INSENSITIVE | SYMOPT_LOAD_LINES | 
									SYMOPT_OMAP_FIND_NEAREST | SYMOPT_NO_PROMPTS | 
									SYMOPT_DEBUG);
		if(SymInitialize(hProc, NULL, TRUE))
		{
			std::vector<std::string> frames;
			// make a copy of th context since it'll be modified by StackWalk64
			CONTEXT ctx = *(tp->pException->ContextRecord);
			STACKFRAME64 frame = {0};
			// initialize the stackframe from the context
			ContextToStackFrame(ctx, frame);
			// get the machine type from the exe header
			PIMAGE_NT_HEADERS pHeaders = ImageNtHeader(GetModuleHandle(NULL));
			WORD machType = pHeaders->FileHeader.Machine;
			// do the hustle
			while(StackWalk64(machType, hProc, GetCurrentThread(), &frame, &ctx, NULL,
					&SymFunctionTableAccess64, &SymGetModuleBase64, NULL))
			{
				std::string symName;
				if(AddressToSymbol(hProc, frame.AddrPC.Offset, symName))
				{
					frames.push_back(symName);
				}
			}
			// display the trace
			EXCEPTION_RECORD& exp = *(tp->pException->ExceptionRecord);
			std::cout << "Stack Trace for exception 0x" << std::hex << exp.ExceptionCode << " at address 0x" << exp.ExceptionAddress << '\n';
			std::copy(frames.begin(), frames.end(), std::ostream_iterator<std::string>(std::cout, "\n"));
		}
		else std::cout << "Couldn't get stacktrace\n";
		CloseHandle(hProc);
	}
	else std::cout << "Couldn't open process handle. Error: " << GetLastError() << '\n';
    return 0;
}

DWORD slot = TlsAlloc();

LONG WINAPI ExceptionFilter(LPEXCEPTION_POINTERS pException)
{
	// This function is called in the context of the thread which 
	// caused the exception so it could be in any state, we spawn
	// off another thread to do the heavy lifting, to mitigate this 
	// as far as possible. Although, who knows what the state of the
	// process is, so 
	HANDLE hThisThread = OpenThread(THREAD_ALL_ACCESS, FALSE, GetCurrentThreadId());
	if(hThisThread)
	{
		ThreadParams tp = {hThisThread, pException};
		HANDLE hWalkThread = CreateThread(NULL, 0, &TraceFunc, &tp, 0, NULL);
		if(hWalkThread)
		{
			WaitForSingleObject(hWalkThread, INFINITE);
			CloseHandle(hWalkThread);
		}
		CloseHandle(hThisThread);
	}
	return EXCEPTION_CONTINUE_SEARCH;
}