#include "screenshot_notifications.h"
#include <atomic>
namespace ks {
namespace {
constexpr wchar_t AppId[]=L"HrynCo.KSwitcher.Screenshot";
constexpr GUID ActivatorId{0xe3f78047,0xf0ad,0x49b6,{0x90,0x74,0x41,0x19,0x31,0x70,0x61,0x67}};
}
bool validScreenshotToken(const std::wstring &id)
{
    if(id.size()!=36) return false;
    for(size_t i=0;i<id.size();++i) {
        if(i==8 || i==13 || i==18 || i==23) { if(id[i]!=L'-') return false; }
        else if(!((id[i]>=L'0' && id[i]<=L'9') || (id[i]>=L'a' && id[i]<=L'f') || (id[i]>=L'A' && id[i]<=L'F'))) return false;
    }
    return true;
}

namespace {
// Windows SDK INotificationActivationCallback ABI (notificationactivationcallback.h).
constexpr GUID CallbackId{0x53e31837,0x6600,0x4a81,{0x93,0x95,0x75,0xcf,0xfe,0x74,0x6f,0x94}};
struct NotificationInput { LPCWSTR key,value; };
struct NotificationCallback : IUnknown {
    virtual HRESULT STDMETHODCALLTYPE Activate(LPCWSTR,LPCWSTR,const NotificationInput *,ULONG)=0;
};
class Activator final : public NotificationCallback {
    std::atomic<ULONG> references{1};
    DWORD thread=GetCurrentThreadId();
    std::wstring target;
    bool quit;
public:
    Activator(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
        if(!out) return E_POINTER;
        *out=nullptr;
        if(id!=IID_IUnknown && id!=CallbackId) return E_NOINTERFACE;
        *out=static_cast<NotificationCallback *>(this); AddRef(); return S_OK;
    }
    ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
    ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
    HRESULT STDMETHODCALLTYPE Activate(LPCWSTR app,LPCWSTR args,const NotificationInput *,ULONG) override {
        if(!app || std::wstring(app)!=AppId || !args || !validScreenshotToken(args)) return E_INVALIDARG;
        auto path=target; auto command=L"\""+path+L"\" --open \"kswitcher-screenshot:"+args+L"\"";
        STARTUPINFOW startup{sizeof(startup)}; PROCESS_INFORMATION process{};
        auto ok=CreateProcessW(path.c_str(),command.data(),nullptr,nullptr,FALSE,CREATE_SUSPENDED,nullptr,nullptr,&startup,&process);
        auto result=ok?S_OK:HRESULT_FROM_WIN32(GetLastError());
        if(ok) {
            AllowSetForegroundWindow(process.dwProcessId); ResumeThread(process.hThread);
            CloseHandle(process.hThread); CloseHandle(process.hProcess);
        }
        if(quit) PostThreadMessageW(thread,WM_QUIT,0,0);
        return result;
    }
};
class ActivatorFactory final : public IClassFactory {
    std::atomic<ULONG> references{1};
    std::wstring target;
    bool quit;
public:
    ActivatorFactory(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
        if(!out) return E_POINTER;
        *out=nullptr;
        if(id!=IID_IUnknown && id!=IID_IClassFactory) return E_NOINTERFACE;
        *out=static_cast<IClassFactory *>(this); AddRef(); return S_OK;
    }
    ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
    ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
    HRESULT STDMETHODCALLTYPE CreateInstance(IUnknown *outer,REFIID id,void **out) override {
        if(outer) return CLASS_E_NOAGGREGATION;
        auto object=new Activator(target,quit); auto result=object->QueryInterface(id,out); object->Release(); return result;
    }
    HRESULT STDMETHODCALLTYPE LockServer(BOOL) override { return S_OK; }
};
}
int runScreenshotToastActivator()
{
    if(FAILED(CoInitializeEx(nullptr,COINIT_APARTMENTTHREADED))) return 1;
    MSG message{}; PeekMessageW(&message,nullptr,0,0,PM_NOREMOVE);
    auto factory=new ActivatorFactory(executable(),true); DWORD registration{};
    auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,&registration);
    factory->Release();
    if(SUCCEEDED(result)) {
        auto timer=SetTimer(nullptr,0,60000,[](HWND,UINT,UINT_PTR,DWORD){PostQuitMessage(0);});
        while(GetMessageW(&message,nullptr,0,0)>0) { TranslateMessage(&message); DispatchMessageW(&message); }
        if(timer) KillTimer(nullptr,timer);
        CoRevokeClassObject(registration);
    }
    CoUninitialize(); return FAILED(result)?1:0;
}
DWORD startScreenshotActivationListener()
{
    auto target=std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe";
    auto factory=new ActivatorFactory(target.wstring(),false);
    DWORD registration{};
    auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,&registration);
    factory->Release();
    return SUCCEEDED(result)?registration:0;
}

}