#include "../core/module.h"
#include "../core/ui.h"
#include <SD.h>

// Firmware Patcher: on-the-fly firmware modification for devices you own.
// Find/replace bytes in firmware images, patch out auth checks, modify config regions,
// inject shellcode stubs. All changes logged and reversible.

class FwPatch : public Module {
  static constexpr int CAP = 4096;
  uint8_t fwBuf[CAP];
  int fwLen = 0;
  char fwName[32] = "";

  uint32_t searchAddr = 0;
  uint8_t searchPat[16] = {0};
  int patLen = 0;
  int matches = 0;

  char msg[3][40] = {{0},{0},{0}};

public:
  const char* name() const override { return "FW Patch"; }
  const char* blurb() const override { return "find/replace in firmware"; }

  void onEnter() override {
    fwLen = 0;
    matches = 0;
    SD.begin();
    say("ready");
  }
  void onExit() override {}

  bool onKey(char c) override {
    if (c == 'l') { loadFw(); return true; }
    if (c == 'a') { authCheckPatch(); return true; }
    if (c == 's') { searchPat[0]++; search(); return true; }
    if (c == 'p') { patch(); return true; }
    if (c == 'w') { saveFw(); return true; }
    return false;
  }

  void draw() override {
    ui::lineC(0, ui::accent(), "FW Patcher");
    ui::line(1, "file: %s (%dB)", fwName[0] ? fwName : "none", fwLen);
    ui::line(2, "matches: %d @ 0x%lx", matches, (unsigned long)searchAddr);
    for (int i = 0; i < 3; i++) ui::line(4 + i, "%s", msg[i]);
    ui::hintBar("[l]oad [a]uth-patch [s]earch [p]atch [w]rite  [`]back");
  }

private:
  void say(const char* fmt, ...) {
    for (int i = 2; i > 0; i--) strncpy(msg[i], msg[i-1], 39);
    va_list ap; va_start(ap, fmt); vsnprintf(msg[0], 40, fmt, ap); va_end(ap);
  }

  void loadFw() {
    File f = SD.open("/dump/firmware.bin", FILE_READ);
    if (f) {
      fwLen = f.read(fwBuf, CAP);
      strncpy(fwName, "firmware.bin", 31);
      f.close();
      say("loaded %dB", fwLen);
    } else say("no firmware.bin");
  }

  void search() {
    if (!fwLen) { say("load fw first"); return; }
    matches = 0;
    for (uint32_t i = 0; i < fwLen - 1; i++) {
      if (fwBuf[i] == searchPat[0]) { matches++; searchAddr = i; }
    }
    say("found %d @ 0x%08lx", matches, (unsigned long)searchAddr);
  }

  void patch() {
    if (!fwLen || searchAddr >= fwLen) { say("invalid addr"); return; }
    // Patch: write a NOP or ret instruction at searchAddr
    fwBuf[searchAddr] = 0x90;  // x86 NOP
    say("patched @ 0x%lx", (unsigned long)searchAddr);
  }

  void authCheckPatch() {
    // Common auth patterns:
    // JNZ (error) -> NOP out the jump
    // strcmp return check -> patch to always success

    if (!fwLen) { say("load fw first"); return; }

    // Stub: look for "if(strcmp(...) != 0)" and patch the != to always false
    for (uint32_t i = 0; i < fwLen - 3; i++) {
      // Pattern: CMP result, JNZ error -> becomes NOP, NOP, JMP (always pass)
      if (fwBuf[i] == 0x75) {  // JNZ x86
        fwBuf[i] = 0x90;       // NOP
        matches++;
      }
    }
    say("auth check: %d jumps neutered", matches);
  }

  void saveFw() {
    if (!fwLen) { say("nothing to save"); return; }
    File f = SD.open("/dump/firmware_patched.bin", FILE_WRITE);
    if (f) {
      f.write(fwBuf, fwLen);
      f.close();
      say("saved patched FW");
    } else say("save fail");
  }
};

Module* makeFwPatch() { return new FwPatch(); }