#include "r3/config/static.hpp"
#include <spdlog/spdlog.h>
#include <windows.h>
#include <DbgHelp.h>
#include <memoryapi.h>

#define GH_BASE_ADDR 0x00400000

static uintptr_t g_gh_translationOffset{};

struct R3Bin {
  R3Bin() { loadOriginal(); }
 
  void loadOriginal() {
    SPDLOG_DEBUG("Loading original binary");

    auto &config = getDefaultConfig();
    std::string path = config.gameRootDir + "/Rayman3.exe";

    module = LoadLibraryA(path.c_str());
    if (!module)
      throw std::runtime_error("Failed to load original binary");

    fixupImports(module);

    g_gh_translationOffset = translationOffset =
        uintptr_t(module) - GH_BASE_ADDR;

    // Now we have to relocate the module to the new base address
  }

  inline void relocate(HMODULE module, void* from, void* to, void* check) {
    void* relocated_addr = (void*)(uintptr_t(from) + translationOffset);
    void* relocated_to = (void*)(uintptr_t(to) + translationOffset);
    void *checkRead{};
    SIZE_T numRead{};
    ReadProcessMemory(GetCurrentProcess(), relocated_addr, &checkRead, sizeof(checkRead), &numRead);
    WriteProcessMemory(GetCurrentProcess(), relocated_addr, relocated_to, sizeof(relocated_to), NULL);
  }
  
  void relocateModule(HMODULE module) {
    #define REL(from, to, original) relocate(module, (void*)(from), (void*)(to), (void*)(original));
    #include "relocations.def"
  }

  // Translate address relative from the original image base address
  void *translateAddress(void *original) {
    uint8_t *runtime_addr =
        reinterpret_cast<uint8_t *>(original) + translationOffset;
    return reinterpret_cast<void *>(runtime_addr);
  }

  void fixupImports(HINSTANCE h) {
    // Find the IAT size
    DWORD ulsize = 0;
    PIMAGE_IMPORT_DESCRIPTOR pImportDesc =
        (PIMAGE_IMPORT_DESCRIPTOR)ImageDirectoryEntryToData(
            h, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ulsize);
    if (!pImportDesc)
      return;

    // Loop names
    for (; pImportDesc->Name; pImportDesc++) {
      PSTR pszModName = (PSTR)((PBYTE)h + pImportDesc->Name);
      if (!pszModName)
        break;

      HINSTANCE hImportDLL = LoadLibraryA(pszModName);
      if (!hImportDLL) {
        // ... (error)
      }

      // Get caller's import address table (IAT) for the callee's functions
      PIMAGE_THUNK_DATA pThunk =
          (PIMAGE_THUNK_DATA)((PBYTE)h + pImportDesc->FirstThunk);

      // Replace current function address with new function address
      for (; pThunk->u1.Function; pThunk++) {
        FARPROC pfnNew = 0;
        size_t rva = 0;
#ifdef _WIN64
        if (pThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG64)
#else
        if (pThunk->u1.Ordinal & IMAGE_ORDINAL_FLAG32)
#endif
        {
          // Ordinal
#ifdef _WIN64
          size_t ord = IMAGE_ORDINAL64(pThunk->u1.Ordinal);
#else
          size_t ord = IMAGE_ORDINAL32(pThunk->u1.Ordinal);
#endif

          PROC *ppfn = (PROC *)&pThunk->u1.Function;
          if (!ppfn) {
            // ... (error)
          }
          rva = (size_t)pThunk;

          char fe[100] = {0};
          sprintf_s(fe, 100, "#%u", ord);
          pfnNew = GetProcAddress(hImportDLL, (LPCSTR)ord);
          if (!pfnNew) {
            // ... (error)
          }
        } else {
          // Get the address of the function address
          PROC *ppfn = (PROC *)&pThunk->u1.Function;
          if (!ppfn) {
            // ... (error)
          }
          rva = (size_t)pThunk;
          PSTR fName = (PSTR)h;
          fName += pThunk->u1.Function;
          fName += 2;
          if (!fName)
            break;
          pfnNew = GetProcAddress(hImportDLL, fName);
          if (!pfnNew) {
            // ... (error)
          }
        }

        // Patch it now...
        auto hp = GetCurrentProcess();
        if (!WriteProcessMemory(hp, (LPVOID *)rva, &pfnNew, sizeof(pfnNew),
                                NULL) &&
            (ERROR_NOACCESS == GetLastError())) {
          DWORD dwOldProtect;
          if (VirtualProtect((LPVOID)rva, sizeof(pfnNew), PAGE_WRITECOPY,
                             &dwOldProtect)) {
            if (!WriteProcessMemory(GetCurrentProcess(), (LPVOID *)rva, &pfnNew,
                                    sizeof(pfnNew), NULL)) {
              // ... (error)
            }
            if (!VirtualProtect((LPVOID)rva, sizeof(pfnNew), dwOldProtect,
                                &dwOldProtect)) {
              // ... (error)
            }
          }
        }
      }
    }
  }

  HINSTANCE module;
  uintptr_t translationOffset;

  static R3Bin &get() {
    static R3Bin instance;
    return instance;
  }
};

uint8_t *gh_map_dbg_mem(size_t addr) {
  R3Bin::get();
  SPDLOG_DEBUG("Mapping debug memory at {}", addr);
  return (uint8_t *)R3Bin::get().translateAddress((void *)addr);
}
void *gh_stub_impl_ptr(void *ptr) {
  R3Bin::get();
  SPDLOG_DEBUG("Forwarding implementation at {}", ptr);
  return (void *)R3Bin::get().translateAddress((void *)ptr);
}
void gh_init_dbg_loader() { R3Bin::get(); }