#include "screenshot_notifications.h"
#include "ocr.h"
#include <algorithm>
#include <chrono>
#include <cstdio>
#include <aclapi.h>
#include <sddl.h>
#include <shlobj.h>
#include <propkey.h>
#include <propvarutil.h>
#include <roapi.h>
#include <winstring.h>
#include <wincodec.h>
#ifdef __MINGW32__
// This SDK aliases BYTE and boolean, producing duplicate template specializations.
// Notifications use the boolean specialization; IReference<BYTE> is unused here.
#define ____FIReference_1_BYTE_INTERFACE_DEFINED__
#endif
#include <windows.ui.notifications.h>
#include <windows.data.xml.dom.h>
#include <wrl/client.h>

namespace ks {
namespace {
using Microsoft::WRL::ComPtr;
namespace Notifications=ABI::Windows::UI::Notifications;
namespace Xml=ABI::Windows::Data::Xml::Dom;
constexpr wchar_t AppId[]=L"HrynCo.KSwitcher.Screenshot";
std::string notificationError;
void check(HRESULT result) { if(FAILED(result)) { char message[100]; snprintf(message,sizeof(message),"Windows screenshot notification failed (0x%08lx)",(unsigned long)result); throw std::runtime_error(message); } }
struct String {
    HSTRING value{};
    explicit String(const std::wstring &s) { check(WindowsCreateString(s.c_str(),(UINT32)s.size(),&value)); }
    ~String() { WindowsDeleteString(value); }
};
std::filesystem::path cacheDirectory()
{
    PWSTR folder{}; check(SHGetKnownFolderPath(FOLDERID_LocalAppData,0,nullptr,&folder));
    auto path=std::filesystem::path(folder)/L"KSwitcher"/L"ScreenshotNotifications";
    CoTaskMemFree(folder); std::filesystem::create_directories(path); return path;
}
std::filesystem::path previewDirectory()
{
    HANDLE token{};
    if(!OpenProcessToken(GetCurrentProcess(),TOKEN_QUERY,&token)) check(HRESULT_FROM_WIN32(GetLastError()));
    DWORD size{}; GetTokenInformation(token,TokenUser,nullptr,0,&size);
    std::vector<BYTE> data(size);
    BOOL read=GetTokenInformation(token,TokenUser,data.data(),size,&size);
    DWORD error=read?ERROR_SUCCESS:GetLastError(); CloseHandle(token); check(HRESULT_FROM_WIN32(error));
    auto user=reinterpret_cast<TOKEN_USER *>(data.data())->User.Sid;
    LPWSTR sidText{}; if(!ConvertSidToStringSidW(user,&sidText)) check(HRESULT_FROM_WIN32(GetLastError()));
    std::wstring sid(sidText); LocalFree(sidText);
    PWSTR folder{}; check(SHGetKnownFolderPath(FOLDERID_ProgramData,0,nullptr,&folder));
    auto path=std::filesystem::path(folder)/(L"KSwitcher.ScreenshotPreviews-"+sid); CoTaskMemFree(folder);
    // The notification broker cannot load previews from some profile directories.
    // Keep this cache outside the profile, private to the current user and SYSTEM.
    auto sddl=L"O:"+sid+L"D:P(A;OICI;FA;;;SY)(A;OICI;FA;;;"+sid+L")";
    PSECURITY_DESCRIPTOR security{};
    if(!ConvertStringSecurityDescriptorToSecurityDescriptorW(sddl.c_str(),SDDL_REVISION_1,&security,nullptr))
        check(HRESULT_FROM_WIN32(GetLastError()));
    SECURITY_ATTRIBUTES attributes{sizeof(attributes),security,FALSE};
    BOOL created=CreateDirectoryW(path.c_str(),&attributes);
    error=created?ERROR_SUCCESS:GetLastError();
    LocalFree(security);
    if(error!=ERROR_ALREADY_EXISTS) check(HRESULT_FROM_WIN32(error));
    // Do not use a pre-created directory or junction controlled by another user.
    DWORD flags=GetFileAttributesW(path.c_str());
    if(flags==INVALID_FILE_ATTRIBUTES || !(flags&FILE_ATTRIBUTE_DIRECTORY) || (flags&FILE_ATTRIBUTE_REPARSE_POINT))
        throw std::runtime_error("Invalid screenshot preview directory");
    PSID owner{}; PSECURITY_DESCRIPTOR existing{};
    error=GetNamedSecurityInfoW(const_cast<wchar_t *>(path.c_str()),SE_FILE_OBJECT,OWNER_SECURITY_INFORMATION,
        &owner,nullptr,nullptr,nullptr,&existing);
    check(HRESULT_FROM_WIN32(error));
    bool owned=EqualSid(owner,user); LocalFree(existing);
    if(!owned) throw std::runtime_error("Screenshot preview directory belongs to another user");
    return path;
}
void cleanCache(const std::filesystem::path &path)
{
    std::vector<std::filesystem::directory_entry> files;
    for(const auto &file:std::filesystem::directory_iterator(path))
        if(file.is_regular_file() && validScreenshotToken(file.path().stem().wstring()) &&
           (file.path().extension()==L".capture" || file.path().extension()==L".png")) files.push_back(file);
    std::sort(files.begin(),files.end(),[](const auto &a,const auto &b){return a.last_write_time()>b.last_write_time();});
    uintmax_t size=0;
    for(size_t i=0;i<files.size();++i) {
        size+=files[i].file_size();
        if(i>=40 || size>256*1024*1024 || std::filesystem::file_time_type::clock::now()-files[i].last_write_time()>std::chrono::hours(24)) {
            std::error_code ignored; std::filesystem::remove(files[i].path(),ignored);
        }
    }
}
void savePreview(const Clip &clip,const std::filesystem::path &path)
{
    auto pixels=decodeImage(clip,360,220);
    for(size_t i=3;i<pixels.bgra.size();i+=4) pixels.bgra[i]=255;
    ComPtr<IWICImagingFactory> factory; check(CoCreateInstance(CLSID_WICImagingFactory,nullptr,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&factory)));
    ComPtr<IWICStream> stream; check(factory->CreateStream(&stream)); check(stream->InitializeFromFilename(path.c_str(),GENERIC_WRITE));
    ComPtr<IWICBitmapEncoder> encoder; check(factory->CreateEncoder(GUID_ContainerFormatPng,nullptr,&encoder));
    check(encoder->Initialize(stream.Get(),WICBitmapEncoderNoCache));
    ComPtr<IWICBitmapFrameEncode> frame; check(encoder->CreateNewFrame(&frame,nullptr)); check(frame->Initialize(nullptr));
    check(frame->SetSize(pixels.width,pixels.height));
    auto format=GUID_WICPixelFormat32bppBGRA; check(frame->SetPixelFormat(&format));
    if(format!=GUID_WICPixelFormat32bppBGRA) throw std::runtime_error("Unsupported preview pixel format");
    check(frame->WritePixels(pixels.height,pixels.width*4,(UINT)pixels.bgra.size(),pixels.bgra.data()));
    check(frame->Commit()); check(encoder->Commit());
}
std::wstring escapeXml(const std::wstring &s)
{
    std::wstring result;
    for(auto ch:s) { if(ch==L'&') result+=L"&amp;"; else if(ch==L'\"') result+=L"&quot;"; else if(ch==L'<') result+=L"&lt;"; else result+=ch; }
    return result;
}
}
std::wstring cacheScreenshotForEditor(const Clip &clip)
{
    // Normalize PNG/DIBV5 and other supported clipboard images at full resolution.
    auto pixels=decodeImage(clip);
    if(pixels.bgra.empty() || pixels.bgra.size()>128*1024*1024-sizeof(BITMAPINFOHEADER))
        throw std::runtime_error("Image is too large for the screenshot editor");
    auto dib=imageDib(pixels);
    auto directory=cacheDirectory(); cleanCache(directory);
    GUID guid{}; check(CoCreateGuid(&guid)); wchar_t value[40]{}; StringFromGUID2(guid,value,40);
    std::wstring id(value+1,36);
    atomicWrite(directory/(id+L".capture"),protect(dib.data));
    return id;
}
Clip loadNotifiedScreenshot(const std::wstring &id)
{
    if(!validScreenshotToken(id)) throw std::runtime_error("Invalid screenshot notification");
    auto directory=cacheDirectory(); cleanCache(directory);
    auto path=directory/(id+L".capture");
    if(!std::filesystem::exists(path)) throw std::runtime_error("This screenshot notification has expired. Take a new screenshot.");
    return Clip{{{CF_DIB,L"",unprotect(readFile(path,160*1024*1024))}}};
}
bool notifyScreenshot(const Clip &clip)
{
    notificationError.clear();
    const char *stage="registration";
    try {
        registerScreenshotNotifications();
        check(SetCurrentProcessExplicitAppUserModelID(AppId));
        stage="notification manager";
        ComPtr<Notifications::IToastNotificationManagerStatics> manager;
        String managerClass(L"Windows.UI.Notifications.ToastNotificationManager");
        check(RoGetActivationFactory(managerClass.value,IID_PPV_ARGS(&manager)));
        ComPtr<Notifications::IToastNotifier> notifier; String appId(AppId);
        stage="notifier";
        check(manager->CreateToastNotifierWithId(appId.value,&notifier));
        stage="settings";
        Notifications::NotificationSetting setting{};
        auto settingResult=notifier->get_Setting(&setting);
        // A first-time sender has no settings record until its first notification.
        if(settingResult!=HRESULT_FROM_WIN32(ERROR_NOT_FOUND)) check(settingResult);
        if(setting!=Notifications::NotificationSetting_Enabled) throw std::runtime_error("Screenshot notifications are disabled in Windows settings");
        stage="cache";
        auto directory=cacheDirectory(); cleanCache(directory);
        GUID guid{}; check(CoCreateGuid(&guid)); wchar_t value[40]{}; StringFromGUID2(guid,value,40);
        std::wstring id(value+1,36);
        auto dib=std::find_if(clip.formats.begin(),clip.formats.end(),[](const auto &format){return format.id==CF_DIB;});
        if(dib==clip.formats.end() || dib->data.size()>128*1024*1024) return false;
        atomicWrite(directory/(id+L".capture"),protect(dib->data));
        stage="preview";
        auto previews=previewDirectory(); cleanCache(previews);
        auto preview=previews/(id+L".png"); savePreview(clip,preview);
        cleanCache(previews);
        cleanCache(directory);
        stage="XML";
        String documentClass(L"Windows.Data.Xml.Dom.XmlDocument"); ComPtr<IInspectable> documentInstance;
        check(RoActivateInstance(documentClass.value,&documentInstance));
        ComPtr<Xml::IXmlDocument> document; check(documentInstance.As(&document));
        ComPtr<Xml::IXmlDocumentIO> documentIO; check(document.As(&documentIO));
        auto uri=L"file:///"+preview.generic_wstring();
        String content(L"<toast launch=\""+id+L"\"><visual><binding template=\"ToastGeneric\">"
            L"<text>Screenshot copied</text><text>Click to edit and annotate</text><image id=\"1\" placement=\"appLogoOverride\" alt=\"Screenshot preview\" src=\""+escapeXml(uri)+L"\"/></binding></visual><audio silent=\"true\"/></toast>");
        check(documentIO->LoadXml(content.value));
        stage="toast factory";
        ComPtr<Notifications::IToastNotificationFactory> toastFactory; String toastClass(L"Windows.UI.Notifications.ToastNotification");
        check(RoGetActivationFactory(toastClass.value,IID_PPV_ARGS(&toastFactory)));
        ComPtr<Notifications::IToastNotification> toast; check(toastFactory->CreateToastNotification(document.Get(),&toast));
        ComPtr<ABI::Windows::Foundation::IPropertyValueStatics> values;
        String propertyClass(L"Windows.Foundation.PropertyValue");
        check(RoGetActivationFactory(propertyClass.value,IID_PPV_ARGS(&values)));
        FILETIME now{}; GetSystemTimeAsFileTime(&now);
        ULARGE_INTEGER timestamp{}; timestamp.LowPart=now.dwLowDateTime; timestamp.HighPart=now.dwHighDateTime;
        ABI::Windows::Foundation::DateTime expires{(INT64)(timestamp.QuadPart+24ull*60*60*10000000)};
        ComPtr<IInspectable> boxed; check(values->CreateDateTime(expires,&boxed));
        ComPtr<ABI::Windows::Foundation::IReference<ABI::Windows::Foundation::DateTime>> expiration;
        check(boxed.As(&expiration)); check(toast->put_ExpirationTime(expiration.Get()));
        stage="show";
        check(notifier->Show(toast.Get()));
        return true;
    } catch(const std::exception &e) { notificationError=std::string(stage)+": "+e.what(); log(wide(notificationError)); return false; } // Clipboard is already populated; fall back to the editor.
}
std::string screenshotNotificationError() { return notificationError; }

std::wstring updateToastXml(const std::wstring &version)
{
    return L"<toast launch=\"check-update\"><visual><binding template=\"ToastGeneric\">"
           L"<text>KSwitcher update available</text><text>Version " + escapeXml(version) +
           L" is available. Click to review and install.</text></binding></visual>"
           L"<actions><action content=\"Review update\" arguments=\"check-update\" activationType=\"foreground\"/>"
           L"<action content=\"Later\" arguments=\"dismiss\" activationType=\"system\"/></actions></toast>";
}
bool notifyUpdateAvailable(const std::wstring &version)
{
    try {
        // Reuse the screenshot notification identity, Windows settings and COM activator.
        registerScreenshotNotifications();
        ComPtr<Notifications::IToastNotificationManagerStatics> manager;
        String managerClass(L"Windows.UI.Notifications.ToastNotificationManager");
        check(RoGetActivationFactory(managerClass.value,IID_PPV_ARGS(&manager)));
        ComPtr<Notifications::IToastNotifier> notifier; String appId(AppId);
        check(manager->CreateToastNotifierWithId(appId.value,&notifier));
        Notifications::NotificationSetting setting{};
        auto settingsResult=notifier->get_Setting(&setting);
        if(settingsResult!=HRESULT_FROM_WIN32(ERROR_NOT_FOUND)) check(settingsResult);
        if(setting!=Notifications::NotificationSetting_Enabled)
            throw std::runtime_error("Windows notifications for KSwitcher Screenshot are disabled");
        String documentClass(L"Windows.Data.Xml.Dom.XmlDocument"); ComPtr<IInspectable> instance;
        check(RoActivateInstance(documentClass.value,&instance));
        ComPtr<Xml::IXmlDocument> document; check(instance.As(&document));
        ComPtr<Xml::IXmlDocumentIO> io; check(document.As(&io));
        String content(updateToastXml(version)); check(io->LoadXml(content.value));
        ComPtr<Notifications::IToastNotificationFactory> factory; String toastClass(L"Windows.UI.Notifications.ToastNotification");
        check(RoGetActivationFactory(toastClass.value,IID_PPV_ARGS(&factory)));
        ComPtr<Notifications::IToastNotification> toast; check(factory->CreateToastNotification(document.Get(),&toast));
        // Repeated submissions replace the existing update card in Notification Center.
        ComPtr<Notifications::IToastNotification2> tagged; check(toast.As(&tagged));
        String tag(L"update"),group(L"KSwitcher"); check(tagged->put_Tag(tag.value)); check(tagged->put_Group(group.value));
        check(notifier->Show(toast.Get()));
        return true;
    } catch(const std::exception &e) { log(L"Update toast: "+wide(e.what())); return false; }
}


}