#include <Windows.h>
#include <psapi.h>

#include <filesystem>
#include <fstream>
#include <iostream>

#include <spdlog/spdlog.h>
#include <theo.hpp>

#include <obf/engine.hpp>
#include <obf/passes/jcc_rewrite_pass.hpp>
#include <obf/passes/next_inst_pass.hpp>
#include <obf/passes/reloc_transform_pass.hpp>

namespace fs = std::filesystem;

int main(int argc, char* argv[]) {
  if (argc < 2)
    return -1;

  // read in lib file...
  std::ifstream f(argv[1], std::ios::binary);
  auto fsize = fs::file_size(fs::path(argv[1]));
  std::vector<std::uint8_t> fdata;
  fdata.resize(fsize);
  f.read((char*)fdata.data(), fsize);

  LoadLibraryA("user32.dll");
  LoadLibraryA("win32u.dll");

  theo::recomp::allocator_t allocator =
      [&](std::uint32_t size,
          coff::section_characteristics_t section_type) -> std::uintptr_t {
    return reinterpret_cast<std::uintptr_t>(VirtualAlloc(
        NULL, size, MEM_COMMIT | MEM_RESERVE,
        section_type.mem_execute ? PAGE_EXECUTE_READWRITE : PAGE_READWRITE));
  };

  theo::recomp::copier_t copier = [&](std::uintptr_t ptr, void* buff,
                                      std::uint32_t size) {
    std::memcpy((void*)ptr, buff, size);
  };

  theo::recomp::resolver_t resolver = [&](std::string sym) -> std::uintptr_t {
    auto loaded_modules = std::make_unique<HMODULE[]>(64);
    std::uintptr_t result = 0u, loaded_module_sz = 0u;
    if (!EnumProcessModules(GetCurrentProcess(), loaded_modules.get(), 512,
                            (PDWORD)&loaded_module_sz))
      return {};

    for (auto i = 0u; i < loaded_module_sz / 8u; i++) {
      wchar_t file_name[MAX_PATH] = L"";
      if (!GetModuleFileNameExW(GetCurrentProcess(), loaded_modules.get()[i],
                                file_name, _countof(file_name)))
        continue;

      if ((result = reinterpret_cast<std::uintptr_t>(
               GetProcAddress(LoadLibraryW(file_name), sym.c_str()))))
        break;
    }
    return result;
  };

  // init enc/dec tables only once... important that this is done before adding
  // obfuscation passes to the engine...
  //
  xed_tables_init();

  // order matters, the order in which the pass is added is the order they
  // will be executed!
  //
  auto engine = theo::obf::engine_t::get();
  engine->add_pass(theo::obf::reloc_transform_pass_t::get());
  engine->add_pass(theo::obf::next_inst_pass_t::get());
  engine->add_pass(theo::obf::jcc_rewrite_pass_t::get());

  std::string entry_name;
  std::cout << "enter the name of the entry point: ";
  std::cin >> entry_name;

  theo::theo_t t(fdata, {allocator, copier, resolver}, entry_name.data());
  auto res = t.decompose();

  if (!res.has_value()) {
    spdlog::error("decomposition failed...\n");
    return -1;
  }

  spdlog::info("decomposed {} symbols...", res.value());
  auto entry_pnt = t.compose();
  spdlog::info("entry point address: {:X}, press enter to execute it...",
               entry_pnt);

  reinterpret_cast<void (*)()>(entry_pnt)();
}