#include <stdexcept>
#include <Windows.h>
#include <Psapi.h>
#include <detours.h>
#include <d3d9.h>
#include <d3dx9.h>

char **game_ptr;

using cast_t = const void*(*)(const void*);
cast_t cast_REDGameInfo_Battle;

void world_to_screen(IDirect3DDevice9 *device, const D3DXVECTOR3 &in, D3DXVECTOR3 *out)
{
	D3DVIEWPORT9 viewport;
	device->GetViewport(&viewport);

	const auto *game = *game_ptr;
	const auto *world = *(char**)(game + 0x50);
	const auto *world_info = **(char***)(world + 0x3C);
	const auto *game_info = *(char**)(world_info + 0x4A8);
	const auto *camera = *(char**)(game_info+0x428);

	D3DXVECTOR3 camera_pos;
	camera_pos.x = *(float*)(camera + 0x384);
	camera_pos.y = *(float*)(camera + 0x388);
	camera_pos.z = *(float*)(camera + 0x38C);

	D3DXVec3Subtract(out, &in, &camera_pos);

	const auto clipx = (float)(viewport.Width);
	const auto clipy = (float)(viewport.Height);
	out->x = (clipx / 2.F) - out->x * ((clipx / 2.F) / tan(54.F * 3.14159F / 360.F)) / out->y;
	out->y = (clipy / 2.F) + out->z * ((clipx / 2.F) / tan(54.F * 3.14159F / 360.F)) / out->y;
}

void draw_rect(IDirect3DDevice9 *device, float x1, float y1, float x2, float y2, D3DCOLOR color)
{
	struct vertex
	{
		float x, y, z, rhw;
		DWORD color;
	};

	vertex vertices[] =
	{
		{ x1, y1, 0.F, 0.F, color },
		{ x1, y2, 0.F, 0.F, color },
		{ x2, y1, 0.F, 0.F, color },
		{ x2, y2, 0.F, 0.F, color },
	};

	device->SetRenderState(D3DRS_ALPHABLENDENABLE, true);
	device->SetRenderState(D3DRS_DESTBLEND, D3DBLEND_INVSRCALPHA);
	device->SetPixelShader(nullptr);
	device->SetFVF(D3DFVF_XYZRHW | D3DFVF_DIFFUSE);
	device->SetTexture(0, nullptr);
	device->DrawPrimitiveUP(D3DPT_TRIANGLESTRIP, 2, vertices, sizeof(vertex));
}

void draw_hitboxes(IDirect3DDevice9 *device, const char *pawn)
{
	const auto posx = *(float*)(pawn + 0x54);
	const auto posy = *(float*)(pawn + 0x5C);
	const auto flip = *(float*)(pawn + 0x70);

	const auto *asw_data = *(char**)(pawn + 0x484);
	if (asw_data == nullptr)
		return;

	const auto *hitbox_data = *(char**)(asw_data + 0x58);
	if (hitbox_data == nullptr)
		return;

	const auto hurtbox_count = *(int*)(asw_data + 0xA0);
	const auto hitbox_count = *(int*)(asw_data + 0xA4);

	for (auto i = 0; i < hitbox_count + hurtbox_count; i++)
	{
		const auto type = *(int*)(hitbox_data);

		// Convert from Arcsys 2D engine coords to UE coords
		const auto box_x = posx + *(float*)(hitbox_data + 4) * .43F * -flip;
		const auto box_y = posy - *(float*)(hitbox_data + 8) * .43F - 20.F;
		const auto box_width = *(float*)(hitbox_data + 12) * .43F * -flip;
		const auto box_height = *(float*)(hitbox_data + 16) * .43F;

		D3DXVECTOR3 ul, lr;
		world_to_screen(device, D3DXVECTOR3(box_x, 0.F, box_y), &ul);
		world_to_screen(device, D3DXVECTOR3(box_x + box_width, 0.F, box_y - box_height), &lr);

		const auto inner_color =
			type == 0 ?
			D3DCOLOR_ARGB(64, 0, 255, 0) :
			D3DCOLOR_ARGB(64, 255, 0, 0);

		const auto outer_color =
			type == 0 ?
			D3DCOLOR_ARGB(255, 0, 255, 0) :
			D3DCOLOR_ARGB(255, 255, 0, 0);

		draw_rect(device, ul[0], ul[1], lr[0], lr[1], inner_color);

		draw_rect(device, ul[0], ul[1], ul[0] + 1, lr[1], outer_color);
		draw_rect(device, ul[0], ul[1], lr[0], ul[1] + 1, outer_color);
		draw_rect(device, lr[0], ul[1], lr[0] + 1, lr[1], outer_color);
		draw_rect(device, ul[0], lr[1], lr[0], lr[1] + 1, outer_color);

		hitbox_data += 20;
	}
}

using EndScene_t = HRESULT(__stdcall*)(IDirect3DDevice9*);
EndScene_t orig_EndScene;
HRESULT __stdcall hook_EndScene(IDirect3DDevice9 *device)
{
	const auto *game = *game_ptr;
	if (game == nullptr)
		return orig_EndScene(device);

	const auto *world = *(char**)(game + 0x50);
	if (world == nullptr)
		return orig_EndScene(device);

	const auto *world_info = **(char***)(world + 0x3C);
	if (world_info == nullptr)
		return orig_EndScene(device);

	const auto *game_info = *(char**)(world_info + 0x4A8);
	if (game_info == nullptr)
		return orig_EndScene(device);

	// Make sure it's a REDGameInfo_Battle
	if (cast_REDGameInfo_Battle(game_info) == nullptr)
		return orig_EndScene(device);

	const auto *pawn1 = *(char**)(game_info+0x408);
	if (pawn1 == nullptr)
		return orig_EndScene(device);

	const auto *pawn2 = *(char**)(game_info+0x40C);
	if (pawn2 == nullptr)
		return orig_EndScene(device);

	draw_hitboxes(device, pawn1);
	draw_hitboxes(device, pawn2);
	
	return orig_EndScene(device);
}

bool get_module_bounds(const char *name, uintptr_t *start, uintptr_t *end)
{
	const auto module = GetModuleHandle(name);
	if(module == nullptr)
		return false;

	MODULEINFO info;
	GetModuleInformation(GetCurrentProcess(), module, &info, sizeof(info));
	*start = (uintptr_t)(info.lpBaseOfDll);
	*end = *start + info.SizeOfImage;
	return true;
}

uintptr_t sigscan(const char *name, const char *sig, const char *mask)
{
	uintptr_t start, end;
	if (!get_module_bounds(name, &start, &end))
		throw std::runtime_error("Module not loaded");

	const auto last_scan = end - strlen(mask) + 1;
	for (auto addr = start; addr < last_scan; addr++) {
		for (size_t i = 0;; i++) {
			if (mask[i] == '\0')
				return addr;
			if (mask[i] != '?' && sig[i] != *(char*)(addr + i))
				break;
		}
	}
	throw std::runtime_error("Sigscan failed");
}

BOOL WINAPI DllMain(
	_In_ HINSTANCE hinstDLL,
	_In_ DWORD     fdwReason,
	_In_ LPVOID    lpvReserved
	)
{
	if (fdwReason != DLL_PROCESS_ATTACH)
		return FALSE;

	game_ptr = *(char***)(sigscan(
		"GuiltyGearXrd.exe",
		"\x8B\x0D\x00\x00\x00\x00\x89\x5C\x24\x20\xE8",
		"xx????xxxxx") + 0x2);

	const auto cast_ref = sigscan(
		"GuiltyGearXrd.exe",
		"\x8B\x88\x00\x00\x00\x00\x51\xC7\x44\x24\x00\x00\x00\x00\x00\xE8",
		"xx????xxxx?????x") + 0xF;

	cast_REDGameInfo_Battle = (cast_t)(cast_ref + *(intptr_t*)(cast_ref + 1) + 5);

	const auto *dev_vtable = *(void***)(sigscan(
		"d3d9.dll",
		"\xC7\x06\x00\x00\x00\x00\x89\x86\x00\x00\x00\x00\x89\x86",
		"xx????xx????xx") + 0x2);

	orig_EndScene = (EndScene_t)(DetourFunction(
		(byte*)(dev_vtable[42]),
		(byte*)(hook_EndScene)));

	return TRUE;
}