//========= Copyright Valve Corporation, All rights reserved. ============//
//
// Purpose: 
//
//=============================================================================//

#include <windows.h>
#include <dbghelp.h>
#include "vmpi.h"
#include "cmdlib.h"
#include "vmpi_tools_shared.h"
#include "tier1/strtools.h"
#include "mpi_stats.h"
#include "iphelpers.h"
#include "tier0/minidump.h"


// ----------------------------------------------------------------------------- //
// Globals.
// ----------------------------------------------------------------------------- //

static bool g_bReceivedDirectoryInfo = false;	// Have we gotten the qdir info yet?

static bool g_bReceivedDBInfo = false;
static CDBInfo g_DBInfo;
static unsigned long g_JobPrimaryID;

static int g_nDisconnects = 0;	// Tracks how many remote processes have disconnected ungracefully.


// ----------------------------------------------------------------------------- //
// Shared dispatch code.
// ----------------------------------------------------------------------------- //

bool SharedDispatch(MessageBuffer *pBuf, int iSource, int iPacketID)
{
	char *pInPos = &pBuf->data[2];

	switch (pBuf->data[1])
	{
	case VMPI_SUBPACKETID_DIRECTORIES:
	{
		Q_strncpy(gamedir, pInPos, sizeof(gamedir));
		pInPos += strlen(pInPos) + 1;

		Q_strncpy(qdir, pInPos, sizeof(qdir));

		g_bReceivedDirectoryInfo = true;
	}
	return true;

	case VMPI_SUBPACKETID_DBINFO:
	{
		g_DBInfo = *((CDBInfo*)pInPos);
		pInPos += sizeof(CDBInfo);
		g_JobPrimaryID = *((unsigned long*)pInPos);

		g_bReceivedDBInfo = true;
	}
	return true;

	case VMPI_SUBPACKETID_CRASH:
	{
		char const chCrashInfoType = *pInPos;
		pInPos += 2;
		switch (chCrashInfoType)
		{
		case 't':
			Warning("\nWorker '%s' dead: %s\n", VMPI_GetMachineName(iSource), pInPos);
			break;
		case 'f':
		{
			int iFileSize = *reinterpret_cast< int const * >(pInPos);
			pInPos += sizeof(iFileSize);

			// Temp folder
			char const *szFolder = NULL;
			if (!szFolder) szFolder = getenv("TEMP");
			if (!szFolder) szFolder = getenv("TMP");
			if (!szFolder) szFolder = "c:";

			// Base module name
			char chModuleName[_MAX_PATH], *pModuleName = chModuleName;
			::GetModuleFileName(NULL, chModuleName, sizeof(chModuleName) / sizeof(chModuleName[0]));

			if (char *pch = strrchr(chModuleName, '.'))
				*pch = 0;
			if (char *pch = strrchr(chModuleName, '\\'))
				*pch = 0, pModuleName = pch + 1;

			// Current time
			time_t currTime = ::time(NULL);
			struct tm * pTime = ::localtime(&currTime);

			// Number of minidumps this run
			static int s_numMiniDumps = 0;
			++s_numMiniDumps;

			// Prepare the filename
			char chSaveFileName[2 * _MAX_PATH] = { 0 };
			sprintf(chSaveFileName, "%s\\vmpi_%s_on_%s_%d%.2d%2d%.2d%.2d%.2d_%d.mdmp",
				szFolder,
				pModuleName,
				VMPI_GetMachineName(iSource),
				pTime->tm_year + 1900,	/* Year less 2000 */
				pTime->tm_mon + 1,		/* month (0 - 11 : 0 = January) */
				pTime->tm_mday,			/* day of month (1 - 31) */
				pTime->tm_hour,			/* hour (0 - 23) */
				pTime->tm_min,		    /* minutes (0 - 59) */
				pTime->tm_sec,		    /* seconds (0 - 59) */
				s_numMiniDumps
				);

			if (FILE *fDump = fopen(chSaveFileName, "wb"))
			{
				fwrite(pInPos, 1, iFileSize, fDump);
				fclose(fDump);

				Warning("\nSaved worker crash minidump '%s', size %d byte(s).\n",
					chSaveFileName, iFileSize);
			}
			else
			{
				Warning("\nReceived worker crash minidump size %d byte(s), failed to save.\n", iFileSize);
			}
		}
		break;
		}
	}
	return true;
	}

	return false;
}

CDispatchReg g_SharedDispatchReg(VMPI_SHARED_PACKET_ID, SharedDispatch);



// ----------------------------------------------------------------------------- //
// Module interfaces.
// ----------------------------------------------------------------------------- //

void SendQDirInfo()
{
	char cPacketID[2] = { VMPI_SHARED_PACKET_ID, VMPI_SUBPACKETID_DIRECTORIES };

	MessageBuffer mb;
	mb.write(cPacketID, 2);
	mb.write(gamedir, strlen(gamedir) + 1);
	mb.write(qdir, strlen(qdir) + 1);

	VMPI_SendData(mb.data, mb.getLen(), VMPI_PERSISTENT);
}


void RecvQDirInfo()
{
	while (!g_bReceivedDirectoryInfo)
		VMPI_DispatchNextMessage();
}


void SendDBInfo(const CDBInfo *pInfo, unsigned long jobPrimaryID)
{
	char cPacketInfo[2] = { VMPI_SHARED_PACKET_ID, VMPI_SUBPACKETID_DBINFO };
	const void *pChunks[] = { cPacketInfo, pInfo, &jobPrimaryID };
	int chunkLengths[] = { 2, sizeof(CDBInfo), sizeof(jobPrimaryID) };

	VMPI_SendChunks(pChunks, chunkLengths, ARRAYSIZE(pChunks), VMPI_PERSISTENT);
}


void RecvDBInfo(CDBInfo *pInfo, unsigned long *pJobPrimaryID)
{
	while (!g_bReceivedDBInfo)
		VMPI_DispatchNextMessage();

	*pInfo = g_DBInfo;
	*pJobPrimaryID = g_JobPrimaryID;
}

// If the file is successfully opened, read and sent returns the size of the file in bytes
// otherwise returns 0 and nothing is sent
int VMPI_SendFileChunk(const void *pvChunkPrefix, int lenPrefix, tchar const *ptchFileName)
{
	HANDLE hFile = NULL;
	HANDLE hMapping = NULL;
	void const *pvMappedData = NULL;
	int iResult = 0;

	hFile = ::CreateFile(ptchFileName, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (!hFile || (hFile == INVALID_HANDLE_VALUE))
		goto done;

	hMapping = ::CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
	if (!hMapping || (hMapping == INVALID_HANDLE_VALUE))
		goto done;

	pvMappedData = ::MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);
	if (!pvMappedData)
		goto done;

	int iMappedFileSize = ::GetFileSize(hFile, NULL);
	if (INVALID_FILE_SIZE == iMappedFileSize)
		goto done;

	// Send the data over VMPI
	if (VMPI_Send3Chunks(
		pvChunkPrefix, lenPrefix,
		&iMappedFileSize, sizeof(iMappedFileSize),
		pvMappedData, iMappedFileSize,
		VMPI_MASTER_ID))
		iResult = iMappedFileSize;

	// Fall-through for cleanup code to execute
done:
	if (pvMappedData)
		::UnmapViewOfFile(pvMappedData);

	if (hMapping && (hMapping != INVALID_HANDLE_VALUE))
		::CloseHandle(hMapping);

	if (hFile && (hFile != INVALID_HANDLE_VALUE))
		::CloseHandle(hFile);

	return iResult;
}

void VMPI_HandleCrash(const char *pMessage, void *pvExceptionInfo, bool bAssert)
{
	static LONG crashHandlerCount = 0;
	if (InterlockedIncrement(&crashHandlerCount) == 1)
	{
		Msg("\nFAILURE: '%s' (assert: %d)\n", pMessage, bAssert);

		// Send a message to the master.
		char crashMsg[4] = { VMPI_SHARED_PACKET_ID, VMPI_SUBPACKETID_CRASH, 't', ':' };

		VMPI_Send2Chunks(
			crashMsg,
			sizeof(crashMsg),
			pMessage,
			strlen(pMessage) + 1,
			VMPI_MASTER_ID);

		// Now attempt to create a minidump with the given exception information
		if (pvExceptionInfo)
		{
			struct _EXCEPTION_POINTERS *pvExPointers = (struct _EXCEPTION_POINTERS *) pvExceptionInfo;
			tchar tchMinidumpFileName[_MAX_PATH] = { 0 };
			bool bSucceededWritingMinidump = WriteMiniDumpUsingExceptionInfo(
				pvExPointers->ExceptionRecord->ExceptionCode,
				pvExPointers,
				(MINIDUMP_TYPE)(MiniDumpWithDataSegs | MiniDumpWithIndirectlyReferencedMemory | MiniDumpWithProcessThreadData),
				// ( MINIDUMP_TYPE )( MiniDumpWithDataSegs | MiniDumpWithFullMemory | MiniDumpWithHandleData | MiniDumpWithUnloadedModules | MiniDumpWithIndirectlyReferencedMemory | MiniDumpWithProcessThreadData | MiniDumpWithPrivateReadWriteMemory  ),
				// ( MINIDUMP_TYPE )( MiniDumpNormal ),
				tchMinidumpFileName);
			if (bSucceededWritingMinidump)
			{
				crashMsg[2] = 'f';
				VMPI_SendFileChunk(crashMsg, sizeof(crashMsg), tchMinidumpFileName);
				::DeleteFile(tchMinidumpFileName);
			}
		}

		// Let the messages go out.
		Sleep(500);
	}

	InterlockedDecrement(&crashHandlerCount);
}


// This is called if we crash inside our crash handler. It just terminates the process immediately.
LONG __stdcall VMPI_SecondExceptionFilter(struct _EXCEPTION_POINTERS *ExceptionInfo)
{
	TerminateProcess(GetCurrentProcess(), 2);
	return EXCEPTION_EXECUTE_HANDLER; // (never gets here anyway)
}


void VMPI_ExceptionFilter(unsigned long uCode, void *pvExceptionInfo)
{
	// This is called if we crash inside our crash handler. It just terminates the process immediately.
	SetUnhandledExceptionFilter(VMPI_SecondExceptionFilter);

	//DWORD code = ExceptionInfo->ExceptionRecord->ExceptionCode;

#define ERR_RECORD( name ) { name, #name }
	struct
	{
		int code;
		char *pReason;
	} errors[] =
	{
		ERR_RECORD(EXCEPTION_ACCESS_VIOLATION),
		ERR_RECORD(EXCEPTION_ARRAY_BOUNDS_EXCEEDED),
		ERR_RECORD(EXCEPTION_BREAKPOINT),
		ERR_RECORD(EXCEPTION_DATATYPE_MISALIGNMENT),
		ERR_RECORD(EXCEPTION_FLT_DENORMAL_OPERAND),
		ERR_RECORD(EXCEPTION_FLT_DIVIDE_BY_ZERO),
		ERR_RECORD(EXCEPTION_FLT_INEXACT_RESULT),
		ERR_RECORD(EXCEPTION_FLT_INVALID_OPERATION),
		ERR_RECORD(EXCEPTION_FLT_OVERFLOW),
		ERR_RECORD(EXCEPTION_FLT_STACK_CHECK),
		ERR_RECORD(EXCEPTION_FLT_UNDERFLOW),
		ERR_RECORD(EXCEPTION_ILLEGAL_INSTRUCTION),
		ERR_RECORD(EXCEPTION_IN_PAGE_ERROR),
		ERR_RECORD(EXCEPTION_INT_DIVIDE_BY_ZERO),
		ERR_RECORD(EXCEPTION_INT_OVERFLOW),
		ERR_RECORD(EXCEPTION_INVALID_DISPOSITION),
		ERR_RECORD(EXCEPTION_NONCONTINUABLE_EXCEPTION),
		ERR_RECORD(EXCEPTION_PRIV_INSTRUCTION),
		ERR_RECORD(EXCEPTION_SINGLE_STEP),
		ERR_RECORD(EXCEPTION_STACK_OVERFLOW),
		ERR_RECORD(EXCEPTION_ACCESS_VIOLATION),
	};

	int nErrors = sizeof(errors) / sizeof(errors[0]);
	int i = 0;
	char *pchReason = NULL;
	char chUnknownBuffer[32];
	for (i; (i < nErrors) && !pchReason; i++)
	{
		if (errors[i].code == uCode)
			pchReason = errors[i].pReason;
	}

	if (i == nErrors)
	{
		sprintf(chUnknownBuffer, "Error code 0x%08X", uCode);
		pchReason = chUnknownBuffer;
	}

	VMPI_HandleCrash(pchReason, pvExceptionInfo, true);

	TerminateProcess(GetCurrentProcess(), 1);
}


void HandleMPIDisconnect(int procID, const char *pReason)
{
	int nLiveWorkers = VMPI_GetCurrentNumberOfConnections() - g_nDisconnects - 1;

	// We ran into the size limit before and it wasn't readily apparent that the size limit had
	// been breached, so make sure to show errors about invalid packet sizes..
	bool bOldSuppress = g_bSuppressPrintfOutput;
	g_bSuppressPrintfOutput = (Q_stristr(pReason, "invalid packet size") == 0);

	Warning("\n\n--- WARNING: lost connection to '%s' (%s).\n", VMPI_GetMachineName(procID), pReason);

	if (g_bMPIMaster)
	{
		Warning("%d workers remain.\n\n", nLiveWorkers);

		++g_nDisconnects;
		/*
		if ( VMPI_GetCurrentNumberOfConnections() - g_nDisconnects <= 1 )
		{
		Error( "All machines disconnected!" );
		}
		*/
	}
	else
	{
		VMPI_HandleAutoRestart();
		Error("Worker quitting.");
	}

	g_bSuppressPrintfOutput = bOldSuppress;
}