#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;

class PEModifier {
private:
  std::vector<BYTE> fileData;
  IMAGE_DOS_HEADER *dosHeader;
  IMAGE_NT_HEADERS *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_HEADERS *)(fileData.data() + dosHeader->e_lfanew);
    if (ntHeaders->Signature != IMAGE_NT_SIGNATURE)
      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(".rsrc");
    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;

    // Find end of rsrc section (aligned)
    DWORD sectionAlignment = ntHeaders->OptionalHeader.SectionAlignment;
    DWORD fileAlignment = ntHeaders->OptionalHeader.FileAlignment;

    // Calculate where to place debug data in the rsrc section
    DWORD currentSectionEnd =
        rsrcSection->PointerToRawData + rsrcSection->SizeOfRawData;
    DWORD availableSpace =
        rsrcSection->Misc.VirtualSize - rsrcSection->SizeOfRawData;

    if (debug_data_size > availableSpace) {
      printf("Error: Not enough space in .rsrc section! Need %zu bytes, have "
             "%lu\n",
             debug_data_size, availableSpace);
      return false;
    }

    // Calculate RVA for debug data
    DWORD debugDataRVA =
        rsrcSection->VirtualAddress + rsrcSection->SizeOfRawData;
    DWORD debugDataFileOffset = currentSectionEnd;

    // 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);

    // Extend file data if needed
    if (debugDataFileOffset + debug_data_size > fileData.size()) {
      fileData.resize(debugDataFileOffset + debug_data_size);
    }

    // Write debug data to rsrc section
    memcpy(fileData.data() + debugDataFileOffset, debugData.data(),
           debug_data_size);

    // Update rsrc section size
    rsrcSection->SizeOfRawData += debug_data_size;

    // Create/update 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;

    // Update debug directory in optional header
    ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG]
        .VirtualAddress = debugDataRVA;
    ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].Size =
        sizeof(IMAGE_DEBUG_DIRECTORY);

    // Actually, we need to place the debug directory entry somewhere too
    // Let's put it right before the debug data
    DWORD debugDirRVA = debugDataRVA - sizeof(IMAGE_DEBUG_DIRECTORY);
    DWORD debugDirFileOffset =
        debugDataFileOffset - sizeof(IMAGE_DEBUG_DIRECTORY);

    // Adjust our calculations
    rsrcSection->SizeOfRawData += sizeof(IMAGE_DEBUG_DIRECTORY);

    // Write debug directory entry
    memcpy(fileData.data() + debugDirFileOffset, &debugDir,
           sizeof(IMAGE_DEBUG_DIRECTORY));

    // Update optional header to point to 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\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) + ".debug";
  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;
}