#include <windows.h>
#include <algorithm>
#include <shlwapi.h>
#include <vector>
#include <iostream>
#include <cstdio>
#include <exception>
#include <sstream>

static LPCWSTR SOFTWARE_HIVE = L"C:\\Windows\\System32\\Config\\software";

void PrintSoftwareSize()
{
	WIN32_FILE_ATTRIBUTE_DATA wfad;
	if(GetFileAttributesExW(SOFTWARE_HIVE, GetFileExInfoStandard, &wfad))
	{
		ULONGLONG size = (static_cast<ULONGLONG>(wfad.nFileSizeHigh) << 32) | wfad.nFileSizeLow;
		std::wcout << SOFTWARE_HIVE << L" size = " << size << L" bytes";
		WCHAR niceSize[32];
		if(StrFormatByteSizeW(size, niceSize, ARRAYSIZE(niceSize)))
		{
			std::wcout << L" (" << niceSize << L')';
		}
		std::wcout << L'\n';
	}
	else
	{
		std::wcout << L"Failed to get size of " << SOFTWARE_HIVE << L'\n';
		ExitProcess(2);
	}
}

template<class Iter>
void CloseKeyRange(Iter begin, Iter end)
{
	std::for_each(begin, end, &RegCloseKey);
}

void MakeTree(const ULONG treeDepth, const ULONG siblings)
{
	HKEY hKeySoftware = NULL;
	LONG error = 0;
	const size_t numBytesInFilePath = (wcslen(_wpgmptr) + 1) * sizeof(WCHAR);
	if((error = RegOpenKeyEx(HKEY_LOCAL_MACHINE, L"Software", 0, KEY_CREATE_SUB_KEY | KEY_QUERY_VALUE, &hKeySoftware)) == 0)
	{
		std::vector<HKEY> treeKeys;
		treeKeys.reserve(60000);
		try
		{
			HKEY hTreeKey = NULL;
			if((error = RegCreateKeyW(hKeySoftware, L"TestTreeRoot", &hTreeKey)) == 0)
			{
				std::cout << "Creating tree\n";
				treeKeys.push_back(hTreeKey);
				for(ULONG depth = 0; depth < treeDepth; ++depth)
				{
					size_t numKeys = treeKeys.size();
					for(ULONG sib = 0; sib < siblings; ++sib)
					{
						std::wostringstream treePath;
						treePath << L"TestTreeChildDepth" << sib;
						for(UINT keyIndex = 0; keyIndex < numKeys; ++keyIndex)
						{
							HKEY hNewKey;
							if((error = RegCreateKeyW(treeKeys[keyIndex], treePath.str().c_str(), &hNewKey) == 0))
							{
								treeKeys.push_back(hNewKey);
								RegSetValueW(hNewKey, NULL, REG_SZ, _wpgmptr, numBytesInFilePath);
							}
							else
							{
								throw std::exception();
							}
						}
					}
					CloseKeyRange(treeKeys.begin(), treeKeys.begin() + numKeys);
					treeKeys.erase(treeKeys.begin(), treeKeys.begin() + numKeys);
				}
				std::cout << "Tree creation succeeded\n";
				RegFlushKey(hKeySoftware);
			}
			else
			{
				std::cout << "Root creation failed with " << error << '\n';
				ExitProcess(1);
			}
		}
		catch(const std::exception&){} // long live exceptions for flow control 
		CloseKeyRange(treeKeys.begin(), treeKeys.end());
		PrintSoftwareSize();
		std::cout << "Deleting tree\n";
		if((error = SHDeleteKeyW(hKeySoftware, L"TestTreeRoot")) == 0)
		{
			RegFlushKey(hKeySoftware);
			RegCloseKey(hKeySoftware);
			std::cout << "Tree deleted\n";
		}
		else
		{
			RegCloseKey(hKeySoftware);
			std::cout << "Tree deletion failed - " << error << '\n';
			ExitProcess(1);
		}
	}
	else
	{
		std::cout << "Couldn't open HKLM\\Software because of error " << error << '\n';
		ExitProcess(1);
	}
}

int __cdecl wmain(int argc, wchar_t** argv)
//int WINAPI WinMain(HINSTANCE, HINSTANCE, LPSTR, int)
{
	if(argc >= 3)
	{
		ULONG treeDepth = _wtol(argv[1]);
		ULONG siblings = _wtol(argv[2]);
		PrintSoftwareSize();
		MakeTree(treeDepth, siblings);
		PrintSoftwareSize();
		MakeTree(treeDepth, siblings);
		PrintSoftwareSize();
	}
	else
	{
		std::cout << "Usage: RegSizeTest [tree depth] [number of siblings]\n";
	}
	return 0;
}