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

#define IN
#define OUT
#define INOUT

//
//Callback function:
//
DWORD
WINAPI
SampleCaptureCallback(
	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, capturing to see if callee intends to
			//capture the file or not.
			//
			//If you do not intend to capture this file, then assign FALSE in msg_back
			//and still return WIM_MSG_SUCCESS.
			//Default is TRUE.
			//

			//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
SampleCaptureCleanup ( 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;
	DWORD  oFlag = 0, oAccess = 0;
	FARPROC callback = (FARPROC) SampleCaptureCallback;

	TCHAR *wimFile	= TEXT("C:\\sample_image.wim");  //destination .wim file
	TCHAR *tmpDir	 = TEXT("C:\\tmp"); 		 //temporary directory: OPTIONAL
	TCHAR *captureDir = TEXT("C:\\src"); 		 //capture directory or drive

	//generic capture/append call sequence
	//
	//w = WIMCreateFile()
	//WIMSetTemporaryPath() - optional
	//i = WIMCaptureImage()
	//WIMCloseHandle(i)
	//WIMCloseHandle(w)
	//

	//Set up access mode and open flag
	//
	if (argc < 1 || argc > 2) {
		SetLastError(ERROR_INVALID_PARAMETER);
		return 1;
}

	if (argc == 1) { 					//capture
		oFlag   = WIM_CREATE_ALWAYS;
		oAccess = WIM_GENERIC_WRITE;
}
	else if (*argv[1] != TEXT('a')) {  //append
		printf("need 'a' to append\n");
		SetLastError(ERROR_INVALID_PARAMETER);
		return 2;
}
	else {
		oFlag   = WIM_OPEN_EXISTING;
		oAccess = WIM_GENERIC_WRITE | WIM_GENERIC_READ;
}

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

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

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

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

	//set temporary directory to work in
	//OPTIONAL, but recommended for large WIM file
	//
	//WIMSetTemporaryPath (hWim, tmpDir); 	 //OK to fail.

	//Now capture or append image
	//
	hImg = WIMCaptureImage ( hWim,
							 captureDir,  //capture directory or drive
							 0); 	//WIM_FLAG_VERIFY

	if ( !hImg )
	{
		printf ("Cannot capture/append image\n");
		SampleCaptureCleanup (hWim, hImg, callback);
		return 5;
}

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

	return 0;
}