// 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);
 if (!dd) {
 OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to lock DROPDESCRIPTION hGlobal");
 ReleaseStgMedium(&medium);
 return E_FAIL;
 }
 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;
}