#include "screenshot_notifications.h"
#include <atomic>
#include <shlobj.h>
#include <propkey.h>
#include <propvarutil.h>
#include <wrl/client.h>
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}};
}
namespace {
using Microsoft::WRL::ComPtr;
void check(HRESULT result) {
    if(FAILED(result)) throw std::runtime_error("Screenshot registration failed: " + std::to_string((unsigned long)result));
}
void registryString(const std::wstring &path,const wchar_t *name,const std::wstring &value)
{
    HKEY key{};
    if(RegCreateKeyExW(HKEY_CURRENT_USER,path.c_str(),0,nullptr,0,KEY_SET_VALUE,nullptr,&key,nullptr)!=ERROR_SUCCESS)
        throw std::runtime_error("Cannot register screenshot notifications");
    auto result=RegSetValueExW(key,name,0,REG_SZ,(const BYTE *)value.c_str(),DWORD((value.size()+1)*sizeof(wchar_t)));
    RegCloseKey(key);
    if(result!=ERROR_SUCCESS) throw std::runtime_error("Cannot register screenshot notification activation");
}
}
void registerScreenshotNotifications()
{
    auto exe=(std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe").wstring();
    if(!std::filesystem::exists(exe)) throw std::runtime_error("Screenshot helper is missing");
    registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",nullptr,L"\""+exe+L"\" --toast-activated");
    registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",L"ServerExecutable",exe);
    auto identity=std::wstring(L"Software\\Classes\\AppUserModelId\\")+AppId;
    registryString(identity,L"DisplayName",L"KSwitcher Screenshot");
    registryString(identity,L"IconUri",exe);
    registryString(identity,L"CustomActivator",L"{e3f78047-f0ad-49b6-9074-411931706167}");
    // A shortcut supplies desktop notification identity and its COM activation callback.
    ComPtr<IShellLinkW> link; check(CoCreateInstance(CLSID_ShellLink,nullptr,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&link)));
    check(link->SetPath(exe.c_str())); check(link->SetIconLocation(exe.c_str(),0));
    ComPtr<IPropertyStore> properties; check(link.As(&properties));
    PROPVARIANT id{}; check(InitPropVariantFromString(AppId,&id));
    auto result=properties->SetValue(PKEY_AppUserModel_ID,id); PropVariantClear(&id); check(result);
    GUID activator=ActivatorId;
    PROPVARIANT clsid{}; clsid.vt=VT_CLSID; clsid.puuid=&activator;
    check(properties->SetValue(PKEY_AppUserModel_ToastActivatorCLSID,clsid));
    check(properties->Commit());
    PWSTR programs{}; check(SHGetKnownFolderPath(FOLDERID_Programs,0,nullptr,&programs));
    auto shortcut=std::filesystem::path(programs)/L"KSwitcher Screenshot.lnk"; CoTaskMemFree(programs);
    ComPtr<IPersistFile> file; check(link.As(&file)); check(file->Save(shortcut.c_str(),TRUE));
    // Notify the shell that notification identity/activation registration changed.
    SHChangeNotify(SHCNE_ASSOCCHANGED,SHCNF_IDLIST,nullptr,nullptr);
}
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;
}

}