#include <iostream>
#include <windows.h>
#include <TlHelp32.h>
#include <string>
#include <Psapi.h>
#include <functional>
#include <stdlib.h>

#ifndef Hook_DLL_Name
#define Hook_DLL_Name L""
#endif

using std::wstring;
using log_function = std::function<void(const wchar_t*)>;


static wstring GetDllFullPath(wstring& dllname)
{
	wchar_t buffer[MAX_PATH];

	DWORD length = GetCurrentDirectoryW(MAX_PATH, buffer);

	if (length != 0)
	{
		return wstring(buffer) + L"\\" + dllname;
	}

	return wstring(L"");
}

static bool ExistFile(wstring& inFileAddress)
{
	bool result = false;

	DWORD fileAttributes = GetFileAttributesW(inFileAddress.c_str());

	if (fileAttributes != INVALID_FILE_ATTRIBUTES)
	{
		result = true;
	}

	return result;
}

static DWORD GetProcessIdByName(wstring& processName)
{
	DWORD result = 0;

	HANDLE snapshotHandle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	if (snapshotHandle == INVALID_HANDLE_VALUE)
	{
		return 0;
	}

	PROCESSENTRY32W processEntry;
	processEntry.dwSize = sizeof(PROCESSENTRY32W);

	if (Process32FirstW(snapshotHandle, &processEntry))
	{
		do
		{
			if (processName == processEntry.szExeFile)
			{
				result = processEntry.th32ProcessID;

				break;
			}
		} while (Process32NextW(snapshotHandle, &processEntry));
	}

	CloseHandle(snapshotHandle);

	return result;
}

static HMODULE GetRemoteModuleHandle(HANDLE hProcess, const wstring& dllPath)
{
	HMODULE hModules[1024];
	DWORD cbNeeded;

	// Get a list of all the modules in the target process
	if (EnumProcessModulesEx(hProcess, hModules, sizeof(hModules), &cbNeeded, LIST_MODULES_ALL))
	{
		for (unsigned int index = 0; index < (cbNeeded / sizeof(HMODULE)); ++index)
		{
			wchar_t szModName[MAX_PATH];

			// Get the full path to the module's file
			if (GetModuleFileNameExW(hProcess, hModules[index], szModName, sizeof(szModName) / sizeof(wchar_t)))
			{
				// Check if this is the module we want to unload
				if (dllPath == szModName)
				{
					return hModules[index];
				}
			}
		}
	}

	return NULL;
}

static bool UnloadDLL(HANDLE hProcess, wstring& dllPath, log_function logger)
{
	// Get the base address of the loaded DLL in the target process
	HMODULE hModule = GetRemoteModuleHandle(hProcess, dllPath);

	if (!hModule)
	{
		logger(L"Failed to find module handle for the DLL in the target process.");

		return false;
	}

	// Get the address of FreeLibrary in kernel32.dll
	HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");

	if (!hKernel32)
	{
		logger(L"Failed to get handle of 'Kernel32.dll'.");

		CloseHandle(hProcess);

		return false;
	}

	// Get the address of FreeLibrary in kernel32.dll
	FARPROC pFreeLibrary = GetProcAddress(hKernel32, "FreeLibrary");

	if (!pFreeLibrary)
	{
		logger(L"Failed to get address of 'FreeLibrary'.");

		CloseHandle(hProcess);

		return false;
	}

	// Create a remote thread in the target process to unload the DLL
	HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)pFreeLibrary, hModule, 0, NULL);

	if (hThread == NULL)
	{
		logger(L"Failed to create remote thread to unload DLL.");

		CloseHandle(hProcess);

		return false;
	}

	// Wait for the remote thread to finish
	WaitForSingleObject(hThread, INFINITE);

	CloseHandle(hThread);

	return true;
}

static bool InjectDLL(DWORD processID, wstring& dllPath, log_function logger)
{
	if (ExistFile(dllPath))
	{
		HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processID);

		if (hProcess == NULL)
		{
			logger(L"Failed to open target process.");

			return false;
		}

		// Unload the dll if was loaded in the target process
		UnloadDLL(hProcess, dllPath, logger);

		// Allocate memory in the target process for the DLL path
		DWORD dllPathSize = DWORD((dllPath.length() + 1) * sizeof(wchar_t));

		LPVOID pDllPath = VirtualAllocEx(hProcess, NULL, dllPathSize, MEM_COMMIT, PAGE_READWRITE);

		if (pDllPath == NULL)
		{
			logger(L"Failed to allocate memory in target process.");

			CloseHandle(hProcess);

			return false;
		}

		// Write the DLL path to the allocated memory
		if (!WriteProcessMemory(hProcess, pDllPath, (LPVOID)dllPath.c_str(), dllPathSize, NULL))
		{
			logger(L"Failed to write DLL path to target process memory.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		// Get the address of LoadLibraryA in kernel32.dll
		HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");

		if (!hKernel32)
		{
			logger(L"Failed to get handle of 'Kernel32.dll'.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		FARPROC pLoadLibraryW = GetProcAddress(hKernel32, "LoadLibraryW");

		if (!pLoadLibraryW)
		{
			logger(L"Failed to get address of 'LoadLibraryW'.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		// Create a remote thread in the target process to load the DLL
		HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)pLoadLibraryW, pDllPath, 0, NULL);

		if (hThread == NULL)
		{
			logger(L"Failed to create remote thread.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		// Wait for the remote thread to finish
		WaitForSingleObject(hThread, INFINITE);

		// Clean up
		VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
		CloseHandle(hThread);
		CloseHandle(hProcess);

		return true;
	}
	else
	{
		logger(L"The dll file does not exist.");
	}

	return false;
}

static bool OpenAndInjectDLL(wstring& exepath, wstring& dllPath, log_function logger, bool waitExit)
{
	if (ExistFile(dllPath))
	{
		STARTUPINFOW si;
		PROCESS_INFORMATION pi;

		ZeroMemory(&si, sizeof(si));
		si.cb = sizeof(si);
		ZeroMemory(&pi, sizeof(pi));

		// 尝试创建一个新进程
		if (!CreateProcessW(NULL,   // 没有启动应用程序的路径，使用命令行
			(wchar_t*)exepath.c_str(), // 要启动的应用程序名
			NULL,           // 进程的安全属性
			NULL,           // 线程的安全属性
			FALSE,          // 不继承句柄
			CREATE_SUSPENDED,// 创建标志
			NULL,           // 使用父进程的环境块
			NULL,           // 使用父进程的当前目录
			&si,            // STARTUPINFO结构
			&pi))          // PROCESS_INFORMATION结构
		{
			return 0;
		}

		HANDLE hProcess = pi.hProcess;

		if (hProcess == NULL)
		{
			logger(L"Failed to open process.");

			return false;
		}

		// Allocate memory in the target process for the DLL path
		DWORD dllPathSize = DWORD((dllPath.length() + 1) * sizeof(wchar_t));

		LPVOID pDllPath = VirtualAllocEx(hProcess, NULL, dllPathSize, MEM_COMMIT, PAGE_READWRITE);

		if (pDllPath == NULL)
		{
			logger(L"Failed to allocate memory in target process.");

			CloseHandle(hProcess);
			return false;
		}

		// Write the DLL path to the allocated memory
		if (!WriteProcessMemory(hProcess, pDllPath, (LPVOID)dllPath.c_str(), dllPathSize, NULL))
		{
			logger(L"Failed to write DLL path to target process memory.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		// Get the address of LoadLibraryA in kernel32.dll
		HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");

		if (!hKernel32)
		{
			logger(L"Failed to get handle of 'Kernel32.dll'.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		FARPROC pLoadLibraryW = GetProcAddress(hKernel32, "LoadLibraryW");

		if (!pLoadLibraryW)
		{
			logger(L"Failed to get address of 'LoadLibraryW'.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}

		// Create a remote thread in the target process to load the DLL
		HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)pLoadLibraryW, pDllPath, 0, NULL);

		if (hThread == NULL)
		{
			logger(L"Failed to create remote thread.");

			VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
			CloseHandle(hProcess);

			return false;
		}
		std::cout << pi.hThread << std::endl;
		Sleep(1);
		ResumeThread(pi.hThread);
		Sleep(10);
		ResumeThread(hThread);
		Sleep(1);

		// Wait for the remote thread to finish
		WaitForSingleObject(hThread, INFINITE);

		// Clean up
		VirtualFreeEx(hProcess, pDllPath, 0, MEM_RELEASE);
		CloseHandle(hThread);
		if (waitExit)
		{
			WaitForSingleObject(hProcess, INFINITE);
		}
		CloseHandle(hProcess);

		return true;
	}
	else
	{
		logger(L"The dll file does not exist.");
	}

	return false;
}

static bool InjectDLL(wstring& processName, wstring& dllPath, wstring& exePath, log_function logger)
{
	bool result = false;

	DWORD processId = GetProcessIdByName(processName);
	if (processId != 0)
	{
		result = InjectDLL(processId, dllPath, logger);
	}
	else
	{
		logger(L"The process name was not found.");
	}

	return result;
}

static bool UnloadDLL(wstring& processName, wstring& dllPath, log_function logger)
{
	if (ExistFile(dllPath))
	{
		DWORD processId = GetProcessIdByName(processName);

		if (processId != 0)
		{
			HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processId);
			if (hProcess == NULL)
			{
				logger(L"UnloadDLL: Failed to open target process.");

				return false;
			}

			return UnloadDLL(hProcess, dllPath, logger);
		}
		else
		{
			logger(L"UnloadDLL: The process name was not found.");
		}
	}

	return false;
}

void HookWindowsApi(wstring& dllPath, wstring& processName, wstring& exePath)
{
	auto logger = [](const wchar_t* msg)
		{
			std::wcout << msg << std::endl;
		};
	// Inject the DLL
	if (OpenAndInjectDLL(exePath, dllPath, logger, true))
	{
		std::cout << "DLL injection success." << std::endl;
	}
	else
	{
		std::cout << "DLL injection failed." << std::endl;
	}

	std::cout << "DLL injection stop." << std::endl;
	UnloadDLL(processName, dllPath, logger);
}

int main()
{
	//SetConsoleOutputCP(CP_UTF8);
	//SetConsoleMode(GetStdHandle(STD_OUTPUT_HANDLE), ENABLE_PROCESSED_OUTPUT | ENABLE_WRAP_AT_EOL_OUTPUT);

	HWND hWnd = GetConsoleWindow();
	ShowWindow(hWnd, SW_HIDE);

	auto cmdLine = GetCommandLineW();

	int argc;
	auto argv = CommandLineToArgvW(cmdLine, &argc);
	auto exeName = L"AGame.exe";
	if (argc != 2)
	{
		std::cout << "use default exe path." << std::endl;
	}
	else
	{
		exeName = argv[1];
	}

	wstring dllName = Hook_DLL_Name;
	wstring dllPath = GetDllFullPath(dllName);
	wstring processName = wstring(exeName);
	
	HookWindowsApi(dllPath, processName, GetDllFullPath(wstring(exeName)));

	return 0;
}