// **************************************************************************
// * This file is part of the FreeFileSync project. It is distributed under *
// * GNU General Public License: http://www.gnu.org/licenses/gpl.html       *
// * Copyright (C) Zenju (zenju AT gmx DOT de) - All Rights Reserved        *
// **************************************************************************

#include "shadow.h"
#include <algorithm>
#include <string>
#include <zen/com_ptr.h>
#include <zen/com_error.h>
#include <zen/scope_guard.h>
#include <boost/thread/tss.hpp>

#ifdef USE_SHADOW_XP
#include "xp/inc/vss.h"
#include "xp/inc/vswriter.h"
#include "xp/inc/vsbackup.h"

#elif defined USE_SHADOW_2003
#include "Server 2003/inc/vss.h"
#include "Server 2003/inc/vswriter.h"
#include "Server 2003/inc/vsbackup.h"

#elif defined USE_SHADOW_WINDOWS7
#include <vss.h>      //
#include <vswriter.h> //part of Windows SDK for Windows 7
#include <vsbackup.h> //
#pragma comment(lib, "VssApi.lib")

#else
#error adapt!
#endif

using namespace zen;


struct shadow::ShadowData
{
    ShadowData(const ComPtr<IVssBackupComponents>& backupComp,
               const std::wstring& shadowVolume) : backupComp_(backupComp), shadowVolume_(shadowVolume) {}

    ComPtr<IVssBackupComponents> backupComp_;
    std::wstring shadowVolume_;
};


namespace
{
std::wstring formatVssError(HRESULT hr) //at least the one's from IVssBackupComponents::AddToSnapshotSet; return empty if no format found
{
    switch (hr)
    {
        case VSS_E_BAD_STATE:
            return L"VSS_E_BAD_STATE";
        case VSS_E_MAXIMUM_NUMBER_OF_VOLUMES_REACHED:
            return L"VSS_E_MAXIMUM_NUMBER_OF_VOLUMES_REACHED";
        case VSS_E_MAXIMUM_NUMBER_OF_SNAPSHOTS_REACHED:
            return L"VSS_E_MAXIMUM_NUMBER_OF_SNAPSHOTS_REACHED";
        case VSS_E_OBJECT_NOT_FOUND:
            return L"VSS_E_OBJECT_NOT_FOUND";
        case VSS_E_PROVIDER_NOT_REGISTERED:
            return L"VSS_E_PROVIDER_NOT_REGISTERED";
        case VSS_E_PROVIDER_VETO:
            return L"VSS_E_PROVIDER_VETO";
        case VSS_E_VOLUME_NOT_SUPPORTED:
            return L"VSS_E_VOLUME_NOT_SUPPORTED";
        case VSS_E_VOLUME_NOT_SUPPORTED_BY_PROVIDER:
            return L"VSS_E_VOLUME_NOT_SUPPORTED_BY_PROVIDER";
        case VSS_E_UNEXPECTED_PROVIDER_ERROR:
            return L"VSS_E_UNEXPECTED_PROVIDER_ERROR";
        default:
            return std::wstring();
    }
}


shadow::ShadowData createShadowCopy(const wchar_t* volumeName) //throw ComError
{
    ComPtr<IVssBackupComponents> backupComp;
    {
        HRESULT hr = ::CreateVssBackupComponents(backupComp.init());
        if (FAILED(hr))
        {
            if (hr == E_ACCESSDENIED)
                throw ComError(L"The caller does not have sufficient backup privileges or is not an administrator.", hr);
            throw ComError(L"Error calling \"CreateVssBackupComponents\".", hr);
        }
    }

    ZEN_COM_CHECK(backupComp->InitializeForBackup()); //throw ComError
    ZEN_COM_CHECK(backupComp->SetBackupState(false, false, VSS_BT_FULL)); //throw ComError

    auto waitForComFuture = [](IVssAsync& fut)
    {
        ZEN_COM_CHECK(fut.Wait());

        HRESULT hr = S_OK;
        ZEN_COM_CHECK(fut.QueryStatus(&hr, nullptr)); //check if the async operation succeeded...
        if (FAILED(hr))
            throw ComError(L"Error calling \"fut->QueryStatus\".", hr);
    };

    ComPtr<IVssAsync> gatherAsync;
    ZEN_COM_CHECK(backupComp->GatherWriterMetadata(gatherAsync.init()));
    waitForComFuture(*gatherAsync); //failure can happen if XP-version of VSS is used on Windows Vista (which needs at least VSS-Server2003 build)

    VSS_ID snapshotSetId = {};
    ZEN_COM_CHECK(backupComp->StartSnapshotSet(&snapshotSetId));
    ScopeGuard guardSnapShot = makeGuard([&] { backupComp->AbortBackup(); });
    //Quote: "This method must be called if a backup operation terminates after the creation of a
    //shadow copy set with "StartSnapshotSet" and before "DoSnapshotSet" returns."

    VSS_ID SnapShotId = {};
    {
        HRESULT hr = backupComp->AddToSnapshotSet(const_cast<wchar_t*>(volumeName), GUID_NULL, &SnapShotId);
        if (FAILED(hr))
        {
            if (hr == VSS_E_VOLUME_NOT_SUPPORTED)
                throw ComError(L"Volume Shadow Copy Service is not supported on this volume!");
            const std::wstring vssError = formatVssError(hr);
            if (!vssError.empty())
                throw ComError(L"Error calling \"backupComp->AddToSnapshotSet\": " + vssError);
            else
                throw ComError(L"Error calling \"backupComp->AddToSnapshotSet\".", hr);
        }
    }

    ComPtr<IVssAsync> prepareAsync;
    ZEN_COM_CHECK(backupComp->PrepareForBackup(prepareAsync.init()));
    waitForComFuture(*prepareAsync);

    ComPtr<IVssAsync> snapshotAsync;
    ZEN_COM_CHECK(backupComp->DoSnapshotSet(snapshotAsync.init()));
    guardSnapShot.dismiss();
    waitForComFuture(*snapshotAsync);

    VSS_SNAPSHOT_PROP props = {};
    ZEN_COM_CHECK(backupComp->GetSnapshotProperties(SnapShotId, &props));
    ZEN_ON_SCOPE_EXIT(::VssFreeSnapshotProperties(&props));

    //finally: write volume name of newly created shadow copy
    return shadow::ShadowData(backupComp, props.m_pwszSnapshotDeviceObject);
}

boost::thread_specific_ptr<std::wstring> lastErrorMessage; //use "thread_local" in C++11
}


shadow::ShadowHandle shadow::createShadowCopy(const wchar_t* volumeName)
{
    try
    {
        ShadowData result = ::createShadowCopy(volumeName); //throw ComError
        return new ShadowData(result); //shadow handle owned by caller! std::bad_alloc?
    }
    catch (const zen::ComError& e)
    {
        lastErrorMessage.reset(new std::wstring(e.toString()));
        return nullptr;
    }
}


const wchar_t* shadow::getShadowVolume(shadow::ShadowHandle handle)
{
    return handle ? handle->shadowVolume_.c_str() : nullptr; //better fail in client code than here!
}


void shadow::releaseShadowCopy(ShadowHandle handle)
{
    delete handle;
}


const wchar_t* shadow::getLastError()
{
    return !lastErrorMessage.get() ? L"" : lastErrorMessage->c_str();
}