#include <stdio.h>
#include <windows.h>
#include <tchar.h>
#include <wimgapi.h>

#define IN
#define INOUT

//
//Callback function:
//
DWORD
WINAPI
SampleApplyCallback(
	IN	DWORD msgId, //message ID
	IN	WPARAM param1,   //usually file name
	INOUT   LPARAM param2,   //usually error code
	IN	void  *unused
	)
{
	//First parameter: full file path for if WIM_MSG_PROCESS, message string for others
	TCHAR *message  = (TCHAR *) param1;
	TCHAR *filePath = (TCHAR *) param1;
	DWORD percent   = (DWORD)   param1;


	//Second parameter: message back to caller if WIM_MSG_PROCESS, error code for others
	DWORD errorCode = param2;
	DWORD *msg_back = (DWORD *) param2;


	switch ( msgId )
	{
		case WIM_MSG_PROCESS:

			//This message is sent for each file, applying to see if callee intends to apply
			//the file or not.
			//
			//If you do not intend to apply the file, then assign FALSE in msg_back
			//and still return WIM_MSG_SUCCESS
			//Default is TRUE.
			//
			//Required for displaying dir/file name (listOnly below)

			//In this example, print out the file name being applied
			//
			_tprintf(TEXT("FilePath: %s\n"), filePath);

			break;

		case WIM_MSG_ERROR:

			//This message is sent upon failure error case
			//
			printf("ERROR: %s [err = %d]\n", message, errorCode);
			break;

		case WIM_MSG_RETRY:

			//This message is sent when the file is being reapplied because of
			//network timeout. Retry is done up to five times.
			//
			printf("RETRY: %s [err = %d]\n", message, errorCode);
			break;

		case WIM_MSG_INFO:

			//This message is sent when informational message is available
			//
			printf("INFO: %s [err = %d]\n", message, errorCode);
			break;

		case WIM_MSG_WARNING:

			//This message is sent when warning message is available
			//
			printf("WARNING: %s [err = %d]\n", message, errorCode);
			break;
}

	return WIM_MSG_SUCCESS;
}


void
SampleApplyCleanup ( HANDLE hWim, HANDLE hImg, FARPROC callback )
{
	if (hImg) {
		WIMCloseHandle (hImg);
}

	if (hWim) {
		WIMCloseHandle (hWim);
}

	if (callback) {
		WIMUnregisterMessageCallback( NULL, callback );
}
}

//
//Main Function:
//
__cdecl
wmain(DWORD argc, TCHAR *argv[])
{
	HANDLE hWim = NULL;
	HANDLE hImg = NULL;
	DWORD  created = 0;
	FARPROC callback = (FARPROC) SampleApplyCallback;

	TCHAR *wimFile  = TEXT("C:\\capture.wim");  //destination .wim file
	TCHAR *tmpDir   = TEXT("C:\\tmp"); 	//temporary directory
	TCHAR *applyDir = TEXT("C:\\restore"); //apply directory or drive

	BOOL  listOnly = FALSE;  //just list dir/files in image
	DWORD imgIndex = 0;

	WIM_INFO hWimInfo = {0};

	//Need image index to apply or to display
	//
	if (argc < 2 || argc > 3 ||
		((imgIndex = _tstoi(argv[1])) == 0)) {
		printf("need image Index (1-based)\n");

		SetLastError(ERROR_INVALID_PARAMETER);
		return 2;
}

	//display dir/file listing instead of apply
	//
	if ((argc == 3)) {

		if (*argv[2] != _T('d')) {
			printf("need 'd' for directory/file listing\n");

			SetLastError(ERROR_INVALID_PARAMETER);
			return 3;
	}

		listOnly = TRUE;
}

	//Register callback
	//
	if (WIMRegisterMessageCallback( NULL,
									callback,
									NULL ) == INVALID_CALLBACK_VALUE) {
		printf ("Cannot set callback\n");
		return 1;
}

	//generic apply call sequence
	//
	//w = WIMCreateFile()
	//i = WIMLoadImage()
	//WIMApplyImage (i)
	//WIMCloseHandle(i)
	//WIMCloseHandle(w)
	//

	//Call SampleCleanup() upon exit from here
	//
	hWim = WIMCreateFile ( wimFile, 		 //existing .wim file to append to
						 WIM_GENERIC_READ, //access mode
						 WIM_OPEN_EXISTING,   //open flag
						 WIM_FLAG_VERIFY,  //recommended flag for file corruption check
						 0,
						 &created );

	if ( !hWim ) {
		printf ("Cannot open .wim file\n");

		SampleApplyCleanup(hWim, hImg, callback);
		return 2;
}

	if (!WIMGetAttributes( hWim, &hWimInfo, sizeof(hWimInfo))) {
		printf ("WIMGetAttributes failed\n");

		SampleApplyCleanup(hWim, hImg, callback);
		return 3;
}

	if (hWimInfo.ImageCount < imgIndex) {
		printf("There is no index %d in the WIM file\n", imgIndex);

		SampleApplyCleanup(hWim, hImg, callback);
		return 4;
}

	//set temporary directory to work in
	//REQUIRED for apply
	//
	if (!WIMSetTemporaryPath (hWim, tmpDir)) {
		printf("cannot set temp path to work in\n");

		SampleApplyCleanup(hWim, hImg, callback);
		return 5;
}

	//Now apply the image
	//
	hImg = WIMLoadImage ( hWim, imgIndex );
	if ( !hImg ) {
		printf("Cannot load imagex %d from src WIM file\n", imgIndex);

		SampleApplyCleanup(hWim, hImg, callback);
		return 6;
}

	if (!WIMApplyImage ( hImg,
						 applyDir, 							//apply directory or drive
						 (listOnly) ? WIM_FLAG_NO_APPLY : 0)) { // or WIM_FLAG_VERIFY
		printf("Applying image failed\n");
		SampleApplyCleanup(hWim, hImg, callback);
		return 7;
}

	//Now we are done
	//
	SampleApplyCleanup (hWim, hImg, callback);

	return 0;
}