#include "FrameworkContextMenuExt.h"
#include "Reg.h"
#include <shlwapi.h>
#include <Strsafe.h>
#pragma comment(lib, "shlwapi.lib")

extern HINSTANCE    g_hInst;
extern long         g_cDllRef;

static const int IDM_SENDTOPSCAN = 0;

static const LPWSTR PScanCommandDisplayName = L"Send To PScan";

static const LPWSTR PScanPipeName = L"\\\\.\\pipe\\PScanPipe";
static const LPWSTR PScanMessageAddFiles = L"AddFiles";


FrameWorkContextMenuExt::FrameWorkContextMenuExt()
    : m_cRef(1)
{
    InterlockedIncrement(&g_cDllRef);
}

FrameWorkContextMenuExt::~FrameWorkContextMenuExt()
{
    InterlockedDecrement(&g_cDllRef);
}

void WriteUTF7Integer(HANDLE _hPipe, unsigned _value)
{
    unsigned long written = 0;
    if(_value < 128)
    {
        unsigned char bValue = static_cast<unsigned char>(_value);
        WriteFile(_hPipe, &bValue, sizeof(bValue), &written, nullptr);
    }
    else
    {
        unsigned short sValue = static_cast<unsigned short>(((_value << 1) & ~0xFF) | 0x80 | (_value & 0x7F));
        WriteFile(_hPipe, &sValue, sizeof(sValue), &written, nullptr);
    }
}

void WriteString(HANDLE _hPipe, LPWSTR _string)
{   
    char filename[1024];
    WideCharToMultiByte(CP_UTF8, 0, _string, -1, filename, ARRAYSIZE(filename), nullptr, nullptr);

    unsigned length = static_cast<unsigned>(strlen(filename));
    WriteUTF7Integer(_hPipe, length);
    unsigned long written = 0;
    WriteFile(_hPipe, filename, length, &written, nullptr);
}

void FrameWorkContextMenuExt::SendFilesToPScan()
{
    HANDLE hPipe = CreateFile(PScanPipeName, GENERIC_WRITE, 0, nullptr, OPEN_EXISTING, 0, nullptr);
    if(hPipe == INVALID_HANDLE_VALUE)
    {
        LaunchPScan();

        // Retry until the pipe is created
        for(unsigned i = 0; i < 10; ++i)
        {
            hPipe = CreateFile(PScanPipeName, GENERIC_WRITE, 0, nullptr, OPEN_EXISTING, 0, nullptr);
            if(hPipe != INVALID_HANDLE_VALUE)
                break;

            Sleep(500);
        }
        if(hPipe == INVALID_HANDLE_VALUE)
            return;
    }

    WriteString(hPipe, PScanMessageAddFiles);
    
    unsigned numFiles = static_cast<unsigned>(m_selectedFiles.size());
    unsigned long written = 0;
    WriteFile(hPipe, &numFiles, sizeof(numFiles), &written, nullptr);

    for(unsigned i = 0; i < numFiles; ++i)
    {
        WriteString(hPipe, m_selectedFiles[i].Name);
    }

    CloseHandle(hPipe);
}

unsigned FrameWorkContextMenuExt::GetDataRoot(wchar_t *_dest, unsigned _dstSize) const
{
    wchar_t rootDir[MAX_PATH];
    if(FAILED(GetRegistryKeyAndValue(HKEY_CURRENT_USER, L"Software\\Ubisoft\\UbiArt\\UAF", L"rootDir", rootDir, sizeof(rootDir))))
        return 0;

    StringCchCatW(rootDir, ARRAYSIZE(rootDir), L"\\lastvalidroot.txt");
    
    HANDLE hFile = CreateFile(rootDir, GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if(hFile == INVALID_HANDLE_VALUE)
        return 0;

    char dataroot[MAX_PATH] = {0};
    unsigned long read = 0;
    ReadFile(hFile, dataroot, MAX_PATH - 1, &read, nullptr);

    CloseHandle(hFile);

    MultiByteToWideChar(CP_UTF8, 0, dataroot, -1, _dest, _dstSize);
    return read;
}

void FrameWorkContextMenuExt::LaunchPScan() const
{
    wchar_t rootDir[MAX_PATH];
    if(FAILED(GetRegistryKeyAndValue(HKEY_CURRENT_USER, L"Software\\Ubisoft\\UbiArt\\UAF", L"rootDir", rootDir, sizeof(rootDir))))
        return;

    wchar_t appPath[MAX_PATH];
    StringCchCopy(appPath, ARRAYSIZE(appPath), rootDir);
    StringCchCatW(appPath, ARRAYSIZE(appPath), L"\\pscansharpui.exe");

    STARTUPINFO si;
    ZeroMemory(&si, sizeof(si));
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi;

    if(SUCCEEDED(CreateProcess(appPath, L"", nullptr, nullptr, FALSE, 0, nullptr, rootDir, &si, &pi)))
    {
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    }
}

#pragma region IUnknown

IFACEMETHODIMP FrameWorkContextMenuExt::QueryInterface(REFIID riid, void **ppv)
{
    static const QITAB qit[] = 
    {
        QITABENT(FrameWorkContextMenuExt, IContextMenu),
        QITABENT(FrameWorkContextMenuExt, IShellExtInit), 
        { 0 },
    };
    return QISearch(this, qit, riid, ppv);
}

IFACEMETHODIMP_(ULONG) FrameWorkContextMenuExt::AddRef()
{
    return InterlockedIncrement(&m_cRef);
}

IFACEMETHODIMP_(ULONG) FrameWorkContextMenuExt::Release()
{
    ULONG cRef = InterlockedDecrement(&m_cRef);
    if (!cRef)
        delete this;

    return cRef;
}

#pragma endregion

#pragma region IShellExtInit

IFACEMETHODIMP FrameWorkContextMenuExt::Initialize(LPCITEMIDLIST pidlFolder, LPDATAOBJECT pDataObj, HKEY hKeyProgID)
{
    if (NULL == pDataObj)
    {
        return E_INVALIDARG;
    }

    HRESULT hr = E_FAIL;

    FORMATETC fe = { CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL };
    STGMEDIUM stm;

    wchar_t dataRoot[MAX_PATH];
    unsigned dataRootLen = GetDataRoot(dataRoot, ARRAYSIZE(dataRoot));
    for(unsigned i = 0; i < dataRootLen; ++i)
    {
        if(dataRoot[i] == L'/')
            dataRoot[i] = L'\\';
    }

    if (SUCCEEDED(pDataObj->GetData(&fe, &stm)))
    {
        HDROP hDrop = static_cast<HDROP>(GlobalLock(stm.hGlobal));
        if (hDrop != NULL)
        {
            UINT nFiles = DragQueryFile(hDrop, 0xFFFFFFFF, NULL, 0);
            for(unsigned i = 0; i < nFiles; ++i)
            {
                File selectedFile;
                if (0 != DragQueryFile(hDrop, i, selectedFile.Name, ARRAYSIZE(selectedFile.Name)))
                {
                    if(dataRootLen > 0 && _wcsnicmp(dataRoot, selectedFile.Name, dataRootLen))
                        continue;

                    m_selectedFiles.push_back(selectedFile);
                    hr = S_OK;
                }
            }

            GlobalUnlock(stm.hGlobal);
        }

        ReleaseStgMedium(&stm);
    }

    return hr;
}

#pragma endregion

#pragma region IContextMenu

IFACEMETHODIMP FrameWorkContextMenuExt::QueryContextMenu(HMENU hMenu, UINT indexMenu, UINT idCmdFirst, UINT idCmdLast, UINT uFlags)
{
    if (CMF_DEFAULTONLY & uFlags)
        return MAKE_HRESULT(SEVERITY_SUCCESS, 0, USHORT(0));

    MENUITEMINFO mii = { sizeof(mii) };
    mii.fMask = MIIM_BITMAP | MIIM_STRING | MIIM_FTYPE | MIIM_ID | MIIM_STATE;
    mii.wID = idCmdFirst + IDM_SENDTOPSCAN;
    mii.fType = MFT_STRING;
    mii.dwTypeData = PScanCommandDisplayName;
    mii.fState = MFS_ENABLED;
    mii.hbmpItem = nullptr;
    if (!InsertMenuItem(hMenu, indexMenu, TRUE, &mii))
        return HRESULT_FROM_WIN32(GetLastError());

    // Add a separator.
    MENUITEMINFO sep = { sizeof(sep) };
    sep.fMask = MIIM_TYPE;
    sep.fType = MFT_SEPARATOR;
    if (!InsertMenuItem(hMenu, indexMenu + 1, TRUE, &sep))
        return HRESULT_FROM_WIN32(GetLastError());

    return MAKE_HRESULT(SEVERITY_SUCCESS, 0, USHORT(IDM_SENDTOPSCAN + 1));
}


IFACEMETHODIMP FrameWorkContextMenuExt::InvokeCommand(LPCMINVOKECOMMANDINFO pici)
{
    BOOL fUnicode = FALSE;

    if (pici->cbSize == sizeof(CMINVOKECOMMANDINFOEX))
    {
        if (pici->fMask & CMIC_MASK_UNICODE)
        {
            fUnicode = TRUE;
        }
    }

    if (!fUnicode && HIWORD(pici->lpVerb))
    {
        if (StrCmpIA(pici->lpVerb, "sendtopscan") == 0)
        {
            SendFilesToPScan();
        }
        else
        {
            return E_FAIL;
        }
    }
    else if (fUnicode && HIWORD(((CMINVOKECOMMANDINFOEX*)pici)->lpVerbW))
    {
        if (StrCmpIW(((CMINVOKECOMMANDINFOEX*)pici)->lpVerbW, L"sendtopscan") == 0)
        {
            SendFilesToPScan();
        }
        else
        {
            return E_FAIL;
        }
    }
    else
    {
        if (LOWORD(pici->lpVerb) == IDM_SENDTOPSCAN)
        {
            SendFilesToPScan();
        }
        else
        {
            return E_FAIL;
        }
    }

    return S_OK;
}

IFACEMETHODIMP FrameWorkContextMenuExt::GetCommandString(UINT_PTR idCommand, UINT uFlags, UINT *pwReserved, LPSTR pszName, UINT cchMax)
{
    HRESULT hr = E_INVALIDARG;

    if (idCommand == IDM_SENDTOPSCAN)
    {
        switch (uFlags)
        {
        case GCS_HELPTEXTW:
            hr = StringCchCopy(reinterpret_cast<PWSTR>(pszName), cchMax, PScanCommandDisplayName);
            break;

        case GCS_VERBW:
            hr = StringCchCopy(reinterpret_cast<PWSTR>(pszName), cchMax, PScanCommandDisplayName);
            break;

        default:
            hr = S_OK;
        }
    }

    return hr;
}

#pragma endregion