#include <windows.h>
#include <detours.h>
#include <iostream>
#include <string>
#include <shlobj_core.h>
#include <cstdlib>

using std::wstring;

typedef HRESULT(WINAPI* SHGetKnownFolderPathF)(REFKNOWNFOLDERID, DWORD, HANDLE, PWSTR*);
typedef HRESULT(WINAPI* SHGetFolderPathWF)(HWND, int, HANDLE, DWORD, LPWSTR);

SHGetKnownFolderPathF Original_SHGetKnownFolderPath = nullptr;
SHGetFolderPathWF Original_SHGetFolderPathW = nullptr;

static wstring target_datapath;

HRESULT Hooked_SHGetKnownFolderPath(REFKNOWNFOLDERID rfid, DWORD dwFlags, HANDLE hToken, PWSTR* ppszPath)
{
	if (rfid == FOLDERID_LocalAppDataLow)
	{
		if (ppszPath != nullptr)
		{
			*ppszPath = new wchar_t(target_datapath.size());
			std::wcscpy(*ppszPath, target_datapath.c_str());
		}

		return S_OK;
	}
	return Original_SHGetKnownFolderPath(rfid, dwFlags, hToken, ppszPath);
}

HRESULT Hooked_SHGetFolderPathW(HWND hwnd, int csidl, HANDLE hToken, DWORD dwFlags, LPWSTR pszPath)
{
	if ((csidl & CSIDL_LOCAL_APPDATA) != 0)
	{
		pszPath = (LPWSTR)target_datapath.c_str();
		return S_OK;
	}
	return Original_SHGetFolderPathW(hwnd, csidl, hToken, dwFlags, pszPath);
}

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

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

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

	return result;
}
wchar_t buffer[MAX_PATH];

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
	if (ul_reason_for_call == DLL_PROCESS_ATTACH)
	{
		DetourTransactionBegin();
		DetourUpdateThread(GetCurrentThread());
		
		const char* env_path_value = getenv("WIN_HOOK_ENV1");
		if (env_path_value == nullptr)
		{
			DWORD length = GetCurrentDirectoryW(MAX_PATH, buffer);

			if (length != 0)
			{
				target_datapath = wstring(buffer);
			}
		}
		else
		{
			mbstowcs(buffer, env_path_value, strlen(env_path_value));
			target_datapath = wstring(buffer);
		}

		Original_SHGetKnownFolderPath = (SHGetKnownFolderPathF)DetourFindFunction("Shell32.dll", "SHGetKnownFolderPath");
		DetourAttach(&(PVOID&)Original_SHGetKnownFolderPath, Hooked_SHGetKnownFolderPath);

		//Original_SHGetFolderPathW = (SHGetFolderPathWF)DetourFindFunction("Shell32.dll", "SHGetFolderPathW");
		//DetourAttach(&(PVOID&)Original_SHGetFolderPathW, Hooked_SHGetFolderPathW);

		DetourTransactionCommit();
	}
	else if (ul_reason_for_call == DLL_PROCESS_DETACH)
	{
		// Unhook the CreateFileW function
		DetourTransactionBegin();
		DetourUpdateThread(GetCurrentThread());
		DetourDetach(&(PVOID&)Original_SHGetKnownFolderPath, Hooked_SHGetKnownFolderPath);

		DetourDetach(&(PVOID&)Original_SHGetFolderPathW, Hooked_SHGetFolderPathW);
		DetourTransactionCommit();
	}

	return TRUE;
}