// SPDX-License-Identifier: LGPL-2.1-or-later
#pragma once
#include <windows.h>
#include <cstddef>
#include <cstdint>
#include <memory>

namespace native_control {
inline bool query(HWND window, UINT message, WPARAM wp, LPARAM lp, DWORD_PTR *result) {
  return SendMessageTimeoutW(window, message, wp, lp, SMTO_ABORTIFHUNG,
                             100, result) != 0;
}

// Common-control messages above WM_USER do not marshal embedded pointers.
// Own the target allocation and process handle together, including failures.
class RemoteBuffer {
  struct Allocation {
    HANDLE process = nullptr;
    void *memory = nullptr;
    BOOL wow64 = FALSE;
    bool pending = false;
    LRESULT result = 0;
    ~Allocation() {
      if (memory) VirtualFreeEx(process, memory, 0, MEM_RELEASE);
      if (process) CloseHandle(process);
    }
  };
  inline static thread_local std::weak_ptr<Allocation> outstanding_;
  static void CALLBACK completed(HWND, UINT, ULONG_PTR data, LRESULT result) {
    std::unique_ptr<std::shared_ptr<Allocation>> retained(
        reinterpret_cast<std::shared_ptr<Allocation> *>(data));
    (*retained)->result = result;
    (*retained)->pending = false;
  }
public:
  RemoteBuffer(HWND window, size_t size) {
    // One late application request may retain its allocation. Do not create
    // another on every polling tick while that application is unresponsive.
    if (auto previous = outstanding_.lock(); previous && previous->pending) return;
    allocation_ = std::make_shared<Allocation>();
    DWORD processId = 0;
    GetWindowThreadProcessId(window, &processId);
    allocation_->process = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ |
                           PROCESS_VM_WRITE | PROCESS_QUERY_LIMITED_INFORMATION,
                           FALSE, processId);
    if (!allocation_->process || !IsWow64Process(allocation_->process, &allocation_->wow64)) return;
    allocation_->memory = VirtualAllocEx(allocation_->process, nullptr, size,
                             MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    // Every byte of the allocation must be addressable by the target.
    if (allocation_->memory && allocation_->wow64 &&
        (address() > UINT32_MAX || size - 1 > UINT32_MAX - address())) {
      VirtualFreeEx(allocation_->process, allocation_->memory, 0, MEM_RELEASE);
      allocation_->memory = nullptr;
    }
  }
  RemoteBuffer(const RemoteBuffer &) = delete;
  RemoteBuffer &operator=(const RemoteBuffer &) = delete;

  explicit operator bool() const { return allocation_ && allocation_->memory; }
  bool is32Bit() const { return allocation_->wow64 != FALSE; }
  uintptr_t address(size_t offset = 0) const {
    return reinterpret_cast<uintptr_t>(allocation_->memory) + offset;
  }
  bool write(const void *buffer, size_t size) const {
    return !allocation_->pending && WriteProcessMemory(allocation_->process, allocation_->memory, buffer, size, nullptr) != 0;
  }
  bool read(void *buffer, size_t size, size_t offset = 0) const {
    return !allocation_->pending && ReadProcessMemory(allocation_->process, reinterpret_cast<const void *>(address(offset)),
                              buffer, size, nullptr) != 0;
  }

  bool query(HWND window, UINT message, WPARAM wp, DWORD_PTR *result) {
    if (!*this || allocation_->pending) return false;
    allocation_->pending = true;
    outstanding_ = allocation_;
    auto retained = new std::shared_ptr<Allocation>(allocation_);
    if (!SendMessageCallbackW(window, message, wp, static_cast<LPARAM>(address()),
                             completed, reinterpret_cast<ULONG_PTR>(retained))) {
      allocation_->pending = false;
      delete retained;
      return false;
    }
    const ULONGLONG deadline = GetTickCount64() + 100;
    while (allocation_->pending && GetTickCount64() < deadline) {
      // Dispatch sent-message completions without consuming queued UI work.
      // Reentrant polls cannot allocate while outstanding_ remains pending.
      MSG message{};
      PeekMessageW(&message, nullptr, 0, 0, PM_NOREMOVE);
      if (allocation_->pending) Sleep(1);
    }
    if (allocation_->pending) return false;
    *result = static_cast<DWORD_PTR>(allocation_->result);
    return true;
  }

private:
  std::shared_ptr<Allocation> allocation_;
};
} // namespace native_control