// Support back to Vista
#define _WIN32_WINNT _WIN32_WINNT_VISTA
#include <sdkddkver.h>

// Use WRL to define a classic COM class
#define __WRL_CLASSIC_COM__
#include <wrl.h>

#include <windows.h>
#include <shlobj.h>
#include <shlwapi.h>
#include <olectl.h>
#include <strsafe.h>

#include "pyshellext_h.h"

#define DDWM_UPDATEWINDOW (WM_USER+3)

static HINSTANCE hModule;
static CLIPFORMAT cfDropDescription;
static CLIPFORMAT cfDragWindow;

static const LPCWSTR CLASS_SUBKEY = L"Software\\Classes\\CLSID\\{BEA218D2-6950-497B-9434-61683EC065FE}";
static const LPCWSTR DRAG_MESSAGE = L"Open with %1";

using namespace Microsoft::WRL;

HRESULT FilenameListCchLengthA(LPCSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) {
    HRESULT hr = S_OK;
    size_t count = 0;
    size_t length = 0;

    while (pszSource && pszSource[0]) {
        size_t oneLength;
        hr = StringCchLengthA(pszSource, cchMax - length, &oneLength);
        if (FAILED(hr)) {
            return hr;
        }
        count += 1;
        length += oneLength + (strchr(pszSource, ' ') ? 3 : 1);
        pszSource = &pszSource[oneLength + 1];
    }

    *pcchCount = count;
    *pcchLength = length;
    return hr;
}

HRESULT FilenameListCchLengthW(LPCWSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) {
    HRESULT hr = S_OK;
    size_t count = 0;
    size_t length = 0;

    while (pszSource && pszSource[0]) {
        size_t oneLength;
        hr = StringCchLengthW(pszSource, cchMax - length, &oneLength);
        if (FAILED(hr)) {
            return hr;
        }
        count += 1;
        length += oneLength + (wcschr(pszSource, ' ') ? 3 : 1);
        pszSource = &pszSource[oneLength + 1];
    }

    *pcchCount = count;
    *pcchLength = length;
    return hr;
}

HRESULT FilenameListCchCopyA(STRSAFE_LPSTR pszDest, size_t cchDest, LPCSTR pszSource, LPCSTR pszSeparator) {
    HRESULT hr = S_OK;
    size_t count = 0;
    size_t length = 0;

    while (pszSource[0]) {
        STRSAFE_LPSTR newDest;

        hr = StringCchCopyExA(pszDest, cchDest, pszSource, &newDest, &cchDest, 0);
        if (FAILED(hr)) {
            return hr;
        }
        pszSource += (newDest - pszDest) + 1;
        pszDest = PathQuoteSpacesA(pszDest) ? newDest + 2 : newDest;

        if (pszSource[0]) {
            hr = StringCchCopyExA(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0);
            if (FAILED(hr)) {
                return hr;
            }
            pszDest = newDest;
        }
    }

    return hr;
}

HRESULT FilenameListCchCopyW(STRSAFE_LPWSTR pszDest, size_t cchDest, LPCWSTR pszSource, LPCWSTR pszSeparator) {
    HRESULT hr = S_OK;
    size_t count = 0;
    size_t length = 0;

    while (pszSource[0]) {
        STRSAFE_LPWSTR newDest;

        hr = StringCchCopyExW(pszDest, cchDest, pszSource, &newDest, &cchDest, 0);
        if (FAILED(hr)) {
            return hr;
        }
        pszSource += (newDest - pszDest) + 1;
        pszDest = PathQuoteSpacesW(pszDest) ? newDest + 2 : newDest;

        if (pszSource[0]) {
            hr = StringCchCopyExW(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0);
            if (FAILED(hr)) {
                return hr;
            }
            pszDest = newDest;
        }
    }

    return hr;
}


class PyShellExt : public RuntimeClass<
    RuntimeClassFlags<ClassicCom>,
    IDropTarget,
    IPersistFile
>
{
    LPOLESTR target, target_dir;
    DWORD target_mode;

    IDataObject *data_obj;

public:
    PyShellExt() : target(NULL), target_dir(NULL), target_mode(0), data_obj(NULL) {
        OutputDebugString(L"PyShellExt::PyShellExt");
    }

    ~PyShellExt() {
        if (target) {
            CoTaskMemFree(target);
        }
        if (target_dir) {
            CoTaskMemFree(target_dir);
        }
        if (data_obj) {
            data_obj->Release();
        }
    }

private:
    HRESULT UpdateDropDescription(IDataObject *pDataObj) {
        STGMEDIUM medium;
        FORMATETC fmt = {
            cfDropDescription,
            NULL,
            DVASPECT_CONTENT,
            -1,
            TYMED_HGLOBAL
        };

        auto hr = pDataObj->GetData(&fmt, &medium);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to get DROPDESCRIPTION format");
            return hr;
        }
        if (!medium.hGlobal) {
            OutputDebugString(L"PyShellExt::UpdateDropDescription - DROPDESCRIPTION format had NULL hGlobal");
            ReleaseStgMedium(&medium);
            return E_FAIL;
        }
        auto dd = (DROPDESCRIPTION*)GlobalLock(medium.hGlobal);
        StringCchCopy(dd->szMessage, sizeof(dd->szMessage) / sizeof(dd->szMessage[0]), DRAG_MESSAGE);
        StringCchCopy(dd->szInsert, sizeof(dd->szInsert) / sizeof(dd->szInsert[0]), PathFindFileNameW(target));
        dd->type = DROPIMAGE_MOVE;

        GlobalUnlock(medium.hGlobal);
        ReleaseStgMedium(&medium);

        return S_OK;
    }

    HRESULT GetDragWindow(IDataObject *pDataObj, HWND *phWnd) {
        HRESULT hr;
        HWND *pMem;
        STGMEDIUM medium;
        FORMATETC fmt = {
            cfDragWindow,
            NULL,
            DVASPECT_CONTENT,
            -1,
            TYMED_HGLOBAL
        };

        hr = pDataObj->GetData(&fmt, &medium);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::GetDragWindow - failed to get DragWindow format");
            return hr;
        }
        if (!medium.hGlobal) {
            OutputDebugString(L"PyShellExt::GetDragWindow - DragWindow format had NULL hGlobal");
            ReleaseStgMedium(&medium);
            return E_FAIL;
        }

        pMem = (HWND*)GlobalLock(medium.hGlobal);
        if (!pMem) {
            OutputDebugString(L"PyShellExt::GetDragWindow - failed to lock DragWindow hGlobal");
            ReleaseStgMedium(&medium);
            return E_FAIL;
        }

        *phWnd = *pMem;

        GlobalUnlock(medium.hGlobal);
        ReleaseStgMedium(&medium);

        return S_OK;
    }

    HRESULT GetArguments(IDataObject *pDataObj, LPCWSTR *pArguments) {
        HRESULT hr;
        DROPFILES *pdropfiles;

        STGMEDIUM medium;
        FORMATETC fmt = {
            CF_HDROP,
            NULL,
            DVASPECT_CONTENT,
            -1,
            TYMED_HGLOBAL
        };

        hr = pDataObj->GetData(&fmt, &medium);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::GetArguments - failed to get CF_HDROP format");
            return hr;
        }
        if (!medium.hGlobal) {
            OutputDebugString(L"PyShellExt::GetArguments - CF_HDROP format had NULL hGlobal");
            ReleaseStgMedium(&medium);
            return E_FAIL;
        }

        pdropfiles = (DROPFILES*)GlobalLock(medium.hGlobal);
        if (!pdropfiles) {
            OutputDebugString(L"PyShellExt::GetArguments - failed to lock CF_HDROP hGlobal");
            ReleaseStgMedium(&medium);
            return E_FAIL;
        }

        if (pdropfiles->fWide) {
            LPCWSTR files = (LPCWSTR)((char*)pdropfiles + pdropfiles->pFiles);
            size_t len, count;
            hr = FilenameListCchLengthW(files, 32767, &len, &count);
            if (SUCCEEDED(hr)) {
                LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
                if (args) {
                    hr = FilenameListCchCopyW(args, 32767, files, L" ");
                    if (SUCCEEDED(hr)) {
                        *pArguments = args;
                    } else {
                        CoTaskMemFree(args);
                    }
                } else {
                    hr = E_OUTOFMEMORY;
                }
            }
        } else {
            LPCSTR files = (LPCSTR)((char*)pdropfiles + pdropfiles->pFiles);
            size_t len, count;
            hr = FilenameListCchLengthA(files, 32767, &len, &count);
            if (SUCCEEDED(hr)) {
                LPSTR temp = (LPSTR)CoTaskMemAlloc(sizeof(CHAR) * (len + 1));
                if (temp) {
                    hr = FilenameListCchCopyA(temp, 32767, files, " ");
                    if (SUCCEEDED(hr)) {
                        int wlen = MultiByteToWideChar(CP_ACP, 0, temp, (int)len, NULL, 0);
                        if (wlen) {
                            LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (wlen + 1));
                            if (MultiByteToWideChar(CP_ACP, 0, temp, (int)len, args, wlen + 1)) {
                                *pArguments = args;
                            } else {
                                OutputDebugString(L"PyShellExt::GetArguments - failed to convert multi-byte to wide-char path");
                                CoTaskMemFree(args);
                                hr = E_FAIL;
                            }
                        } else {
                            OutputDebugString(L"PyShellExt::GetArguments - failed to get length of wide-char path");
                            hr = E_FAIL;
                        }
                    }
                    CoTaskMemFree(temp);
                } else {
                    hr = E_OUTOFMEMORY;
                }
            }
        }

        GlobalUnlock(medium.hGlobal);
        ReleaseStgMedium(&medium);

        return hr;
    }

    HRESULT NotifyDragWindow(HWND hwnd) {
        LRESULT res;

        if (!hwnd) {
            return S_FALSE;
        }

        res = SendMessage(hwnd, DDWM_UPDATEWINDOW, 0, NULL);

        if (res) {
            OutputDebugString(L"PyShellExt::NotifyDragWindow - failed to post DDWM_UPDATEWINDOW");
            return E_FAIL;
        }

        return S_OK;
    }

public:
    // IDropTarget implementation

    STDMETHODIMP DragEnter(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
        HWND hwnd;

        OutputDebugString(L"PyShellExt::DragEnter");

        pDataObj->AddRef();
        data_obj = pDataObj;

        *pdwEffect = DROPEFFECT_MOVE;

        if (FAILED(UpdateDropDescription(data_obj))) {
            OutputDebugString(L"PyShellExt::DragEnter - failed to update drop description");
        }
        if (FAILED(GetDragWindow(data_obj, &hwnd))) {
            OutputDebugString(L"PyShellExt::DragEnter - failed to get drag window");
        }
        if (FAILED(NotifyDragWindow(hwnd))) {
            OutputDebugString(L"PyShellExt::DragEnter - failed to notify drag window");
        }

        return S_OK;
    }

    STDMETHODIMP DragLeave() {
        return S_OK;
    }

    STDMETHODIMP DragOver(DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
        return S_OK;
    }

    STDMETHODIMP Drop(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) {
        LPCWSTR args;

        OutputDebugString(L"PyShellExt::Drop");
        *pdwEffect = DROPEFFECT_NONE;

        if (pDataObj != data_obj) {
            OutputDebugString(L"PyShellExt::Drop - unexpected data object");
            return E_FAIL;
        }

        data_obj->Release();
        data_obj = NULL;

        if (SUCCEEDED(GetArguments(pDataObj, &args))) {
            OutputDebugString(args);
            ShellExecute(NULL, NULL, target, args, target_dir, SW_NORMAL);

            CoTaskMemFree((LPVOID)args);
        } else {
            OutputDebugString(L"PyShellExt::Drop - failed to get launch arguments");
        }

        return S_OK;
    }

    // IPersistFile implementation

    STDMETHODIMP GetCurFile(LPOLESTR *ppszFileName) {
        HRESULT hr;
        size_t len;

        if (!ppszFileName) {
            return E_POINTER;
        }

        hr = StringCchLength(target, STRSAFE_MAX_CCH - 1, &len);
        if (FAILED(hr)) {
            return E_FAIL;
        }

        *ppszFileName = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
        if (!*ppszFileName) {
            return E_OUTOFMEMORY;
        }

        hr = StringCchCopy(*ppszFileName, len + 1, target);
        if (FAILED(hr)) {
            CoTaskMemFree(*ppszFileName);
            *ppszFileName = NULL;
            return E_FAIL;
        }

        return S_OK;
    }

    STDMETHODIMP IsDirty() {
        return S_FALSE;
    }

    STDMETHODIMP Load(LPCOLESTR pszFileName, DWORD dwMode) {
        HRESULT hr;
        size_t len;

        OutputDebugString(L"PyShellExt::Load");
        OutputDebugString(pszFileName);

        hr = StringCchLength(pszFileName, STRSAFE_MAX_CCH - 1, &len);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::Load - failed to get string length");
            return hr;
        }

        if (target) {
            CoTaskMemFree(target);
        }
        if (target_dir) {
            CoTaskMemFree(target_dir);
        }

        target = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
        if (!target) {
            OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY");
            return E_OUTOFMEMORY;
        }
        target_dir = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1));
        if (!target_dir) {
            OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY");
            return E_OUTOFMEMORY;
        }

        hr = StringCchCopy(target, len + 1, pszFileName);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::Load - failed to copy string");
            return hr;
        }

        hr = StringCchCopy(target_dir, len + 1, pszFileName);
        if (FAILED(hr)) {
            OutputDebugString(L"PyShellExt::Load - failed to copy string");
            return hr;
        }
        if (!PathRemoveFileSpecW(target_dir)) {
            OutputDebugStringW(L"PyShellExt::Load - failed to remove filespec from target");
            return E_FAIL;
        }

        OutputDebugString(target);
        target_mode = dwMode;
        OutputDebugString(L"PyShellExt::Load - S_OK");
        return S_OK;
    }

    STDMETHODIMP Save(LPCOLESTR pszFileName, BOOL fRemember) {
        return E_NOTIMPL;
    }

    STDMETHODIMP SaveCompleted(LPCOLESTR pszFileName) {
        return E_NOTIMPL;
    }

    STDMETHODIMP GetClassID(CLSID *pClassID) {
        *pClassID = CLSID_PyShellExt;
        return S_OK;
    }
};

CoCreatableClass(PyShellExt);

STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) {
    return Module<InProc>::GetModule().GetClassObject(rclsid, riid, ppv);
}

STDAPI DllCanUnloadNow() {
    return Module<InProc>::GetModule().Terminate() ? S_OK : S_FALSE;
}

STDAPI DllRegisterServer() {
    LONG res;
    SECURITY_ATTRIBUTES secattr = { sizeof(SECURITY_ATTRIBUTES), NULL, FALSE };
    LPSECURITY_ATTRIBUTES psecattr = NULL;
    HKEY key, ipsKey;
    WCHAR modname[MAX_PATH];
    DWORD modname_len;

    OutputDebugString(L"PyShellExt::DllRegisterServer");
    if (!hModule) {
        OutputDebugString(L"PyShellExt::DllRegisterServer - module handle was not set");
        return SELFREG_E_CLASS;
    }
    modname_len = GetModuleFileName(hModule, modname, MAX_PATH);
    if (modname_len == 0 ||
        (modname_len == MAX_PATH && GetLastError() == ERROR_INSUFFICIENT_BUFFER)) {
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to get module file name");
        return SELFREG_E_CLASS;
    }

    DWORD disp;
    res = RegCreateKeyEx(HKEY_LOCAL_MACHINE, CLASS_SUBKEY, 0, NULL, 0,
        KEY_ALL_ACCESS, psecattr, &key, &disp);
    if (res == ERROR_ACCESS_DENIED) {
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to write per-machine registration. Attempting per-user instead.");
        res = RegCreateKeyEx(HKEY_CURRENT_USER, CLASS_SUBKEY, 0, NULL, 0,
            KEY_ALL_ACCESS, psecattr, &key, &disp);
    }
    if (res != ERROR_SUCCESS) {
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create class key");
        return SELFREG_E_CLASS;
    }

    res = RegCreateKeyEx(key, L"InProcServer32", 0, NULL, 0,
        KEY_ALL_ACCESS, psecattr, &ipsKey, NULL);
    if (res != ERROR_SUCCESS) {
        RegCloseKey(key);
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create InProcServer32 key");
        return SELFREG_E_CLASS;
    }

    res = RegSetValueEx(ipsKey, NULL, 0,
        REG_SZ, (LPBYTE)modname, modname_len * sizeof(modname[0]));

    if (res != ERROR_SUCCESS) {
        RegCloseKey(ipsKey);
        RegCloseKey(key);
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set server path");
        return SELFREG_E_CLASS;
    }

    res = RegSetValueEx(ipsKey, L"ThreadingModel", 0,
        REG_SZ, (LPBYTE)(L"Apartment"), sizeof(L"Apartment"));

    RegCloseKey(ipsKey);
    RegCloseKey(key);
    if (res != ERROR_SUCCESS) {
        OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set threading model");
        return SELFREG_E_CLASS;
    }

    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);

    OutputDebugString(L"PyShellExt::DllRegisterServer - S_OK");
    return S_OK;
}

STDAPI DllUnregisterServer() {
    LONG res_lm, res_cu;

    res_lm = RegDeleteTree(HKEY_LOCAL_MACHINE, CLASS_SUBKEY);
    if (res_lm != ERROR_SUCCESS && res_lm != ERROR_FILE_NOT_FOUND) {
        OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-machine registration");
        return SELFREG_E_CLASS;
    }

    res_cu = RegDeleteTree(HKEY_CURRENT_USER, CLASS_SUBKEY);
    if (res_cu != ERROR_SUCCESS && res_cu != ERROR_FILE_NOT_FOUND) {
        OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-user registration");
        return SELFREG_E_CLASS;
    }

    if (res_lm == ERROR_FILE_NOT_FOUND && res_cu == ERROR_FILE_NOT_FOUND) {
        OutputDebugString(L"PyShellExt::DllUnregisterServer - extension was not registered");
        return SELFREG_E_CLASS;
    }

    SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL);

    OutputDebugString(L"PyShellExt::DllUnregisterServer - S_OK");
    return S_OK;
}

STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) {
    if (reason == DLL_PROCESS_ATTACH) {
        hModule = hinst;

        cfDropDescription = RegisterClipboardFormat(CFSTR_DROPDESCRIPTION);
        if (!cfDropDescription) {
            OutputDebugString(L"PyShellExt::DllMain - failed to get CFSTR_DROPDESCRIPTION format");
        }
        cfDragWindow = RegisterClipboardFormat(L"DragWindow");
        if (!cfDragWindow) {
            OutputDebugString(L"PyShellExt::DllMain - failed to get DragWindow format");
        }

        DisableThreadLibraryCalls(hinst);
    }
    return TRUE;
}