// SPDX-License-Identifier: LGPL-2.1-or-later
#define COBJMACROS
#include "bootstrapper_hook.h"
#include "menu_bridge.h"
#include <initguid.h>
#include <uiautomationcore.h>
#include <uiautomationcoreapi.h>
#include <uiautomationclient.h>
#include <string.h>

typedef LRESULT (WINAPI *ReturnProvider)(HWND, WPARAM, LPARAM, IRawElementProviderSimple *);
static ReturnProvider originalReturn;
static HWND queryWindow, queryBridge;
static LONG querying;

static void read_toggle(IRawElementProviderSimple *root) {
  IRawElementProviderFragmentRoot *fragmentRoot = NULL;
  IRawElementProviderFragment *focus = NULL;
  IRawElementProviderSimple *provider = NULL;
  IUnknown *pattern = NULL;
  IToggleProvider *toggle = NULL;
  VARIANT type = {0}, name = {0}, enabled = {0}, password = {0};
  enum ToggleState state;
  if (FAILED(IRawElementProviderSimple_QueryInterface(root, &IID_IRawElementProviderFragmentRoot,
      (void **)&fragmentRoot)) || !fragmentRoot) goto done;
  if (FAILED(IRawElementProviderFragmentRoot_GetFocus(fragmentRoot, &focus)) || !focus) goto done;
  if (FAILED(IRawElementProviderFragment_QueryInterface(focus, &IID_IRawElementProviderSimple,
      (void **)&provider)) || !provider) goto done;
  if (FAILED(IRawElementProviderSimple_GetPropertyValue(provider, UIA_ControlTypePropertyId, &type)) ||
      type.vt != VT_I4 || type.lVal != UIA_CheckBoxControlTypeId) goto done;
  if (FAILED(IRawElementProviderSimple_GetPatternProvider(provider, UIA_TogglePatternId, &pattern)) ||
      !pattern || FAILED(IUnknown_QueryInterface(pattern, &IID_IToggleProvider, (void **)&toggle)) ||
      !toggle || FAILED(IToggleProvider_get_ToggleState(toggle, &state)) ||
      (state != ToggleState_Off && state != ToggleState_On && state != ToggleState_Indeterminate)) goto done;
  if (FAILED(IRawElementProviderSimple_GetPropertyValue(provider, UIA_NamePropertyId, &name)) ||
      name.vt != VT_BSTR || !name.bstrVal || GetForegroundWindow() != queryWindow) goto done;
  WineAccessBootstrapperState event = {0};
  event.owner = (ULONGLONG)(UINT_PTR)queryWindow;
  event.state = state;
  // Copy all checkbox state on this UI thread. Mixing a newer Toggle value
  // with flags from the asynchronous UIA worker can announce a false change.
  if (SUCCEEDED(IRawElementProviderSimple_GetPropertyValue(provider, UIA_IsEnabledPropertyId, &enabled)) &&
      enabled.vt == VT_BOOL && enabled.boolVal == VARIANT_FALSE)
    event.flags |= WINE_ACCESS_BOOTSTRAPPER_DISABLED;
  if (SUCCEEDED(IRawElementProviderSimple_GetPropertyValue(provider, UIA_IsPasswordPropertyId, &password)) &&
      password.vt == VT_BOOL && password.boolVal == VARIANT_TRUE)
    event.flags |= WINE_ACCESS_BOOTSTRAPPER_PASSWORD;
  event.tick = GetTickCount64();
  UINT length = SysStringLen(name.bstrVal);
  if (length >= ARRAYSIZE(event.name)) goto done;
  memcpy(event.name, name.bstrVal, length * sizeof(WCHAR));
  COPYDATASTRUCT data = {WINE_ACCESS_BOOTSTRAPPER_TAG, sizeof(event), &event};
  DWORD_PTR result;
  SendMessageTimeoutW(queryBridge, WM_COPYDATA, 0, (LPARAM)&data,
      SMTO_ABORTIFHUNG | SMTO_BLOCK, 50, &result);
done:
  VariantClear(&type); VariantClear(&name);
  VariantClear(&enabled); VariantClear(&password);
  if (toggle) IToggleProvider_Release(toggle);
  if (pattern) IUnknown_Release(pattern);
  if (provider) IRawElementProviderSimple_Release(provider);
  if (focus) IRawElementProviderFragment_Release(focus);
  if (fragmentRoot) IRawElementProviderFragmentRoot_Release(fragmentRoot);
}

static LRESULT WINAPI return_provider(HWND window, WPARAM wp, LPARAM lp,
                                     IRawElementProviderSimple *provider) {
  // Borrow the provider only during this original call, on its UI thread.
  // Never retain it, toggle a control, or change the original UIA result.
  if (provider && window == queryWindow && queryBridge && (LONG)lp == UiaRootObjectId)
    read_toggle(provider);
  return originalReturn(window, wp, lp, provider);
}

static BOOL watch_import(void) {
  BYTE *base = (BYTE *)GetModuleHandleW(NULL);
  IMAGE_DOS_HEADER *dos = (IMAGE_DOS_HEADER *)base;
  if (!base || dos->e_magic != IMAGE_DOS_SIGNATURE) return FALSE;
  IMAGE_NT_HEADERS *header = (IMAGE_NT_HEADERS *)(base + dos->e_lfanew);
  if (header->Signature != IMAGE_NT_SIGNATURE) return FALSE;
  IMAGE_DATA_DIRECTORY directory = header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
  if (!directory.VirtualAddress || !directory.Size) return FALSE;
  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, "UiaReturnRawElementProvider")) continue;
      if (addresses->u1.Function == (ULONG_PTR)return_provider) return TRUE;
      HMODULE uia = GetModuleHandleW(L"uiautomationcore.dll"), pinned = NULL;
      if (!uia) return FALSE;
      FARPROC address = GetProcAddress(uia, "UiaReturnRawElementProvider");
      if (!address || addresses->u1.Function != (ULONG_PTR)address) return FALSE;
      memcpy(&originalReturn, &address, sizeof(originalReturn));
      if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_PIN,
          (LPCWSTR)(ULONG_PTR)return_provider, &pinned)) return FALSE;
      DWORD protection;
      if (!VirtualProtect(&addresses->u1.Function, sizeof(addresses->u1.Function), PAGE_READWRITE,
          &protection)) return FALSE;
      InterlockedExchangePointer((PVOID volatile *)&addresses->u1.Function, (PVOID)(ULONG_PTR)return_provider);
      DWORD ignored;
      VirtualProtect(&addresses->u1.Function, sizeof(addresses->u1.Function), protection, &ignored);
      return TRUE;
    }
  }
  return FALSE;
}

void wine_access_query_bootstrapper_state(HWND window, HWND bridge) {
  WCHAR className[128];
  DWORD process = 0;
  if (!bridge || window != GetForegroundWindow() ||
      GetWindowThreadProcessId(window, &process) != GetCurrentThreadId() ||
      process != GetCurrentProcessId() || !GetClassNameW(window, className, ARRAYSIZE(className)) ||
      lstrcmpW(className, L"Blizzard Bootstrapper")) return;
  if (InterlockedCompareExchange(&querying, 1, 0)) return;
  if (watch_import()) {
    queryWindow = window; queryBridge = bridge;
    SendMessageW(window, WM_GETOBJECT, 0, UiaRootObjectId);
    queryWindow = queryBridge = NULL;
  }
  InterlockedExchange(&querying, 0);
}