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

#define IN

//
//Callback function:
//
DWORD
WINAPI
SampleExportCallback(
	IN   DWORD msgId, //message ID
	IN   WPARAM param1,   //usually file name
	IN   LPARAM param2,   //usually error code
	IN   void  *unused
	)
{
	//First parameter: message string
	TCHAR *message  = (TCHAR *) param1;

	//Second parameter: error code
	DWORD errorCode = param2;

	switch ( msgId )
	{
		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 file is being reapplied because of a
			//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
SampleExportCleanup( HANDLE sWim, HANDLE tWim, HANDLE sImg, FARPROC callback)
{
	if ( sWim )
	{
		WIMCloseHandle( sWim );
}
	if ( tWim )
	{
		WIMCloseHandle( tWim );
}
	if ( sImg )
	{
		WIMCloseHandle( sImg );
}
}

//
// Main Function:
//
__cdecl
wmain(DWORD argc, TCHAR *argv[])
{
	HANDLE sWim = NULL,
		 tWim = NULL;
	HANDLE sImg = NULL,
		 gImg = NULL;
	DWORD  sCreated = 0,
		 tCreated = 0,
		 sIndex = 0;

	FARPROC callback = (FARPROC) SampleExportCallback;

	TCHAR *source  = TEXT("C:\\sample_image.wim");  //Source .wim file
	TCHAR *target  = TEXT("C:\\sample_exported.wim");  //Destination .wim file
	TCHAR *tmpDir  = TEXT("C:\\tmp"); //Temporary directory: OPTIONAL

	WIM_INFO sWimInfo = {0};

	if ((argc != 2) ||
		((sIndex = _tstoi(argv[1])) == 0)) {
		printf("need image Index (1-based)\n");
		SetLastError(ERROR_INVALID_PARAMETER);
		return 1;
}

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

	//generic export call sequence
	//
	//s = WIMCreateFile()   src
	//t = WIMCreateFile()   target
	//si = WIMLoadImage(s, imagex_index_to_export);
	//WIMExportImage( si, t, )
	//WIMCloseHandle(si);
	//WIMCloseHandle(t);
	//WIMCloseHandle(s);

	//Call SampleCleanup() upon exit from here
	//

	//Open source .wim file
	//
	sWim = WIMCreateFile ( source, 		//Source .wim file to export from
						 WIM_GENERIC_READ, //Access mode
						 WIM_OPEN_EXISTING,   //Open flag
						 WIM_FLAG_VERIFY,  //Recommened flag
						 0, 			 //No need to set compression
						 &sCreated );

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

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 3;
}

	if (!WIMGetAttributes( sWim, &sWimInfo, sizeof(sWimInfo))) {
		printf ("Cannot get WIM attr\n");

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 4;
}

	_tprintf(TEXT("\nWIM path: %s\n"), sWimInfo.WimPath); //TCHAR
	printf(" image #: %d\n", sWimInfo.ImageCount);
	printf(" comp.  : %s\n", (sWimInfo.CompressionType == WIM_COMPRESS_NONE)?"NO":"YES");
	printf(" part  #: %d / %d\n\n", sWimInfo.PartNumber, sWimInfo.TotalParts);

	//both are 1-based index
	//
	if (sIndex > WIMGetImageCount(sWim)) {  //The same as sWimInfo.ImageCount
		printf ("cannot find image index %d in src WIM file\n", sIndex);

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 5;
}

	//Try to open the destination .wim file
	//
	tWim = WIMCreateFile ( target, 		 //Destination .wim file to export to
						 WIM_GENERIC_WRITE | WIM_GENERIC_READ, // access mode
						 WIM_OPEN_ALWAYS,  //Open flag
						 WIM_FLAG_VERIFY,  //Recommended flag for file corruption check
						 WIM_COMPRESS_XPRESS, //Or WIM_COMPRESS_LZX or WIM_COMPRESS_NONE
						 &tCreated );

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

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 6;
}

	//Set temporary directory to work in
	//REQUIRED
	//
	if (!WIMSetTemporaryPath (sWim, tmpDir) ||
		!WIMSetTemporaryPath (tWim, tmpDir))   {
		printf ("Cannot set temporary path\n");

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 7;
}

	//Now capture or append image
	//
	sImg = WIMLoadImage ( sWim, sIndex );
	if ( !sImg ) {
		printf("Cannot load imagex %d from src WIM file\n", sIndex);

		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 8;
}

	if (!WIMExportImage ( sImg,
						tWim,  //Capture directory or drive
						0) ) { //or WIM_EXPORT_ALLOW_DUPLICATES
		printf("Export failed\n");
		if (GetLastError() == ERROR_ALREADY_EXISTS) {
			printf("the same image already exist in the destination .wim file\n");
			// can use WIM_EXPORT_ALLOW_DUPLICATES if you do not want this
	}
		SampleExportCleanup(sWim, tWim, sImg, callback);
		return 9; 
}

	//Now we are done
	//
	SampleExportCleanup(sWim, tWim, sImg, callback);

	return 0;
}