#include <windows.h>
#include <coffi/coffi.hpp>
#include <CLI11.hpp>
#include <spdlog/spdlog.h>

int main(int argc, char *argv[]) {
  CLI::App app("Patcher");
  std::string inputFile;
  std::string outputFile;
  app.add_option("-i,--input", inputFile, "Input exe file to patch")
      ->required();
  app.add_option("-o,--output", outputFile, "Output patched exe file")
      ->required();

  CLI11_PARSE(app, argc, argv);

  // Load the object file containing the main function
  COFFI::coffi objReader;
  std::string objPath = SRC_OBJECT;
  SPDLOG_INFO("Loading object file: {}", objPath);
  if (!objReader.load(objPath)) {
    spdlog::error("Failed to load object file: {}", objPath);
    return 1;
  }

  // Load the source PE file
  COFFI::coffi peReader;
  SPDLOG_INFO("Loading PE file: {}", inputFile);
  if (!peReader.load(inputFile)) {
    spdlog::error("Failed to load PE file: {}", inputFile);
    return 1;
  }

  // Find the 'main' function in the object file
  auto &symbols = *objReader.get_symbols();
  COFFI::symbol *mainSymbol = nullptr;
  for (auto &sym : symbols) {
    SPDLOG_INFO("Symbol: {}", sym.get_name());
    if (sym.get_name() == "_ref") {
      mainSymbol = &sym;
      break;
    }
  }

  if (!mainSymbol) {
    spdlog::error("Could not find 'main' symbol in object file");
    return 1;
  }

  // Get the section containing the main function
  auto &sections = objReader.get_sections();
  auto mainSection = sections[mainSymbol->get_section_number() - 1];

  // Calculate main function size using next symbol method
  uint32_t mainOffset = mainSymbol->get_value();
  uint32_t mainSize = 0;

  // Find the next symbol in the same section to calculate size
  uint32_t nextSymbolOffset = UINT32_MAX;
  for (auto &sym : symbols) {
    if (sym.get_section_number() == mainSymbol->get_section_number() &&
        sym.get_value() > mainOffset && sym.get_value() < nextSymbolOffset) {
      nextSymbolOffset = sym.get_value();
    }
  }

  if (nextSymbolOffset != UINT32_MAX) {
    mainSize = nextSymbolOffset - mainOffset;
    spdlog::info(
        "Calculated main function size: {} bytes (next symbol at offset {})",
        mainSize, nextSymbolOffset);
  } else {
    // If no next symbol found, use remaining section size
    mainSize = mainSection->get_data_size() - mainOffset;
    spdlog::info("No next symbol found, using remaining section size: {} bytes",
                 mainSize);
  }

  auto mainCodeData = mainSection->get_data();

  spdlog::info("Found main function at offset {} with size {}", mainOffset,
               mainSize);

  // List relocations for the main function section
  auto& sectionRelocations = mainSection->get_relocations();
  spdlog::info("Relocations for section {} (containing main function):", mainSymbol->get_section_number());
  for (auto& reloc : sectionRelocations) {
    spdlog::info("  Relocation at offset 0x{:x}, type: {}, symbol index: {}", 
                 reloc.get_virtual_address(), 
                 reloc.get_type(), 
                 reloc.get_symbol_table_index());
    
    // Get symbol name for this relocation
    spdlog::info("    -> Symbol: {}", reloc.get_symbol());
  }

  spdlog::info("Main function code:");
  std::string s;
  for (uint32_t i = 0; i < mainSize; i++) {
    if (i > 0 && i % 16 == 0) {
      spdlog::info("{}", s);
      s.clear();
    }
    if (s.size() > 0)
      s += " ";
    s += fmt::format("{:02X}", mainCodeData[i]);
  }
  if (s.size() > 0)
    spdlog::info("{}", s);

  // Find .text section in PE file
  auto &peSections = peReader.get_sections();
  COFFI::section *textSection = nullptr;
  for (auto &section : peSections) {
    if (section->get_name() == ".text") {
      textSection = section;
      break;
    }
  }

  if (!textSection) {
    spdlog::error("Could not find .text section in PE file");
    return 1;
  }

  // Get image base and calculate actual VA
  uint64_t imageBase = 0;
  auto winHeader = peReader.get_win_header();
  if (winHeader) {
    imageBase = winHeader->get_image_base();
  }
  
  uint32_t textSectionRVA = textSection->get_virtual_address();
  uint32_t textSectionSize = textSection->get_virtual_size();
  uint64_t textSectionVA = imageBase + textSectionRVA;
  uint64_t textSectionEndVA = textSectionVA + textSectionSize;
  
  spdlog::info("PE Image base: 0x{:x}", imageBase);
  spdlog::info(".text section RVA: 0x{:x}, size: 0x{:x}", textSectionRVA, textSectionSize);
  spdlog::info(".text section VA: 0x{:x} - 0x{:x}", textSectionVA, textSectionEndVA);

  // Find available space at the end of .text section (look for null bytes)
  auto textSectionData = textSection->get_data();
  uint32_t originalTextSize = textSection->get_data_size();
  
  // Search backwards from the end to find contiguous null bytes
  uint32_t availableSpace = 0;
  for (int32_t i = originalTextSize - 1; i >= 0; i--) {
    if (reinterpret_cast<const uint8_t*>(textSectionData)[i] == 0x00) {
      availableSpace++;
    } else {
      break; // Found non-null byte, stop counting
    }
  }
  
  spdlog::info("Found {} bytes of available space (null bytes) at end of .text section", availableSpace);
  
  if (availableSpace < mainSize) {
    spdlog::error("Not enough space in .text section! Need {} bytes, found {} bytes", 
                  mainSize, availableSpace);
    return 1;
  }
  
  // Calculate injection offset (place code at start of null space)
  uint32_t injectionOffset = originalTextSize - availableSpace;
  uint64_t injectionVA = textSectionVA + injectionOffset;
  
  spdlog::info("Injecting {} bytes at .text section offset 0x{:x} (VA: 0x{:x})", 
               mainSize, injectionOffset, injectionVA);

  // Copy the main function code into the available space
  const uint8_t* mainCode = reinterpret_cast<const uint8_t*>(mainCodeData) + mainOffset;
  
  // Create a copy of the section data to modify
  std::vector<uint8_t> newTextData(reinterpret_cast<const uint8_t*>(textSectionData), 
                                   reinterpret_cast<const uint8_t*>(textSectionData) + originalTextSize);
  
  // Copy our code into the null space
  std::memcpy(newTextData.data() + injectionOffset, mainCode, mainSize);
  
  // Now handle relocations - resolve imports and copy constants to .rdata
  spdlog::info("Processing relocations...");
  
  // Find .rdata section for constants
  COFFI::section* rdataSection = nullptr;
  for (auto& section : peSections) {
    if (section->get_name() == ".rdata") {
      rdataSection = section;
      break;
    }
  }
  
  if (!rdataSection) {
    spdlog::error("Could not find .rdata section in PE file");
    return 1;
  }
  
  // Find available space in .rdata section (similar to .text)
  auto rdataData = rdataSection->get_data();
  uint32_t rdataSize = rdataSection->get_data_size();
  uint32_t rdataAvailableSpace = 0;
  
  for (int32_t i = rdataSize - 1; i >= 0; i--) {
    if (reinterpret_cast<const uint8_t*>(rdataData)[i] == 0x00) {
      rdataAvailableSpace++;
    } else {
      break;
    }
  }
  
  spdlog::info("Found {} bytes of available space in .rdata section", rdataAvailableSpace);
  
  // Create rdata copy for modifications
  std::vector<uint8_t> newRdataData(reinterpret_cast<const uint8_t*>(rdataData), 
                                    reinterpret_cast<const uint8_t*>(rdataData) + rdataSize);
  uint32_t rdataInjectionOffset = rdataSize - rdataAvailableSpace;
  uint32_t rdataCurrentOffset = rdataInjectionOffset;
  
  // Get import table information for resolving __imp_ symbols
  auto& directories = peReader.get_directories();
  uint32_t importRVA = 0;
  if (directories.size() > 1) {
    importRVA = directories[1]->get_virtual_address(); // Import table is directory entry 1
  }
  
  // Process each relocation
  for (auto& reloc : sectionRelocations) {
    if (reloc.get_type() != 6) continue; // Only handle type 6 (IMAGE_REL_I386_DIR32)
    
    uint32_t relocOffset = reloc.get_virtual_address();
    std::string symbolName = reloc.get_symbol();
    uint32_t resolvedAddress = 0;
    
    spdlog::info("Processing relocation at offset 0x{:x} for symbol: {}", relocOffset, symbolName);
    
    if (symbolName.starts_with("__imp_")) {
      // This is an import symbol - find it in the PE's import table
      std::string functionName = symbolName.substr(6); // Remove "__imp_" prefix
      
      // Remove @N suffix for stdcall functions
      size_t atPos = functionName.find('@');
      if (atPos != std::string::npos) {
        functionName = functionName.substr(0, atPos);
      }
      
      spdlog::info("Looking for import function: {}", functionName);
      
      // For now, we'll need to implement import table parsing
      // This is a placeholder - you'd need to parse the import table to find the actual IAT address
      spdlog::warn("Import resolution not fully implemented yet for: {}", functionName);
      resolvedAddress = 0x12345678; // Placeholder
      
    } else if (symbolName.starts_with("??_C@")) {
      // This is a string constant - find it in the object file and copy to .rdata
      COFFI::symbol* constSymbol = nullptr;
      for (auto& sym : symbols) {
        if (sym.get_name() == symbolName) {
          constSymbol = &sym;
          break;
        }
      }
      
      if (constSymbol) {
        // Get the string data from the object file
        auto constSection = sections[constSymbol->get_section_number() - 1];
        auto constData = constSection->get_data();
        uint32_t constOffset = constSymbol->get_value();
        
        // Find the string length (null-terminated)
        uint32_t stringLen = 0;
        while (constData[constOffset + stringLen] != 0) {
          stringLen++;
        }
        stringLen++; // Include null terminator
        
        if (rdataCurrentOffset + stringLen > rdataSize) {
          spdlog::error("Not enough space in .rdata for string constant");
          return 1;
        }
        
        // Copy string to .rdata
        std::memcpy(newRdataData.data() + rdataCurrentOffset, 
                   constData + constOffset, stringLen);
        
        // Calculate the new address
        uint32_t rdataRVA = rdataSection->get_virtual_address();
        resolvedAddress = imageBase + rdataRVA + rdataCurrentOffset;
        
        spdlog::info("Copied string constant '{}' to .rdata at RVA 0x{:x} (VA: 0x{:x})", 
                     std::string(constData + constOffset), 
                     rdataRVA + rdataCurrentOffset, resolvedAddress);
        
        rdataCurrentOffset += stringLen;
      }
    }
    
    // Apply the relocation to the injected code
    if (resolvedAddress != 0) {
      uint32_t patchOffset = injectionOffset + relocOffset;
      *reinterpret_cast<uint32_t*>(newTextData.data() + patchOffset) = resolvedAddress;
      spdlog::info("Applied relocation: patched offset 0x{:x} with address 0x{:x}", 
                   patchOffset, resolvedAddress);
    }
  }
  
  // Update both sections
  rdataSection->set_data(reinterpret_cast<const char*>(newRdataData.data()), newRdataData.size());
  
  // Update the .text section with modified data (same size)
  textSection->set_data(reinterpret_cast<const char*>(newTextData.data()), newTextData.size());
  
  spdlog::info("Injected code into existing .text section space (size unchanged: 0x{:x})", originalTextSize);
  
  // Save the modified PE file to output path
  spdlog::info("Saving patched PE file to: {}", outputFile);
  if (!peReader.save(outputFile)) {
    spdlog::error("Failed to save patched PE file to: {}", outputFile);
    return 1;
  }
  
  spdlog::info("Successfully patched PE file! Main function injected at VA: 0x{:x}", injectionVA);

  return 0;
}