#include <stdlib.h>
#include <stdio.h>
#include <stdatomic.h>
#include <windows.h>
#include "CounterOffsets.h"

#define PERFORMANCE_KEY "SYSTEM\\CurrentControlSet\\Services\\Millistream Database Application\\Performance"
#define MAP_NAME "Global\\Millistream Database Application Perfcounters"

#define NELEMS(x) (sizeof x/sizeof x[0])

struct Queuestats {
	unsigned __int64 items;
	unsigned __int64 empty;
	unsigned __int64 updates;
	unsigned __int64 dedup;
	unsigned __int64 failed;
};

struct statistics {
	_Atomic unsigned __int64 receivedmessages;
	_Atomic unsigned __int64 connected_to_feed;
	struct Queuestats queuestat[4];
	unsigned __int64 file_templates_ver;
	unsigned __int64 feed_templates_ver;
};

DWORD g_OpenCount = 0;
DWORD FirstCounterIndex = 0;
DWORD FirstHelpIndex = 0;
HANDLE hMap;
struct statistics *pMap;

__declspec(dllexport) DWORD APIENTRY OpenPerformanceData (LPWSTR __attribute__ ((unused))pContext)
{
	LONG rc = ERROR_SUCCESS;

	if (g_OpenCount == 0) {
		HKEY hkey;

		rc = RegOpenKeyEx (HKEY_LOCAL_MACHINE, PERFORMANCE_KEY, 0, KEY_READ, &hkey);

		if (rc == ERROR_SUCCESS) {
			DWORD DataSize = sizeof FirstCounterIndex;
			rc = RegQueryValueEx (hkey, "First Counter", NULL, NULL, (LPBYTE) &FirstCounterIndex, &DataSize);

			if (rc == ERROR_SUCCESS) {
				DataSize = sizeof FirstHelpIndex;
				rc = RegQueryValueEx (hkey, "First Help", NULL, NULL, (LPBYTE) &FirstHelpIndex, &DataSize);
			}
		}

		RegCloseKey (hkey);

		if (rc == ERROR_SUCCESS) {
			hMap = CreateFileMapping (INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, sizeof *pMap, MAP_NAME);

			if (hMap == NULL) {
				rc = GetLastError ();
				return rc;
			}

			pMap = MapViewOfFile (hMap, FILE_MAP_READ, 0, 0, 0);

			if (pMap == NULL) {
				rc = GetLastError ();
				CloseHandle (hMap);
				return rc;
			}
		}
	}

	if (rc != ERROR_SUCCESS)
		return rc;

	g_OpenCount++;
	return ERROR_SUCCESS;
}

__declspec(dllexport) DWORD APIENTRY CollectPerformanceData (LPWSTR pQuery, LPVOID *ppData, LPDWORD pcbData, LPDWORD pObjectsReturned)
{
	PBYTE pPosition;
	PERF_OBJECT_TYPE *pHeader;
	PERF_COUNTER_DEFINITION *pCounterDef;
	PERF_COUNTER_BLOCK *pCounterBlock;
	DWORD needed_length;
	LPWSTR pCopy;
	ULONGLONG counter_value[26];
	int counter_use[26] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

	static const int counter_offsets[26] = {
		RECEIVED_MESSAGES,
		RECEIVED_MESSAGES_SEC,
		STRANGLED,
		STRANGLED_PERCENT,
		STRANGLED_SEC,
		DATABASEITEMS,
		DATABASEITEMS_SEC,
		DATABASEQUEUE,
		FAILEDDATABASEUPDATES,
		CONNECTED,
		Q1_QUEUE,
		Q1_UPDATES,
		Q1_DEDUP,
		Q1_EMPTY,
		Q2_QUEUE,
		Q2_UPDATES,
		Q2_DEDUP,
		Q2_EMPTY,
		Q3_QUEUE,
		Q3_UPDATES,
		Q3_DEDUP,
		Q3_EMPTY,
		Q4_QUEUE,
		Q4_UPDATES,
		Q4_DEDUP,
		Q4_EMPTY
	};

	static const DWORD counter_types[26] = {
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_BULK_COUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_BULK_COUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_BULK_COUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT,
		PERF_COUNTER_LARGE_RAWCOUNT
	};
	
	int num_counters = 0;
	int i;

	*pObjectsReturned = 0;

	if (g_OpenCount == 0) {
		*pcbData = 0;
		return ERROR_SUCCESS;
	}

	pCopy = _wcsdup (pQuery);
	pCopy = _wcslwr (pCopy);

	if (wcsstr (pCopy, L"global")) {
		for (i = 0; i < (int)(NELEMS (counter_use)); i++)
			counter_use[i] = 1;

		num_counters = NELEMS (counter_use);
	} else {
		LPWSTR p = pCopy;

		do {
			DWORD offset = _wtoi (p);

			if (offset == FirstCounterIndex) {
				for (i = 0; i < (int)(NELEMS (counter_use)); i++)
					counter_use[i] = 1;

				num_counters = NELEMS (counter_use);
				break;
			}

			for (i = 0; i < (int)(NELEMS (counter_use)); i++) {
				if (counter_offsets[i] + FirstCounterIndex == offset) {
					if (counter_use[i] == 0) {
						counter_use[i] = 1;
						num_counters++;
					}
				}
			}

			p = wcschr (p + 1, L' ');
		} while (p != NULL);
	}

	free (pCopy);

	if (num_counters == 0) {
		*pcbData = 0;
		return ERROR_SUCCESS;
	}

	needed_length = sizeof *pHeader + sizeof *pCounterDef * num_counters + sizeof *pCounterBlock + sizeof (ULONGLONG) * num_counters + 7;

	if (needed_length > *pcbData) {
		*pcbData = 0;
		return ERROR_MORE_DATA;
	}

	/* Fetch counter values from the MDA process */
	counter_value[0] = pMap->receivedmessages;
	counter_value[1] = pMap->receivedmessages;
	counter_value[9] = pMap->connected_to_feed;

	counter_value[10] = pMap->queuestat[0].items;
	counter_value[11] = pMap->queuestat[0].updates;
	counter_value[12] = pMap->queuestat[0].dedup;
	counter_value[13] = pMap->queuestat[0].empty;
		counter_value[2] = counter_value[12];
		counter_value[5] = counter_value[11];
		counter_value[7] = counter_value[10];
		counter_value[8] = pMap->queuestat[0].failed;

	counter_value[14] = pMap->queuestat[1].items;
	counter_value[15] = pMap->queuestat[1].updates;
	counter_value[16] = pMap->queuestat[1].dedup;
	counter_value[17] = pMap->queuestat[1].empty;
		counter_value[2] += counter_value[16];
		counter_value[5] += counter_value[15];
		counter_value[7] += counter_value[14];
		counter_value[8] += pMap->queuestat[1].failed;

	counter_value[18] = pMap->queuestat[2].items;
	counter_value[19] = pMap->queuestat[2].updates;
	counter_value[20] = pMap->queuestat[2].dedup;
	counter_value[21] = pMap->queuestat[2].empty;
		counter_value[2] += counter_value[20];
		counter_value[5] += counter_value[19];
		counter_value[7] += counter_value[18];
		counter_value[8] += pMap->queuestat[2].failed;

	counter_value[22] = pMap->queuestat[3].items;
	counter_value[23] = pMap->queuestat[3].updates;
	counter_value[24] = pMap->queuestat[3].dedup;
	counter_value[25] = pMap->queuestat[3].empty;
		counter_value[2] += counter_value[24];
		counter_value[5] += counter_value[23];
		counter_value[7] += counter_value[22];
		counter_value[8] += pMap->queuestat[3].failed;

	counter_value[4] = counter_value[2];
	counter_value[6] = counter_value[5];

	/* calculate STRANGLED_PERCENT */
	if (counter_value[0] != 0 && counter_value[2] != 0)
		counter_value[3] = counter_value[2] * 100 / counter_value[0];
	else
		counter_value[3] = 0;

	pPosition = (PBYTE) *ppData;

	pHeader = (struct _PERF_OBJECT_TYPE *) pPosition;
	pPosition += sizeof *pHeader;

	pCounterDef = (struct _PERF_COUNTER_DEFINITION *) pPosition;
	pPosition += sizeof *pCounterDef * num_counters;

	pCounterBlock = (struct _PERF_COUNTER_BLOCK *) pPosition;
	pCounterBlock->ByteLength = sizeof *pCounterBlock;
	pPosition += sizeof *pCounterBlock;

	if (((DWORD_PTR)pPosition % 8) != 0) {
		DWORD_PTR offset = 8 - (((DWORD_PTR)pPosition) % 8);

		pPosition += offset;
		pCounterBlock->ByteLength += offset;
		needed_length -= 7;
		needed_length += offset;
	}

	*pObjectsReturned = 1;
	*pcbData = needed_length;

	memset (pHeader, 0, sizeof *pHeader);
	memset (pCounterDef, 0, sizeof *pCounterDef);

	pHeader->HeaderLength = sizeof *pHeader;
	pHeader->DefinitionLength = pHeader->HeaderLength + sizeof *pCounterDef * num_counters;
	pHeader->TotalByteLength = needed_length;
	pHeader->ObjectNameTitleIndex = FirstCounterIndex;
	pHeader->ObjectHelpTitleIndex = FirstHelpIndex;
	pHeader->DetailLevel = PERF_DETAIL_ADVANCED;
	pHeader->NumCounters = num_counters;
	pHeader->DefaultCounter = -1;
	pHeader->NumInstances = PERF_NO_INSTANCES;
	pHeader->CodePage = 0;

	QueryPerformanceCounter (&pHeader->PerfTime);
	QueryPerformanceFrequency (&pHeader->PerfFreq);

	for (i = 0; i < (int)(NELEMS (counter_use)); i++) {
		if (counter_use[i] == 0)
			continue;

		pCounterDef->ByteLength = sizeof *pCounterDef;
		pCounterDef->CounterNameTitleIndex = FirstCounterIndex + counter_offsets[i];
		pCounterDef->CounterHelpTitleIndex = FirstHelpIndex + counter_offsets[i];
		pCounterDef->DefaultScale = 0;
		pCounterDef->DetailLevel = PERF_DETAIL_ADVANCED;
		pCounterDef->CounterType = counter_types[i];
		pCounterDef->CounterSize = sizeof (ULONGLONG);
		pCounterDef->CounterOffset = pCounterBlock->ByteLength;
		pCounterDef++;

		*((ULONGLONG *)pPosition) = counter_value[i];
		pPosition += sizeof (ULONGLONG);
		pCounterBlock->ByteLength += sizeof (ULONGLONG);
	}

	*ppData = pPosition;

	return ERROR_SUCCESS;
}

__declspec(dllexport) DWORD APIENTRY ClosePerformanceData ()
{
	g_OpenCount--;

	if (g_OpenCount == 0) {
		UnmapViewOfFile (pMap);
		CloseHandle (hMap);
	}

	return ERROR_SUCCESS;
}
