// SPDX-License-Identifier: LGPL-2.1-or-later
#define COBJMACROS
#include "notification_hook.h"
#include "menu_bridge.h"
#include <oleauto.h>
#include <tlhelp32.h>
#include <uiautomationcore.h>
#include <uiautomationcoreapi.h>
#include <string.h>

// Wine currently drops UiaRaiseNotificationEvent. Replace direct imports in
// applications already reached by our native hook; leave other UIA APIs alone.
// Keep the DLL mapped for the application's lifetime once an import points here.
static HRESULT WINAPI raise_notification(IRawElementProviderSimple *provider,
    enum NotificationKind kind, enum NotificationProcessing processing,
    BSTR displayString, BSTR activityId) {
  (void)kind; (void)activityId;
  if (!provider) return E_INVALIDARG;
  HWND bridge = FindWindowW(WINE_ACCESS_MENU_CLASS, NULL);
  if (!bridge) return S_OK;
  WineAccessNotificationEvent event = {0};
  event.process = GetCurrentProcessId();
  event.processing = processing;
  UINT length = displayString ? SysStringLen(displayString) : 0;
  if (length >= ARRAYSIZE(event.text)) length = ARRAYSIZE(event.text) - 1;
  if (length && displayString[length - 1] >= 0xd800 &&
      displayString[length - 1] <= 0xdbff) --length;
  if (length) memcpy(event.text, displayString, length * sizeof(WCHAR));
  COPYDATASTRUCT data = {WINE_ACCESS_NOTIFICATION_TAG, sizeof(event), &event};
  DWORD_PTR result = 0;
  // The receiver only copies and queues. No reader calls on an app's thread.
  SendMessageTimeoutW(bridge, WM_COPYDATA, 0, (LPARAM)&data,
      SMTO_ABORTIFHUNG | SMTO_BLOCK, 50, &result);
  return S_OK;
}

static void watch_module(BYTE *base) {
  IMAGE_DOS_HEADER *dos = (IMAGE_DOS_HEADER *)base;
  if (dos->e_magic != IMAGE_DOS_SIGNATURE) return;
  IMAGE_NT_HEADERS *header = (IMAGE_NT_HEADERS *)(base + dos->e_lfanew);
  if (header->Signature != IMAGE_NT_SIGNATURE) return;
  IMAGE_DATA_DIRECTORY directory =
      header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
  if (!directory.VirtualAddress || !directory.Size) return;
  IMAGE_IMPORT_DESCRIPTOR *imports =
      (IMAGE_IMPORT_DESCRIPTOR *)(base + directory.VirtualAddress);
  for (; imports->Name; ++imports) {
    if (lstrcmpiA((char *)(base + imports->Name), "uiautomationcore.dll") ||
        !imports->OriginalFirstThunk) continue;
    IMAGE_THUNK_DATA *names = (IMAGE_THUNK_DATA *)(base + imports->OriginalFirstThunk);
    IMAGE_THUNK_DATA *addresses = (IMAGE_THUNK_DATA *)(base + imports->FirstThunk);
    for (; names->u1.AddressOfData; ++names, ++addresses) {
      if (IMAGE_SNAP_BY_ORDINAL(names->u1.Ordinal)) continue;
      IMAGE_IMPORT_BY_NAME *name = (IMAGE_IMPORT_BY_NAME *)(base + names->u1.AddressOfData);
      if (strcmp((char *)name->Name, "UiaRaiseNotificationEvent") ||
          addresses->u1.Function == (ULONG_PTR)raise_notification) continue;
      // The Wine stub even accepts a missing required provider. Leave an
      // implementation which rejects it (including native UIA) untouched.
      typedef HRESULT (WINAPI *RaiseNotification)(IRawElementProviderSimple *,
          enum NotificationKind, enum NotificationProcessing, BSTR, BSTR);
      RaiseNotification original = (RaiseNotification)(ULONG_PTR)addresses->u1.Function;
      if (original(NULL, NotificationKind_ActionCompleted,
                   NotificationProcessing_All, NULL, NULL) != S_OK) continue;
      HMODULE pinnedModule = NULL;
      if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
          GET_MODULE_HANDLE_EX_FLAG_PIN, (LPCWSTR)(ULONG_PTR)raise_notification,
          &pinnedModule)) continue;
      DWORD protection;
      if (!VirtualProtect(&addresses->u1.Function, sizeof(addresses->u1.Function),
                          PAGE_READWRITE, &protection)) continue;
      InterlockedExchangePointer((PVOID volatile *)&addresses->u1.Function,
                                 (PVOID)(ULONG_PTR)raise_notification);
      DWORD ignored;
      VirtualProtect(&addresses->u1.Function, sizeof(addresses->u1.Function),
                     protection, &ignored);
    }
  }
}

void wine_access_watch_notifications(void) {
  static DWORD lastScan;
  static LONG scanning;
  DWORD now = GetTickCount();
  if (InterlockedCompareExchange(&scanning, 1, 0)) return;
  if (lastScan && now - lastScan < 1000) {
    InterlockedExchange(&scanning, 0);
    return;
  }
  lastScan = now;
  HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetCurrentProcessId());
  if (snapshot != INVALID_HANDLE_VALUE) {
    MODULEENTRY32W module = {0};
    module.dwSize = sizeof(module);
    if (Module32FirstW(snapshot, &module)) {
      do {
        // Hold a reference while walking imports: another thread may unload it.
        HMODULE reference = NULL;
        if (GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
                              (LPCWSTR)module.modBaseAddr, &reference)) {
          watch_module((BYTE *)reference);
          FreeLibrary(reference);
        }
      } while (Module32NextW(snapshot, &module));
    }
    CloseHandle(snapshot);
  }
  InterlockedExchange(&scanning, 0);
}