// **************************************************************************
// * 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 (zhnmju123 AT gmx DOT de) - All Rights Reserved    *
// **************************************************************************

#include "file_op.h"
#include <algorithm>
#include <string>
#include <vector>

#define WIN32_LEAN_AND_MEAN
#include <zen/com_ptr.h>
#include <zen/com_error.h>
#include <zen/scope_guard.h>

#include <boost/thread/tss.hpp>

#include <RestartManager.h>
#pragma comment(lib, "Rstrtmgr.lib")

#define STRICT_TYPED_ITEMIDS //better type safety for IDLists
#include <Shlobj.h>
#include <shobjidl.h>
#include <shellapi.h> //shell constants such as FO_* values

using namespace zen;


namespace
{
void moveToRecycleBin(const wchar_t* fileNames[],    //throw ComError
                      size_t         fileNo) //size of fileNames array
{
    ComPtr<IFileOperation> fileOp;
    ZEN_CHECK_COM(::CoCreateInstance(CLSID_FileOperation, //throw ComError
                                     nullptr,
                                     CLSCTX_ALL,
                                     IID_PPV_ARGS(fileOp.init())));

    // Set the operation flags.  Turn off  all UI
    // from being shown to the user during the
    // operation.  This includes error, confirmation
    // and progress dialogs.
    ZEN_CHECK_COM(fileOp->SetOperationFlags(FOF_ALLOWUNDO      | //throw ComError
                                            FOF_NOCONFIRMATION |
                                            FOF_SILENT         |
                                            FOF_NOERRORUI      |
                                            FOFX_EARLYFAILURE  |
                                            FOF_NO_CONNECTED_ELEMENTS));

    int operationCount = 0;

    for (size_t i = 0; i < fileNo; ++i)
    {
        //create file/folder item object
        ComPtr<IShellItem> psiFile;
        HRESULT hr = ::SHCreateItemFromParsingName(fileNames[i],
                                                   nullptr,
                                                   IID_PPV_ARGS(psiFile.init()));
        if (FAILED(hr))
        {
            if (hr == HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) || //file not existing anymore
                hr == HRESULT_FROM_WIN32(ERROR_PATH_NOT_FOUND))
                continue;
            throw ComError(std::wstring(L"Error calling \"SHCreateItemFromParsingName\" for file:\n") + L"\"" + fileNames[i] + L"\".", hr);
        }

        ZEN_CHECK_COM(fileOp->DeleteItem(psiFile.get(), nullptr));

        ++operationCount;
    }

    if (operationCount == 0) //calling PerformOperations() without anything to do would result in E_UNEXPECTED
        return;

    //perform actual operations
    ZEN_CHECK_COM(fileOp->PerformOperations());

    //check if errors occured: if FOFX_EARLYFAILURE is not used, PerformOperations() can return with success despite errors!
    BOOL pfAnyOperationsAborted = FALSE;
    ZEN_CHECK_COM(fileOp->GetAnyOperationsAborted(&pfAnyOperationsAborted));

    if (pfAnyOperationsAborted == TRUE)
        throw ComError(L"Operation did not complete successfully.");
}


void copyFile(const wchar_t* sourceFile, //throw ComError
              const wchar_t* targetFile)
{
    ComPtr<IFileOperation> fileOp;
    ZEN_CHECK_COM(::CoCreateInstance(CLSID_FileOperation, //throw ComError
                                     nullptr,
                                     CLSCTX_ALL,
                                     IID_PPV_ARGS(fileOp.init())));

    // Set the operation flags.  Turn off  all UI
    // from being shown to the user during the
    // operation.  This includes error, confirmation
    // and progress dialogs.
    ZEN_CHECK_COM(fileOp->SetOperationFlags(FOF_NOCONFIRMATION | //throw ComError
                                            FOF_SILENT         |
                                            FOFX_EARLYFAILURE  |
                                            FOF_NOERRORUI));
    //create source object
    ComPtr<IShellItem> psiSourceFile;
    {
        HRESULT hr = ::SHCreateItemFromParsingName(sourceFile,
                                                   nullptr,
                                                   IID_PPV_ARGS(psiSourceFile.init()));
        if (FAILED(hr))
            throw ComError(std::wstring(L"Error calling \"SHCreateItemFromParsingName\" for file:\n") + L"\"" + sourceFile + L"\".", hr);
    }

    const size_t pos = std::wstring(targetFile).find_last_of(L'\\');
    if (pos == std::wstring::npos)
        throw ComError(L"Target filename does not contain a path separator.");

    const std::wstring targetFolder(targetFile, pos);
    const std::wstring targetFileNameShort = targetFile + pos + 1;

    //create target folder object
    ComPtr<IShellItem> psiTargetFolder;
    {
        HRESULT hr = ::SHCreateItemFromParsingName(targetFolder.c_str(),
                                                   nullptr,
                                                   IID_PPV_ARGS(psiTargetFolder.init()));
        if (FAILED(hr))
            throw ComError(std::wstring(L"Error calling \"SHCreateItemFromParsingName\" for folder:\n") + L"\"" + targetFolder + L"\".", hr);
    }

    //schedule file copy operation
    ZEN_CHECK_COM(fileOp->CopyItem(psiSourceFile.get(), psiTargetFolder.get(), targetFileNameShort.c_str(), nullptr));

    //perform actual operations
    ZEN_CHECK_COM(fileOp->PerformOperations());

    //check if errors occured: if FOFX_EARLYFAILURE is not used, PerformOperations() can return with success despite errors!
    BOOL pfAnyOperationsAborted = FALSE;
    ZEN_CHECK_COM(fileOp->GetAnyOperationsAborted(&pfAnyOperationsAborted));

    if (pfAnyOperationsAborted == TRUE)
        throw ComError(L"Operation did not complete successfully.");
}


void getFolderClsid(const wchar_t* dirname, CLSID& pathCLSID) //throw ComError
{
    ComPtr<IShellFolder> desktopFolder;
    ZEN_CHECK_COM(::SHGetDesktopFolder(desktopFolder.init())); //throw ComError

    PIDLIST_RELATIVE pidlFolder = nullptr;
    ZEN_CHECK_COM(desktopFolder->ParseDisplayName(nullptr,     // [in]       HWND hwnd,
                                                  nullptr,     // [in]       IBindCtx *pbc,
                                                  const_cast<LPWSTR>(dirname), // [in]       LPWSTR pszDisplayName,
                                                  nullptr,     // [out]      ULONG *pchEaten,
                                                  &pidlFolder, // [out]      PIDLIST_RELATIVE* ppidl,
                                                  nullptr));   // [in, out]  ULONG *pdwAttributes
    ZEN_ON_SCOPE_EXIT(::ILFree(pidlFolder)); //older version: ::CoTaskMemFree

    ComPtr<IPersist> persistFolder;
    ZEN_CHECK_COM(desktopFolder->BindToObject(pidlFolder, // [in]   PCUIDLIST_RELATIVE pidl,
                                              nullptr,    // [in]   IBindCtx *pbc,
                                              IID_PPV_ARGS(persistFolder.init()))); //throw ComError

    ZEN_CHECK_COM(persistFolder->GetClassID(&pathCLSID)); //throw ComError
}


struct Win32Error
{
    Win32Error(DWORD errorCode) : errorCode_(errorCode) {}
    DWORD errorCode_;
};

std::vector<std::wstring> getLockingProcesses(const wchar_t* filename) //throw Win32Error
{
    wchar_t sessionKey[CCH_RM_SESSION_KEY + 1] = {}; //fixes two bugs: http://blogs.msdn.com/b/oldnewthing/archive/2012/02/17/10268840.aspx
    DWORD sessionHandle = 0;
    DWORD rv1 = ::RmStartSession(&sessionHandle, //__out       DWORD *pSessionHandle,
                                 0,              //__reserved  DWORD dwSessionFlags,
                                 sessionKey);    //__out       WCHAR strSessionKey[ ]
    if (rv1 != ERROR_SUCCESS)
        throw Win32Error(rv1);
    ZEN_ON_SCOPE_EXIT(::RmEndSession(sessionHandle));

    DWORD rv2 = ::RmRegisterResources(sessionHandle, //__in      DWORD dwSessionHandle,
                                      1,             //__in      UINT nFiles,
                                      &filename,     //__in_opt  LPCWSTR rgsFilenames[ ],
                                      0,             //__in      UINT nApplications,
                                      nullptr,       //__in_opt  RM_UNIQUE_PROCESS rgApplications[ ],
                                      0,             //__in      UINT nServices,
                                      nullptr);      //__in_opt  LPCWSTR rgsServiceNames[ ]
    if (rv2 != ERROR_SUCCESS)
        throw Win32Error(rv2);

    UINT procInfoSize       = 0;
    UINT procInfoSizeNeeded = 0;
    DWORD rebootReasons     = 0;
    ::RmGetList(sessionHandle, &procInfoSizeNeeded, &procInfoSize, nullptr, &rebootReasons); //get procInfoSizeNeeded
    //fails with "access denied" for C:\pagefile.sys!

    if (procInfoSizeNeeded == 0)
        return std::vector<std::wstring>();

    procInfoSize = procInfoSizeNeeded;
    std::vector<RM_PROCESS_INFO> procInfo(procInfoSize);

    DWORD rv3 = ::RmGetList(sessionHandle,       //__in         DWORD dwSessionHandle,
                            &procInfoSizeNeeded, //__out        UINT *pnProcInfoNeeded,
                            &procInfoSize,       //__inout      UINT *pnProcInfo,
                            &procInfo[0],        //__inout_opt  RM_PROCESS_INFO rgAffectedApps[ ],
                            &rebootReasons);     //__out        LPDWORD lpdwRebootReasons
    if (rv3 != ERROR_SUCCESS)
        throw Win32Error(rv3);
    procInfo.resize(procInfoSize);

    std::vector<std::wstring> output;
    for (auto iter = procInfo.begin(); iter != procInfo.end(); ++iter)
    {
        std::wstring processName = iter->strAppName;

        //try to get process path
        HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, //__in  DWORD dwDesiredAccess,
                                        false,                             //__in  BOOL bInheritHandle,
                                        iter->Process.dwProcessId);        //__in  DWORD dwProcessId
        if (hProcess)
        {
            ZEN_ON_SCOPE_EXIT(::CloseHandle(hProcess));

            FILETIME creationTime = {};
            FILETIME exitTime     = {};
            FILETIME kernelTime   = {};
            FILETIME userTime     = {};
            if (::GetProcessTimes(hProcess,      //__in   HANDLE hProcess,
                                  &creationTime, //__out  LPFILETIME lpCreationTime,
                                  &exitTime,     //__out  LPFILETIME lpExitTime,
                                  &kernelTime,   //__out  LPFILETIME lpKernelTime,
                                  &userTime))    //__out  LPFILETIME lpUserTime
                if (::CompareFileTime(&iter->Process.ProcessStartTime, &creationTime) == 0)
                {
                    DWORD bufferSize = MAX_PATH;
                    std::vector<wchar_t> buffer(bufferSize);
                    if (::QueryFullProcessImageName(hProcess,     //__in     HANDLE hProcess,
                                                    0,            //__in     DWORD dwFlags,
                                                    &buffer[0],   //__out    LPTSTR lpExeName,
                                                    &bufferSize)) //__inout  PDWORD lpdwSize
                        if (bufferSize < buffer.size())
                            processName += std::wstring(L" - ") + L"\"" + &buffer[0] + L"\"";
                }
        }
        output.push_back(processName);
    }
    return output;
}


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


bool fileop::moveToRecycleBin(const wchar_t* fileNames[], size_t fileNo) //size of fileNames array
{
    try
    {
        ::moveToRecycleBin(fileNames, fileNo); //throw ComError
        return true;
    }
    catch (const ComError& e)
    {
        lastErrorMessage.reset(new std::wstring(e.toString()));
        return false;
    }
}


bool fileop::copyFile(const wchar_t* sourceFile,
                      const wchar_t* targetFile)
{
    try
    {
        ::copyFile(sourceFile, targetFile); //throw ComError
        return true;
    }
    catch (const ComError& e)
    {
        lastErrorMessage.reset(new std::wstring(e.toString()));
        return false;
    }
}


bool fileop::checkRecycler(const wchar_t* dirname, bool& isRecycler)
{
    try
    {
        CLSID clsid = {};
        getFolderClsid(dirname, clsid); //throw ComError
        isRecycler = ::IsEqualCLSID(clsid, CLSID_RecycleBin) == TRUE; //silence perf warning
        return true;
    }
    catch (const ComError& e)
    {
        lastErrorMessage.reset(new std::wstring(e.toString()));
        return false;
    }
}


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


bool fileop::getLockingProcesses(const wchar_t* filename, const wchar_t*& procList)
{
    try
    {
        std::vector<std::wstring> result = ::getLockingProcesses(filename); //throw Win32Error

        std::wstring buffer;
        for (auto iter = result.begin(); iter != result.end(); ++iter)
        {
            buffer += *iter;
            buffer += L'\n';
        }
        if (!buffer.empty())
            buffer.resize(buffer.size() - 1); //remove last line break

        auto tmp = new wchar_t [buffer.size() + 1]; //bad_alloc ?
        ::wmemcpy(tmp, buffer.c_str(), buffer.size() + 1); //include 0-termination
        procList = tmp; //ownership passed

        return true;
    }
    catch (const Win32Error& e)
    {
        lastErrorMessage.reset(new std::wstring(formatWin32Msg(e.errorCode_)));
        return false;
    }
}


void fileop::freeString(const wchar_t* str)
{
    delete [] str;
}