#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <DbgHelp.h>
#include <string>
#include <vector>
#include <fstream>
#include <filesystem>

// RSDS debug format
typedef struct {
  DWORD signature; // 'RSDS'
  BYTE guid[16];   // GUID
  DWORD age;       // Age
                   // Followed by null-terminated PDB path
} RSDS_DEBUG_FORMAT;

typedef struct _IMAGE_OPTIONAL_HEADER6433 {
    WORD        Magic;
    BYTE        MajorLinkerVersion;
    BYTE        MinorLinkerVersion;
    DWORD       SizeOfCode;
    DWORD       SizeOfInitializedData;
    DWORD       SizeOfUninitializedData;
    DWORD       AddressOfEntryPoint;
    DWORD       BaseOfCode;
    ULONGLONG   ImageBase;
    DWORD       SectionAlignment;
    DWORD       FileAlignment;
    WORD        MajorOperatingSystemVersion;
    WORD        MinorOperatingSystemVersion;
    WORD        MajorImageVersion;
    WORD        MinorImageVersion;
    WORD        MajorSubsystemVersion;
    WORD        MinorSubsystemVersion;
    DWORD       Win32VersionValue;
    DWORD       SizeOfImage;
    DWORD       SizeOfHeaders;
    DWORD       CheckSum;
    WORD        Subsystem;
    WORD        DllCharacteristics;
    ULONGLONG   SizeOfStackReserve;
    ULONGLONG   SizeOfStackCommit;
    ULONGLONG   SizeOfHeapReserve;
    ULONGLONG   SizeOfHeapCommit;
    DWORD       LoaderFlags;
    DWORD       NumberOfRvaAndSizes;
    IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES];
};

class PEModifier {
private:
  std::vector<BYTE> fileData;
  IMAGE_DOS_HEADER *dosHeader;
  IMAGE_NT_HEADERS32 *ntHeaders;
  IMAGE_SECTION_HEADER *sectionHeaders;

public:
  bool loadPE(const char *filename) {
    std::ifstream file(filename, std::ios::binary);
    if (!file)
      return false;

    file.seekg(0, std::ios::end);
    size_t fileSize = file.tellg();
    file.seekg(0, std::ios::beg);

    fileData.resize(fileSize);
    file.read((char *)fileData.data(), fileSize);
    file.close();

    // Parse PE headers
    dosHeader = (IMAGE_DOS_HEADER *)fileData.data();
    if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
      return false;

    ntHeaders = (IMAGE_NT_HEADERS32 *)(fileData.data() + dosHeader->e_lfanew);
    if (ntHeaders->Signature != IMAGE_NT_SIGNATURE)
      return false;

    // Verify it's actually a 32-bit PE
    if (ntHeaders->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
      printf("Error: This is not a 32-bit PE file!\n");
      return false;
    }

    sectionHeaders =
        (IMAGE_SECTION_HEADER *)((BYTE *)&ntHeaders->OptionalHeader +
                                 ntHeaders->FileHeader.SizeOfOptionalHeader);

    return true;
  }

  IMAGE_SECTION_HEADER *findSection(const char *name) {
    for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
      if (memcmp(sectionHeaders[i].Name, name, strlen(name)) == 0) {
        return &sectionHeaders[i];
      }
    }
    return nullptr;
  }

  bool addDebugDirectory(const std::string &pdbPath) {
    IMAGE_SECTION_HEADER *rsrcSection = findSection(".rdata");
    if (!rsrcSection) {
      printf("Error: .rsrc section not found!\n");
      return false;
    }

    // Calculate debug data size
    size_t pdb_path_len = pdbPath.size() + 1;
    size_t debug_data_size = sizeof(RSDS_DEBUG_FORMAT) + pdb_path_len;
    size_t total_needed = sizeof(IMAGE_DEBUG_DIRECTORY) + debug_data_size + 32; // + padding

    // Check if we have enough space at the end of rsrc section
    if (total_needed > rsrcSection->SizeOfRawData) {
      printf("Error: Not enough space in .rsrc section! Need %zu bytes, section is %lu bytes\n",
             total_needed, rsrcSection->SizeOfRawData);
      return false;
    }

    // Calculate where to place debug directory and data (at the very end of rsrc section)
    DWORD debugDirFileOffset = rsrcSection->PointerToRawData + rsrcSection->SizeOfRawData - total_needed;
    DWORD debugDataFileOffset = debugDirFileOffset + sizeof(IMAGE_DEBUG_DIRECTORY);

    DWORD debugDirRVA = rsrcSection->VirtualAddress + rsrcSection->SizeOfRawData - total_needed;
    DWORD debugDataRVA = debugDirRVA + sizeof(IMAGE_DEBUG_DIRECTORY);

    // Prepare debug data
    std::vector<BYTE> debugData(debug_data_size);

    RSDS_DEBUG_FORMAT rsds = {0};
    rsds.signature = 0x53445352; // 'RSDS'

    // Generate or use existing GUID
    BYTE guid_bytes[] = {0x7E, 0x71, 0xA5, 0x7F, 0x3A, 0x25, 0xB5, 0xB9,
                         0x4C, 0x4C, 0x44, 0x20, 0x50, 0x44, 0x42, 0x2E};
    memcpy(rsds.guid, guid_bytes, 16);
    rsds.age = 0x00000001;

    // Copy RSDS header
    memcpy(debugData.data(), &rsds, sizeof(RSDS_DEBUG_FORMAT));

    // Copy PDB path
    memcpy(debugData.data() + sizeof(RSDS_DEBUG_FORMAT), pdbPath.c_str(),
           pdb_path_len);

    // Create debug directory entry
    IMAGE_DEBUG_DIRECTORY debugDir = {0};
    debugDir.Characteristics = 0x00000000;
    debugDir.TimeDateStamp = ntHeaders->FileHeader.TimeDateStamp;
    debugDir.MajorVersion = 0x0000;
    debugDir.MinorVersion = 0x0002;
    debugDir.Type = IMAGE_DEBUG_TYPE_CODEVIEW;
    debugDir.SizeOfData = debug_data_size;
    debugDir.AddressOfRawData = debugDataRVA;
    debugDir.PointerToRawData = debugDataFileOffset;

    // Overwrite the end of rsrc section with debug directory entry
    memcpy(fileData.data() + debugDirFileOffset, &debugDir,
           sizeof(IMAGE_DEBUG_DIRECTORY));

    // Overwrite with debug data
    memcpy(fileData.data() + debugDataFileOffset, debugData.data(),
           debug_data_size);

    // Update optional header to point to debug directory
    printf("IMAGE_DIRECTORY_ENTRY_DEBUG constant value: %d\n", IMAGE_DIRECTORY_ENTRY_DEBUG);
    printf("Setting debug directory at index %d\n", IMAGE_DIRECTORY_ENTRY_DEBUG);
    printf("Debug directory RVA: 0x%08X\n", debugDirRVA);
    printf("Debug directory size: %d\n", sizeof(IMAGE_DEBUG_DIRECTORY));

    ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG]
        .VirtualAddress = debugDirRVA;
    ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].Size =
        sizeof(IMAGE_DEBUG_DIRECTORY);

    printf("Debug directory added at RVA: 0x%08X (overwriting rsrc end)\n", debugDirRVA);
    printf("Debug data added at RVA: 0x%08X\n", debugDataRVA);

    return true;
  }

  bool savePE(const char *filename) {
    std::string newFilename(filename);
    // Replace .exe with .debug.exe
    newFilename.replace(newFilename.find_last_of(".") + 1, 3, "debug.exe");
    std::ofstream file(newFilename.c_str(), std::ios::binary);
    if (!file)
      return false;

    file.write((char *)fileData.data(), fileData.size());
    file.close();
    return true;
  }
};

int main(int argc, char **argv) {
  if (argc != 2) {
    printf("Usage: %s <path_to_exe>\n", argv[0]);
    return 1;
  }

  const char *exePath = argv[1];
  // Pick the filename only, and replace the extension with .pdb
  std::filesystem::path pdbPath(exePath);
  pdbPath = pdbPath.stem().string();
  pdbPath += ".pdb";

  PEModifier pe;

  if (!pe.loadPE(exePath)) {
    printf("Error: Failed to load PE file: %s\n", exePath);
    return 1;
  }

  if (!pe.addDebugDirectory(pdbPath.string())) {
    printf("Error: Failed to add debug directory\n");
    return 1;
  }

  // Save modified PE
  std::string outputPath = std::string(exePath);
  if (!pe.savePE(outputPath.c_str())) {
    printf("Error: Failed to save modified PE file\n");
    return 1;
  }

  printf("Successfully added debug directory to: %s\n", outputPath.c_str());
  return 0;
}