// BestCFBF by Krishty 2019. Released to the public domain.

#if _DEBUG
#define _ALLOW_RTCc_IN_STL 1 // avoid warnings; we’re just using strlen() anyway
#endif
#if _M_X64
#	pragma pack(push, 8)
#endif
#include <cstring>
#include <objbase.h>
#include <Shlwapi.h>
#if _M_X64
#	pragma pack(pop)
#endif



static void print_stderr(char const text[]) {
	auto size = DWORD(strlen(text));
	DWORD written;
	WriteFile(GetStdHandle(STD_ERROR_HANDLE), text, size, &written, nullptr);
}

static void print_stdout(char const text[]) {
	auto size = DWORD(strlen(text));
	DWORD written;
	WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), text, size, &written, nullptr);
}

[[noreturn]] static void errorAndExit(char const text[], int code) {
	print_stderr(text);
	ExitProcess(code);
}

// Much shorter than the <vector> header:
template <typename T> void push_back(T * & array, size_t & length, T const & newInstance) {
	auto const requiredSize = sizeof(T) * (length + 1);
	if(nullptr == array) {
		array = static_cast<T *>(HeapAlloc(GetProcessHeap(), 0, requiredSize));
	} else {
		array = static_cast<T *>(HeapReAlloc(GetProcessHeap(), 0, array, requiredSize));
	}
	if(nullptr == array) {
		errorAndExit("error: out of memory\n", -18);
	}
	array[length++] = newInstance;
}

struct Stream {
	UINT64     size;
	IStream *  source;
	IStream *  dest;
	IStorage * sourceParent;
	IStorage * destParent;
};

// Copies topology from one storage to another and generates a list of streams whose content must be copied.
static void copyTopology(
	IStorage &     source,
	IStorage &     dest,
	Stream * &     streamsToCopy,
	size_t &       numberOfStreams,
	IStorage * * & storageToClose,
	size_t &       numberOfStorages
) {
	IEnumSTATSTG * childs;
	if(FAILED(source.EnumElements(0, nullptr, 0, &childs))) {
		errorAndExit("error: cannot enumerate\n", -15);
	}

	ULONG childrenActuallyFetched;
	STATSTG child;
	while(S_OK == childs->Next(1, &child, &childrenActuallyFetched)) {
		if(STGTY_STREAM == child.type) {

			Stream stream;
			stream.size = child.cbSize.QuadPart;

			// Open the stream for exclusive reading (prevents snapshots):
			if(FAILED(source.OpenStream(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, 0, &stream.source))) {
				errorAndExit("error: cannot open stream\n", 14);
			}
			stream.sourceParent = &source; source.AddRef(); // closing the storage would destroy the stream (WTF!)

			// Create a new stream with exclusive access (prevents snapshots):
			if(FAILED(dest.CreateStream(child.pwcsName, STGM_DIRECT | STGM_CREATE | STGM_READWRITE | STGM_SHARE_EXCLUSIVE, 0, 0, &stream.dest))) {
				errorAndExit("error: cannot create stream\n", -13);
			}
			stream.destParent = &dest; dest.AddRef(); // closing the storage would destroy the stream (WTF!)

			push_back(streamsToCopy, numberOfStreams, stream);

		} else if(STGTY_STORAGE == child.type) {

			// Open the storage for exclusive reading (prevents snapshots):
			IStorage * sourceStorage;
			if(FAILED(source.OpenStorage(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, nullptr, 0, &sourceStorage))) {
				errorAndExit("error: cannot open storage\n", -12);
			}

			// Create new storage with exclusive access (prevents snapshots):
			IStorage * destinationStorage;
			if(FAILED(dest.CreateStorage(child.pwcsName, STGM_DIRECT | STGM_CREATE | STGM_READWRITE | STGM_SHARE_EXCLUSIVE, 0, 0, &destinationStorage))) {
				errorAndExit("error: cannot create storage\n", -11);
			}

			copyTopology(*sourceStorage, *destinationStorage, streamsToCopy, numberOfStreams, storageToClose, numberOfStorages);

			// Keep the storage open until all data is written – otherwise “IStream::CopyTo()” will raise “STG_E_REVERTED”!
			push_back(storageToClose, numberOfStorages, destinationStorage);
			push_back(storageToClose, numberOfStorages, sourceStorage);

		} else {
			errorAndExit("error: unknown stream type\n", -10);
		}
	}
	childs->Release();
}

static bool operator !=(FILETIME const l, FILETIME const r) {
	return l.dwHighDateTime != r.dwHighDateTime || l.dwLowDateTime != r.dwLowDateTime;
}

static bool verifyIdentical(IStorage & l, IStorage & r) {
	IEnumSTATSTG * childs;
	if(FAILED(l.EnumElements(0, nullptr, 0, &childs))) {
		errorAndExit("error: cannot enumerate\n", -15);
	}

	ULONG childrenActuallyFetched;
	STATSTG child;
	while(S_OK == childs->Next(1, &child, &childrenActuallyFetched)) {
		if(STGTY_STREAM == child.type) {

			// Open the streams for exclusive reading (prevents snapshots):

			IStream * leftStream;
			if(FAILED(l.OpenStream(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, 0, &leftStream))) {
				return false;
			}

			IStream * rightStream;
			if(FAILED(r.OpenStream(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, 0, &rightStream))) {
				leftStream->Release();
				return false;
			}

			STATSTG rightStat;
			if(FAILED(rightStream->Stat(&rightStat, STATFLAG_NONAME))) {
				leftStream->Release();
				rightStream->Release();
				return false;
			}
			if(child.cbSize.QuadPart != rightStat.cbSize.QuadPart || child.clsid != rightStat.clsid || child.ctime != rightStat.ctime || child.mtime != rightStat.mtime) {
				leftStream->Release();
				rightStream->Release();
				return false;
			}

			leftStream->Release();
			rightStream->Release();

		} else if(STGTY_STORAGE == child.type) {

			// Open the storages for exclusive reading (prevents snapshots):

			IStorage * leftStorage;
			if(FAILED(l.OpenStorage(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, nullptr, 0, &leftStorage))) {
				return false;
			}

			IStorage * rightStorage;
			if(FAILED(r.OpenStorage(child.pwcsName, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_EXCLUSIVE, nullptr, 0, &rightStorage))) {
				leftStorage->Release();
				return false;
			}

			if(not verifyIdentical(*leftStorage, *rightStorage)) {
				return false;
			}

			leftStorage->Release();
			rightStorage->Release();

		} else {
			return false;
		}
	}
	childs->Release();

	return true;
}

static void optimizeCFBF(
	WCHAR const sourcePath[],
	WCHAR const destPath[],
	bool const largeSectors
) {

	IStorage * source;
	if(FAILED(StgOpenStorage(sourcePath, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_DENY_WRITE, nullptr, 0, &source))) {
		errorAndExit("cannot open source stream\n", -9);
	}

	STGOPTIONS opt = { 1, 0, 4096 };
	IStorage * destination;
	if(FAILED(StgCreateStorageEx(destPath, STGM_CREATE | STGM_READWRITE | STGM_SHARE_EXCLUSIVE, STGFMT_DOCFILE, 0, largeSectors ? &opt : nullptr, nullptr, IID_IStorage, (void * *)&destination))) {
		errorAndExit("cannot create target file\n", -8);
	}

	// 0. copy the class ID – why?!
	STATSTG StatStg;
	source->Stat(&StatStg, 0);
	destination->SetClass(StatStg.clsid);

	// 1. copy the topology
	//    • prevents fragmentation of the directory table
	//    • keep handles open, NEVER re-open storage
	//      – opening a second time adds unknown data, making the resulting file larger(!)
	Stream *     streams = nullptr;
	size_t       numberOfStreams = 0;
	IStorage * * storageToClose = nullptr;
	size_t       numberOfStorageToClose = 0;
	copyTopology(*source, *destination, streams, numberOfStreams, storageToClose, numberOfStorageToClose);

	// 2. copy data of small streams (less than 4096 B)
	//    • prevents fragmentation of the mini streams
	//    • keeps small streams like Summary Information close to the beginning of the file
	for(auto stream = streams; stream < streams + numberOfStreams; ++stream) {
		if(stream->size < 4096)
		{
			ULARGE_INTEGER nread, nwritten;
			ULARGE_INTEGER size; size.QuadPart = stream->size;
			if(FAILED(stream->source->CopyTo(stream->dest, size, &nread, &nwritten)) || stream->size != nread.QuadPart || stream->size != nwritten.QuadPart) {
				errorAndExit("cannot copy small stream\n", -7);
			}
		}
	}

	// 3. copy data of large streams (4096 B or more)
	//    • prevents fragmentation of normal streams
	for(auto stream = streams; stream < streams + numberOfStreams; ++stream) {
		if(stream->size >= 4096)
		{
			ULARGE_INTEGER nread, nwritten;
			ULARGE_INTEGER size; size.QuadPart = stream->size;
			if(FAILED(stream->source->CopyTo(stream->dest, size, &nread, &nwritten)) || stream->size != nread.QuadPart || stream->size != nwritten.QuadPart) {
				errorAndExit("cannot copy large stream\n", -6);
			}
		}
	}

	// 4. commit the changes
	//    • STGC_OVERWRITE throws away old versions/snapshots (makes some sectors difference)
	//    • STGC_CONSOLIDATE
	for(auto stream = streams; stream < streams + numberOfStreams; ++stream) {
		stream->source->Release();
		stream->dest->Release();
		stream->sourceParent->Release();
		stream->destParent->Release();
	}
	for(auto it = storageToClose; it < storageToClose + numberOfStorageToClose; ++it) {
		(*it)->Release();
	}
    destination->Commit(STGC_OVERWRITE | STGC_CONSOLIDATE);
	destination->Release();
	source->Release();

	// 5. Make sure we didn’t miss out on anything:
	IStorage * l;
	if(FAILED(StgOpenStorage(sourcePath, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_DENY_WRITE, nullptr, 0, &l))) {
		errorAndExit("cannot re-open source stream\n", -17);
	}
	IStorage * r;
	if(FAILED(StgOpenStorage(destPath, nullptr, STGM_DIRECT | STGM_READ | STGM_SHARE_DENY_WRITE, nullptr, 0, &r))) {
		errorAndExit("cannot re-open target stream\n", -17);
	}
	if(not verifyIdentical(*l, *r)) {
		errorAndExit("result is broken\n", -17);
	}
	l->Release();
	r->Release();

	print_stdout("OK\n");
}

#include "version for C.h"

int __stdcall entry() {
	int argc;
	auto argv = CommandLineToArgvW(GetCommandLineW(), &argc);

	WCHAR const * sourcePath = nullptr;
	WCHAR const * destPath   = nullptr;
	auto          v4         = false;
	for(auto toParam = argv + 1; toParam < argv + argc; ++toParam) { // parameter [0] is the executable name
		auto const param = *toParam;
		if(param[0] == '-' && param[1] == 'v' && param[2] == '4' && param[3] == '\0') {
			v4 = true;
		} else {
			if(nullptr == sourcePath) {
				sourcePath = *toParam;
			} else {
				destPath = *toParam;
			}
		}
	}

	if(nullptr == sourcePath || nullptr == destPath) {
		print_stdout(
			"Papa's Best CFBF Optimizer " VERSION " from https://papas-best.com\n"
			"usage: BestCFBF [-v4] <in path> <out path>\n"
			"  input and output must be different files\n"
			"  -v4 uses 4096-B sectors; 512 B otherwise\n"
		);
		return -16;
	}

	CoInitialize(nullptr);

	optimizeCFBF(sourcePath, destPath, v4);

	CoUninitialize();

	ExitProcess(0);
}

#pragma comment(linker, "/ENTRY:entry")
#pragma comment(lib, "Kernel32.lib")
#pragma comment(lib, "Shell32.lib")
#if _DEBUG
	#pragma comment(lib, "vcruntimed")
	#pragma comment(lib, "ucrtd")
#endif