hrynco-kswitcher/src/screenshot_main.cpp

#include "screenshot.h"
#include "screenshot_notifications.h"
#include <commctrl.h>
#include <shellapi.h>
#include <roapi.h>

namespace {
constexpr wchar_t ControllerClass[]=L"KSwitcher.Screenshot.Controller";
LRESULT CALLBACK controller(HWND window,UINT message,WPARAM wp,LPARAM lp)
{
    auto host=(ks::ScreenshotHost *)GetWindowLongPtrW(window,GWLP_USERDATA);
    if(message==WM_NCCREATE) {
        host=(ks::ScreenshotHost *)((CREATESTRUCTW *)lp)->lpCreateParams;
        SetWindowLongPtrW(window,GWLP_USERDATA,(LONG_PTR)host);
    }
    if(host && message==ks::ScreenshotRequest) {
        host->autoOpenEditor=(wp&1)!=0;
        ks::setLoggingEnabled((wp&2)!=0);
        try { ks::showScreenshot(*host); }
        catch(const std::exception &e) {
            MessageBoxW(window,ks::wide(e.what()).c_str(),L"Screen capture",MB_OK|MB_ICONERROR);
            PostMessageW(window,ks::ScreenshotIdle,0,0);
        }
        return TRUE;
    }
    if(host && message==WM_COPYDATA) {
        auto data=(COPYDATASTRUCT *)lp;
        if(!data || data->dwData!=73 || data->cbData!=37*sizeof(wchar_t) || !data->lpData) return FALSE;
        std::wstring token((const wchar_t *)data->lpData,36);
        if(!ks::validScreenshotToken(token)) return FALSE;
        try { ks::showScreenshotEditor(*host,ks::loadNotifiedScreenshot(token)); }
        catch(const std::exception &e) {
            MessageBoxW(window,ks::wide(e.what()).c_str(),L"Screenshot",MB_OK|MB_ICONINFORMATION);
            PostMessageW(window,ks::ScreenshotIdle,0,0);
        }
        return TRUE;
    }
    if(host && message==ks::ScreenshotIdle && !host->windows) {
        DestroyWindow(window);
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProcW(window,message,wp,lp);
}
}

int WINAPI wWinMain(HINSTANCE instance,HINSTANCE,PWSTR,int)
{
    std::wstring token;
    bool autoOpenEditor=false, writeLog=false;
    int argc{}; auto argv=CommandLineToArgvW(GetCommandLineW(),&argc);
    if(argc>1 && std::wstring(argv[argc-1])==L"--write-log") {
        writeLog=true; --argc;
        ks::setLoggingEnabled(true);
    }
    if(argc>=2 && std::wstring(argv[1])==L"--toast-activated") {
        LocalFree(argv); return ks::runScreenshotToastActivator();
    }
    if(argc==2 && std::wstring(argv[1])==L"--register-notifications") {
        LocalFree(argv); CoInitializeEx(nullptr,COINIT_APARTMENTTHREADED);
        int result=0; try { ks::registerScreenshotNotifications(); } catch(...) { result=1; }
        CoUninitialize(); return result;
    }
    if(argc==2 && (std::wstring(argv[1])==L"--auto-open-editor=0" ||
                   std::wstring(argv[1])==L"--auto-open-editor=1")) {
        autoOpenEditor=std::wstring(argv[1])==L"--auto-open-editor=1";
    } else if(argc>1) {
        if(argc==3 && std::wstring(argv[1])==L"--open") {
            std::wstring uri=argv[2]; const std::wstring prefix=L"kswitcher-screenshot:";
            if(uri.starts_with(prefix)) token=uri.substr(prefix.size());
        }
        if(!ks::validScreenshotToken(token)) { LocalFree(argv); return 1; }
    }
    LocalFree(argv);
    COPYDATASTRUCT open{73,DWORD((token.size()+1)*sizeof(wchar_t)),token.data()};
    HANDLE mutex=CreateMutexW(nullptr,FALSE,L"Local\\KSwitcher.Screenshot.Singleton");
    if(!mutex) return 1;
    // Wait through startup/shutdown races; the controller accepts requests only while alive.
    bool acquired=false;
    for(int attempt=0;attempt<100;++attempt) {
        auto wait=WaitForSingleObject(mutex,0);
        if(wait==WAIT_OBJECT_0 || wait==WAIT_ABANDONED) { acquired=true; break; }
        if(auto existing=FindWindowW(ControllerClass,nullptr)) {
            DWORD pid{}; GetWindowThreadProcessId(existing,&pid);
            AllowSetForegroundWindow(pid);
            DWORD_PTR accepted{};
            if(SendMessageTimeoutW(existing,token.empty()?ks::ScreenshotRequest:WM_COPYDATA,token.empty()?(autoOpenEditor?1:0)|(writeLog?2:0):0,
                token.empty()?0:(LPARAM)&open,SMTO_ABORTIFHUNG,2000,&accepted) && accepted) {
                CloseHandle(mutex); return 0;
            }
        }
        Sleep(50);
    }
    if(!acquired) { CloseHandle(mutex); return 1; }
    RoInitialize(RO_INIT_SINGLETHREADED);
    INITCOMMONCONTROLSEX controls{sizeof(controls),ICC_STANDARD_CLASSES|ICC_BAR_CLASSES};
    InitCommonControlsEx(&controls);
    ks::ScreenshotHost host{instance};
    host.notify=ks::notifyScreenshot;
    WNDCLASSW cls{}; cls.hInstance=instance; cls.lpfnWndProc=controller; cls.lpszClassName=ControllerClass;
    RegisterClassW(&cls);
    host.window=CreateWindowW(ControllerClass,L"KSwitcher Screenshot",0,0,0,0,0,nullptr,nullptr,instance,&host);
    int result=1;
    if(host.window) {
        // Toast protocol activation may run unelevated while the editor was launched by an elevated resident.
        // WM_COPYDATA accepts only a validated cache GUID, never a path or executable command.
        ChangeWindowMessageFilterEx(host.window,WM_COPYDATA,MSGFLT_ALLOW,nullptr);
        if(token.empty()) PostMessageW(host.window,ks::ScreenshotRequest,(autoOpenEditor?1:0)|(writeLog?2:0),0);
        else SendMessageW(host.window,WM_COPYDATA,0,(LPARAM)&open);
        MSG message{};
        while((result=GetMessageW(&message,nullptr,0,0))>0) {
            TranslateMessage(&message); DispatchMessageW(&message);
        }
        result=result<0?1:0;
    }
    RoUninitialize();
    ReleaseMutex(mutex); CloseHandle(mutex);
    return result;
}