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

#define PI 3.14159F

char **game_ptr;

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

void angle_vectors(
	const float p, const float y, const float r,
	float *forward, float *right, float *up)
{
	const auto sp = sin(-p);
	const auto cp = cos(-p);

	const auto sy = sin(y);
	const auto cy = cos(y);

	const auto sr = sin(-r);
	const auto cr = cos(-r);

	if (forward != nullptr) {
		forward[0] = cp * cy;
		forward[1] = cp * sy;
		forward[2] = -sp;
	}
	if (right != nullptr) {
		right[0] = -1 * sr * sp * cy + -1 * cr * -sy;
		right[1] = -1 * sr * sp * sy + -1 * cr * cy;
		right[2] = -1 * sr * cp;
	}
	if (up != nullptr) {
		up[0] = cr * sp * cy + -sr * -sy;
		up[1] = cr * sp * sy + -sr * cy;
		up[2] = cr * cp;
	}
}

float vec_dot(float *a, float *b)
{
	return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}

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);

	D3DXVECTOR3 relative_pos;
	D3DXVec3Subtract(&relative_pos, &in, &camera_pos);

	const auto clipx = (float)(viewport.Width);
	const auto clipy = (float)(viewport.Height);

	const auto pitch = (float)(*(int*)(camera + 0x390)) / 32768.F * PI;
	const auto yaw = (float)(*(int*)(camera + 0x394)) / 32768.F * PI;
	const auto roll = (float)(*(int*)(camera + 0x398)) / 32768.F * PI;

	const auto fov = *(float*)(camera + 0x39C);

	float forward[3], right[3], up[3];
	angle_vectors(pitch, yaw, roll, forward, right, up);

	out->x = vec_dot(relative_pos, right);
	out->y = vec_dot(relative_pos, up);
	out->z = vec_dot(relative_pos, forward);

	out->x = (clipx / 2.F) - out->x * ((clipx / 2.F) / tan(fov * PI / 360.F)) / out->z;
	out->y = (clipy / 2.F) - out->y * ((clipx / 2.F) / tan(fov * PI / 360.F)) / out->z;
}

void draw_rect(
	IDirect3DDevice9 *device,
	const D3DXVECTOR3 &p1,
	const D3DXVECTOR3 &p2,
	const D3DXVECTOR3 &p3,
	const D3DXVECTOR3 &p4,
	D3DCOLOR inner_color,
	D3DCOLOR outer_color)
{
	struct vertex
	{
		float x, y, z, rhw;
		DWORD color;
	};

	D3DXVECTOR3 sp1, sp2, sp3, sp4;
	world_to_screen(device, p1, &sp1);
	world_to_screen(device, p2, &sp2);
	world_to_screen(device, p3, &sp3);
	world_to_screen(device, p4, &sp4);

	device->SetRenderState(D3DRS_ALPHABLENDENABLE, true);
	device->SetRenderState(D3DRS_DESTBLEND, D3DBLEND_INVSRCALPHA);
	device->SetPixelShader(nullptr);
	device->SetFVF(D3DFVF_XYZRHW | D3DFVF_DIFFUSE);
	device->SetTexture(0, nullptr);

	vertex vertices[] =
	{
		{ sp1.x, sp1.y, 0.F, 0.F, inner_color },
		{ sp2.x, sp2.y, 0.F, 0.F, inner_color },
		{ sp3.x, sp3.y, 0.F, 0.F, inner_color },
		{ sp4.x, sp4.y, 0.F, 0.F, inner_color },
	};

	device->DrawPrimitiveUP(D3DPT_TRIANGLESTRIP, 2, vertices, sizeof(vertex));

	vertex outline[] =
	{
		{ sp1.x, sp1.y, 0.F, 0.F, outer_color },
		{ sp2.x, sp2.y, 0.F, 0.F, outer_color },
		{ sp4.x, sp4.y, 0.F, 0.F, outer_color },
		{ sp3.x, sp3.y, 0.F, 0.F, outer_color },
		{ sp1.x, sp1.y, 0.F, 0.F, outer_color },
	};

	device->DrawPrimitiveUP(D3DPT_LINESTRIP, 4, outline, 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;
		const auto box_width = *(float*)(hitbox_data + 12) * .43F * -flip;
		const auto box_height = *(float*)(hitbox_data + 16) * .43F;

		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,
			D3DXVECTOR3(box_x, 0.F,  box_y),
			D3DXVECTOR3(box_x, 0.F,  box_y - box_height),
			D3DXVECTOR3(box_x + box_width, 0.F,  box_y),
			D3DXVECTOR3(box_x + box_width, 0.F,  box_y - box_height),
			inner_color,
			outer_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;
}