#include "shadow.h"

#define WIN32_LEAN_AND_MEAN
#include "windows.h"
#include "inc/vss.h"
#include "inc/vswriter.h"
#include "inc/vsbackup.h"
#include <algorithm>
#include <string>
#include <cstdio>
#include <comdef.h>

//typedef GUID VSS_ID;


void writeString(const wchar_t* input, wchar_t* output, unsigned int outputBufferLen)
{
    const unsigned int newSize = min(wcslen(input) + 1, outputBufferLen); //including null-termination
    memcpy(output, input, newSize * sizeof(wchar_t));
}


std::wstring numberToHexString(const long number)
{
    wchar_t result[100];
    swprintf(result, 100, L"0x%08x", number);
    return std::wstring(result);
}


void writeErrorMsg(const wchar_t* input, HRESULT hr, wchar_t* output, unsigned int outputBufferLen)
{
    std::wstring formattedMsg(input);
    formattedMsg += L" (";
    formattedMsg += numberToHexString(hr);
    formattedMsg += L": ";
    formattedMsg += _com_error(hr).ErrorMessage();
    formattedMsg += L")";

    writeString(formattedMsg.c_str(), output, outputBufferLen);
}


bool shadow::createShadowCopy(const wchar_t* volumeName,
                              wchar_t*		 shadowVolName,
                              unsigned int   shadowBufferLen,
                              void**		 backupHandle,
                              wchar_t*		 errorMessage,
                              unsigned int	 errorBufferLen)
{
    //MessageBox(0, L"backup err", L"", 0); */
    *backupHandle = NULL;
    HRESULT hr = NULL;

    IVssBackupComponents* pBackupComponents = NULL;
    if (FAILED(hr = CreateVssBackupComponents(&pBackupComponents)))
    {
        if (hr == E_ACCESSDENIED)
            writeErrorMsg(L"The caller does not have sufficient backup privileges or is not an administrator.", hr, errorMessage, errorBufferLen);
        else
            writeErrorMsg(L"Error calling \"CreateVssBackupComponents\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    if (FAILED(hr = pBackupComponents->InitializeForBackup()))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"InitializeForBackup\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    IVssAsync* pWriteMetaData = NULL;
    if (FAILED(hr = pBackupComponents->GatherWriterMetadata( &pWriteMetaData )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"\"Shadow.dll\" might be incompatible with this Operating System!\n\
					   Error calling \"GatherWriterMetadata\".", hr, errorMessage, errorBufferLen);
        return false;
    }

    //wait for shadow copy writers to complete
    hr = pWriteMetaData->Wait();
	if (SUCCEEDED(hr))
		pWriteMetaData->QueryStatus(&hr, NULL); //check if the async operation succeeded...
	
    pWriteMetaData->Release();
    if (FAILED(hr))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"ppWriteMetaData->Wait\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    VSS_ID snapshotSetId = {0};
    if (FAILED(hr = pBackupComponents->StartSnapshotSet( &snapshotSetId )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"StartSnapshotSet\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    VSS_ID SnapShotId = {0};
    if (FAILED(hr = pBackupComponents->AddToSnapshotSet(const_cast<wchar_t*>(volumeName), GUID_NULL, &SnapShotId)))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"AddToSnapshotSet\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    if (FAILED(hr = pBackupComponents->SetBackupState( false, false, VSS_BT_FULL )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"SetBackupState\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    IVssAsync* pPrepare = NULL;
    if (FAILED(hr = pBackupComponents->PrepareForBackup( &pPrepare )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"PrepareForBackup\".", hr, errorMessage, errorBufferLen);
        return false;
    }

    hr = pPrepare->Wait();
	if (SUCCEEDED(hr))
		pPrepare->QueryStatus(&hr, NULL); //check if the async operation succeeded...

    pPrepare->Release();
    if (FAILED(hr))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"pPrepare->Wait\".", hr, errorMessage, errorBufferLen);
        return false;
    }


    IVssAsync* pDoShadowCopy = NULL;
    if (FAILED(hr = pBackupComponents->DoSnapshotSet( &pDoShadowCopy )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"DoSnapshotSet\".", hr, errorMessage, errorBufferLen);
        return false;
    }

    hr = pDoShadowCopy->Wait();
	if (SUCCEEDED(hr))
		pDoShadowCopy->QueryStatus(&hr, NULL); //check if the async operation succeeded...

	pDoShadowCopy->Release();
    if (FAILED(hr))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"pPrepare->Wait\".", hr, errorMessage, errorBufferLen);
        return false;
    }

    VSS_SNAPSHOT_PROP props;
    if (FAILED(hr = pBackupComponents->GetSnapshotProperties( SnapShotId, &props )))
    {
        releaseShadowCopy(pBackupComponents);
        writeErrorMsg(L"Error calling \"GetSnapshotProperties\".", hr, errorMessage, errorBufferLen);
        return false;
    }

    //finally: write volume name of newly created shadow copy
    writeString(props.m_pwszSnapshotDeviceObject, shadowVolName, shadowBufferLen);

    VssFreeSnapshotProperties(&props);

    *backupHandle = pBackupComponents;

    return true;
}


void shadow::releaseShadowCopy(void* backupHandle)
{
    if (backupHandle != NULL)
        static_cast<IVssBackupComponents*>(backupHandle)->Release();
}