#include <windows.h>
#include <psapi.h>

#include "TSHooks.hpp"

static ADDR ImageBase;
static ADDR ImageSize;

static void InitScanner()
{
	HMODULE module = GetModuleHandle(NULL);
	if (module) {
		MODULEINFO info;
		GetModuleInformation(GetCurrentProcess(), module, &info, sizeof(MODULEINFO));
		ImageBase = (ADDR)info.lpBaseOfDll;
		ImageSize = info.SizeOfImage;
	}
}

static bool CompareData(BYTE *data, BYTE *pattern, char* mask)
{
	for (; *mask; ++data, ++pattern, ++mask) {
		if (*mask == 'x' && *data != *pattern)
			return false;
	}
	return (*mask) == 0;
}

static ADDR FindPattern(ADDR imageBase, ADDR imageSize, BYTE *pattern, char *mask)
{
	for (ADDR i = imageBase; i < imageBase + imageSize; ++i) {
		if (CompareData((PBYTE)i, pattern, mask))
			return i;
	}
	return 0;
}

// Public

ADDR tsh_ScanFunctionCode(char *pattern, char *mask)
{
	return FindPattern(ImageBase, ImageSize - strlen(mask), (BYTE *)pattern, mask);
}

ADDR tsh_ScanFunctionHex(char *text)
{
	unsigned int len = strlen(text);
	char *patt = new char[len];
	char *mask = new char[len];

	int outidx = 0;
	int val = 0;
	bool uk = false;
	for (unsigned int i = 0; i < len; ++i) {
		char c = text[i];
		if(c == '?') {
			uk = true;
		} else if (c >= '0' && c <= '9') {
			val = (val << 4) + (c - '0');
		} else if (c >= 'A' && c <= 'F') {
			val = (val << 4) + (c - 'A' + 10);
		} else if (c >= 'a' && c <= 'f') {
			val = (val << 4) + (c - 'a' + 10);
		} else if (c == ' ') {
			patt[outidx] = uk ? 0 : val;
			mask[outidx] = uk ? '?' : 'x';
			val = 0;
			uk = false;
			++outidx;
		}
	}

	patt[outidx] = uk ? 0 : val;
	mask[outidx] = uk ? '?' : 'x';
	++outidx;
	patt[outidx] = 0;
	mask[outidx] = 0;

	ADDR res = tsh_ScanFunctionCode(patt, mask);

	delete(patt);
	delete(mask);

	return res;
}

//////////////////////////////////////////////////
// Call Patching and Hooking

void tsh_PatchByte(ADDR location, BYTE value)
{
	DWORD oldProtection;
	VirtualProtect((void *)location, 1, PAGE_EXECUTE_READWRITE, &oldProtection);
	*((BYTE *)location) = value;
	VirtualProtect((void *)location, 1, oldProtection, &oldProtection);
}

void tsh_PatchBytes(unsigned int len, ADDR location, BYTE *repl)
{
	for (unsigned int i = 0; i < len; ++i) {
		tsh_PatchByte(location + i, repl[i]);
	}
}

void tsh_PatchInt(ADDR addr, int rval)
{
	for (unsigned int i = 0; i < 4; ++i) {
		BYTE repl = (rval >> (i*8)) & 0xFF;
		tsh_PatchByte(addr + i, repl);
	}
}

int tsh_PatchAllMatches(unsigned int len, char *pattern, char *mask, char *replace, bool debugprint)
{
	int numpatched = 0;
	for (ADDR i = ImageBase; i < ImageBase + ImageSize - len; ++i){
		if (CompareData((BYTE *)i, (BYTE *)pattern, mask)) {
			if (debugprint) BlPrintf("TSHooks: Patching call at %08x", i);

			++numpatched;
			for (ADDR c = 0; c < len; c++)
				tsh_PatchByte(i+c, replace[c]);
		}
	}
	return numpatched;
}

//////////////////////////////////////////////////
// Init

BlFunctionDefIntern(tsh_BlPrintf);

bool tsh_InitInternal()
{
	InitScanner();

	tsh_BlPrintf = (tsh_tsh_BlPrintfFnT)tsh_ScanFunctionHex((char *)"8D 44 24 08 33 D2 50 FF 74 24 08 33 C9 E8 ? ? ? ? 83 C4 08 C3");
	if (!tsh_BlPrintf)
		return false;

	#ifndef TSFUNCS_EXCLUDE_MINHOOK

	MH_STATUS status = MH_Initialize();
	if (status != MH_OK) {
		tsh_BlPrintf("TSHooks | Failed to initialize MinHook: %s", MH_StatusToString(status));
		return false;
	}
	else if (TSFUNCS_DEBUG)
		tsh_BlPrintf("TSHooks | MinHook initialized");

	#endif

	return true;
}
