#pragma once

#include <windows.h>
#include <stdio.h>

namespace RITZ
{
	template<typename FUNCTIONTYPE>
	class HOOK
	{
	public:
		const static int PATCH_LEN = 5;

		void Hook(FUNCTIONTYPE from, FUNCTIONTYPE to, HMODULE hmod = GetCurrentModule())
		{
			SafeMemcpy(overwritten, from, PATCH_LEN); // Save original function code

			static BYTE jumpcode[] = {0xE9, 0,0,0,0}; // Offset jump code
			*((DWORD*)(jumpcode+1)) = (DWORD)to - (DWORD)from - PATCH_LEN; // Write in jump operand

			SafeMemcpy(from, jumpcode, PATCH_LEN); // Insert hook

			orig = from;
			hook = to;
			this->hmod = hmod;
			hooked = true;
		}

		void Hook(char *from_name, FUNCTIONTYPE to, HMODULE hmod = GetCurrentModule())
		{
			FARPROC procaddr = GetProcAddress(hmod, from_name);
			if(procaddr)
			{
				Hook((FUNCTIONTYPE)procaddr, to, hmod);
			}
		}

		void Unhook()
		{
			SafeMemcpy(orig, overwritten, PATCH_LEN);
			hooked = false;
		}

		void Rehook()
		{
			Hook(orig, hook, this->hmod);
		}

		FUNCTIONTYPE orig; // For calling after unhooking.
		FUNCTIONTYPE hook;

	private:
		BYTE overwritten[PATCH_LEN];
		HMODULE hmod;
		bool hooked;
		
		void SafeMemcpy(void *to, void *from, DWORD size)
		{
			DWORD oldprot_from;
			DWORD oldprot_to;

			VirtualProtect(from, 8, PAGE_EXECUTE_READWRITE, &oldprot_from);
			VirtualProtect(to, 8, PAGE_EXECUTE_READWRITE, &oldprot_to);

			__int64 todata = *((__int64*)to);
			__int64 fromdata = *((__int64*)from);

			// So that we don't overwrite more than 'size',
			// copy the data that we are supposed to keep to the 'fromdata' buffer.
			memcpy(((BYTE*)&fromdata)+size, ((BYTE*)&todata)+size, 8-size);

			__asm
			{
				mov edx, dword ptr ds:[todata+4]
				mov eax, dword ptr ds:[todata]

				mov ecx, dword ptr ds:[fromdata+4]
				mov ebx, dword ptr ds:[fromdata]

				mov esi, to // Load to address into register for cmpxchg8b.
				lock cmpxchg8b qword ptr ds:[esi] // Write
			}

			VirtualProtect(from, 8, oldprot_from, &oldprot_from);
			VirtualProtect(to, 8, oldprot_from, &oldprot_to);
		}

		static HMODULE GetCurrentModule()
		{
			MEMORY_BASIC_INFORMATION mbi;
			static int dummy;
			VirtualQuery(&dummy, &mbi, sizeof(mbi));

			return (HMODULE)mbi.AllocationBase;
		}

	public:
		HOOK()
		{
			hooked = false;
		}

		HOOK(FUNCTIONTYPE from, FUNCTIONTYPE to, HMODULE hmod)
		{
			Hook(from, to, hmod);
		}

		~HOOK()
		{
			if(hooked)
				Unhook();
		}
	};
}