#include "ClientLoader.h"
#include <iostream>
#include <processthreadsapi.h>

#include "../DragonHook/FiestaHook.h"
#include "File.h"
#include "StringHelper.h"

ClientLoader::ClientLoader(const std::string& client_root, const std::string& client_executable, const std::string args)
{
	const auto exeFile = File(client_root + client_executable);
	const auto exeExists = exeFile.exists();
	if(!exeExists)
	{
		std::cerr << "Bad path to client!" << std::endl;
		return;
	}

	this->client_root = client_root;
	this->client_executable = client_executable;
	this->client_args = args;

	std::cout << "Client Path OK!" << std::endl;
}

void ClientLoader::start_client()
{
	STARTUPINFO start_info = { 0 };
	PROCESS_INFORMATION process_info = { nullptr };

	const auto full_command = client_root + client_executable + " " + client_args;

	const bool start_ok = CreateProcess(nullptr, 
		StringHelper::string_to_lpwstr(full_command),
		nullptr, nullptr, TRUE,
		CREATE_NEW_PROCESS_GROUP, nullptr, nullptr,
		&start_info, &process_info);

	if(start_ok)
	{
		std::cout << "Started client." << std::endl;

		const auto base_addr = GetProcessBaseAddress(process_info.dwProcessId);

		DragonHook::FiestaHook::Start(base_addr);
		
		char mask = '?';
		auto addr = DragonHook::FiestaHook::FindPattern(x3_pattern.c_str(), &mask);

		std::cout << "Address: " << addr << std::endl;
		
		WaitForSingleObject(process_info.hProcess, INFINITE);
		std::cout << "Client exited." << std::endl;
		CloseHandle(process_info.hProcess);
		CloseHandle(process_info.hThread);
	}
	else
	{
		const auto err = GetLastError();
		switch(err)
		{
			case 2:
				std::cerr << "File not found!" << std::endl;
				break;
			case 87:
				std::cerr << "Client already open!" << std::endl;
				break;
			default:
				printf("CreateProcess failed (%lu).\n", err);
				break;
		}
	}
}

DWORD_PTR ClientLoader::GetProcessBaseAddress(DWORD processID)
{
	DWORD_PTR baseAddress = 0;
	HANDLE process_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processID);
	DWORD bytesRequired;

	if (process_handle)
	{
		if (EnumProcessModules(process_handle, nullptr, 0, &bytesRequired))
		{
			if (bytesRequired)
			{
				const auto module_array_bytes = static_cast<LPBYTE>(LocalAlloc(LPTR, bytesRequired));

				if (module_array_bytes)
				{
					unsigned int moduleCount = bytesRequired / sizeof(HMODULE);
					const auto module_array = reinterpret_cast<HMODULE*>(module_array_bytes);

					if (EnumProcessModules(process_handle, module_array, bytesRequired, &bytesRequired))
					{
						baseAddress = reinterpret_cast<DWORD_PTR>(module_array[0]);
					}

					LocalFree(module_array_bytes);
				}
			}
		}

		CloseHandle(process_handle);
	}

	return baseAddress;
}