#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <inspectable.h>
#include <winstring.h>
#include <stdint.h>
#include <stdio.h>
#include <wchar.h>
#include "core_discovery.h"
struct Packet { volatile LONG state; LONG hr; DWORD thread; DWORD protocol;
  wchar_t script[8192]; wchar_t result[32768]; };
static const UINT Wake = WM_APP + 0x498;
static Packet* packet=nullptr;
static HANDLE mapping=nullptr;
static IUnknown* core=nullptr;
static IUnknown* operation=nullptr;
static IUnknown* info=nullptr;
static const GUID CoreIID={0x3a3f559a,0xe5e9,0x5338,{0xbb,0x67,0x4e,0xb0,0x50,0x4a,0x4f,0x14}};
static const GUID InfoIID={0x00000036,0,0,{0xc0,0,0,0,0,0,0,0x46}};
static void** table(void* p) { return *(void***)p; }
static HRESULT finishAsync() {
  HRESULT outcome=S_OK;
  if(info) {
    using Status=HRESULT(STDMETHODCALLTYPE*)(void*,INT32*);
    using Action=HRESULT(STDMETHODCALLTYPE*)(void*);
    INT32 status=0;
    HRESULT hr=((Status)table(info)[7])(info,&status);
    if(SUCCEEDED(hr) && status==0) {
      hr=((Action)table(info)[9])(info); // IAsyncInfo.Cancel
      if(FAILED(hr))outcome=hr;
    }
    hr=((Action)table(info)[10])(info); // IAsyncInfo.Close
    if(FAILED(hr) && SUCCEEDED(outcome))outcome=hr;
  }
  return outcome;
}
static void clearAsync() {
  if(info){info->Release();info=nullptr;}
  if(operation){operation->Release();operation=nullptr;}
}
static void fail(HRESULT hr) { clearAsync();packet->hr=hr;InterlockedExchange(&packet->state,-1); }
static HRESULT findCore() {
  auto module=(uintptr_t)GetModuleHandleW(L"Microsoft.Web.WebView2.Core.dll");
  if(!module)return HRESULT_FROM_WIN32(ERROR_MOD_NOT_FOUND);
  CoreLayout layout{};
  if(!discoverCoreLayout(module,layout))return HRESULT_FROM_WIN32(ERROR_REVISION_MISMATCH);
  MEMORY_BASIC_INFORMATION mbi;
  uintptr_t address=0,scanned=0;
  while(VirtualQuery((void*)address,&mbi,sizeof(mbi))) {
    uintptr_t end=(uintptr_t)mbi.BaseAddress+mbi.RegionSize;
    if(end<=address)break;
    address=end;
    if(mbi.State!=MEM_COMMIT||mbi.Type!=MEM_PRIVATE||mbi.Protect!=PAGE_READWRITE||
       mbi.RegionSize>64*1024*1024)continue;
    if(scanned+mbi.RegionSize>256*1024*1024)break;
    scanned+=mbi.RegionSize;
    for(uintptr_t p=(uintptr_t)mbi.BaseAddress;p+0x118<=end;p+=8) {
      if(*(uintptr_t*)p!=layout.primary || *(uintptr_t*)(p+0x108)!=layout.secondary)continue;
      if(table((void*)p)[0]!=(void*)(layout.query))continue;
      auto hr=((IUnknown*)p)->QueryInterface(CoreIID,(void**)&core);
      if(FAILED(hr))continue;
      if(!supportedExecute(table(core)[49])) { core->Release();core=nullptr;continue; }
      return S_OK;
    }
  }
  return HRESULT_FROM_WIN32(ERROR_NOT_FOUND);
}
static void runProbe() {
  if(!packet) {
    wchar_t name[96];swprintf_s(name,L"Local\\WhatsAppMathProbe-v2-%lu",GetCurrentProcessId());
    mapping=OpenFileMappingW(FILE_MAP_ALL_ACCESS,FALSE,name);
    if(!mapping)return;
    packet=(Packet*)MapViewOfFile(mapping,FILE_MAP_ALL_ACCESS,0,0,sizeof(Packet));
    if(!packet){CloseHandle(mapping);mapping=nullptr;return;}
  }
  if(packet->protocol!=0x4d415432){fail(E_INVALIDARG);return;}
  if(packet->thread!=GetCurrentThreadId()) { fail(RPC_E_WRONG_THREAD);return; }
  LONG state=packet->state;
  if(state==3) {InterlockedExchange(&packet->state,4);return;}
  if(state==6) {
    packet->hr=finishAsync();clearAsync();if(core){core->Release();core=nullptr;}
    InterlockedExchange(&packet->state,7);
    UnmapViewOfFile(packet);packet=nullptr;CloseHandle(mapping);mapping=nullptr;
    return;
  }
  if(state==0) {
    if(!core){auto hr=findCore();if(FAILED(hr)){fail(hr);return;}}
    HSTRING text=nullptr;
    auto hr=WindowsCreateString(packet->script,(UINT32)wcsnlen_s(packet->script,8192),&text);
    if(FAILED(hr)){fail(hr);return;}
    using Execute=HRESULT(STDMETHODCALLTYPE*)(void*,HSTRING,IUnknown**);
    hr=((Execute)table(core)[49])(core,text,&operation);
    WindowsDeleteString(text);
    if(FAILED(hr)){fail(hr);return;}
    hr=operation->QueryInterface(InfoIID,(void**)&info);
    if(FAILED(hr)){fail(hr);return;}
    InterlockedExchange(&packet->state,1);
  }
  if(packet->state==1) {
    using Status=HRESULT(STDMETHODCALLTYPE*)(void*,INT32*);
    INT32 status=0;auto hr=((Status)table(info)[7])(info,&status);
    if(FAILED(hr)){fail(hr);return;}
    if(status==0)return;
    if(status!=1){fail(E_ABORT);return;}
    using Results=HRESULT(STDMETHODCALLTYPE*)(void*,HSTRING*);
    HSTRING value=nullptr;hr=((Results)table(operation)[8])(operation,&value);
    if(FAILED(hr)){fail(hr);return;}
    UINT32 length=0;auto raw=WindowsGetStringRawBuffer(value,&length);
    if(length>=32768){WindowsDeleteString(value);fail(HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER));return;}
    wmemcpy(packet->result,raw,length);packet->result[length]=0;
    WindowsDeleteString(value);finishAsync();clearAsync();packet->hr=S_OK;
    InterlockedExchange(&packet->state,2);
  }
}
extern "C" __declspec(dllexport) LRESULT CALLBACK ProbeHook(int code,WPARAM wParam,LPARAM lParam) {
  if(code>=0 && wParam==PM_REMOVE && ((MSG*)lParam)->message==Wake) {
    ((MSG*)lParam)->message=WM_NULL;
    __try {runProbe();}
    __except(EXCEPTION_EXECUTE_HANDLER) {
      if(packet){packet->hr=(LONG)GetExceptionCode();InterlockedExchange(&packet->state,-1);}
    }
  }
  return CallNextHookEx(nullptr,code,wParam,lParam);
}
BOOL WINAPI DllMain(HINSTANCE,DWORD,LPVOID) {return TRUE;}