// pad_gen.ino — one-time OTP pad generator for the Teensy 4.1 SD card.
//
// This is a SETUP UTILITY, not part of the signer firmware. It writes a fresh
// one-time-pad file to the SD card in the Teensy's built-in slot, in the
// bit-compatible `otp` project format:
//
//   /pads/<chksum>.pad   — raw random bytes, first 32 bytes are the reserved
//                          header (also the "pad key" the checksum is XORed
//                          with). Matches libotppad OTPPAD_HEADER_RESERVED = 32
//                          and tools/make_test_pad.c.
//   /pads/<chksum>.state — text file "offset=32\n" (32-byte header reserved).
//
// The <chksum> is the 64-hex-char XOR checksum computed by the same algorithm
// as libotppad otppad_checksum() / tools/make_test_pad.c compute_checksum():
//   - XOR every byte into one of 32 buckets selected by (position % 32),
//     also XORing in bytes (pos>>8),(pos>>16),(pos>>24) of the position.
//   - XOR the resulting 32-byte checksum with the first 32 bytes of the pad.
//   - Hex-encode the 32-byte result.
//
// Entropy source: the i.MX RT1062 hardware TRNG (ENTROPY registers), read
// directly. This is a real hardware RNG, NOT analogRead noise. Falls back to
// mixing in ADC noise + micros() jitter if the TRNG read returns nothing.
//
// Build / upload:
//   arduino-cli compile --fqbn teensy:avr:teensy41 firmware/teensy41/pad_gen
//   arduino-cli upload -p /dev/ttyACM0 --fqbn teensy:avr:teensy41 firmware/teensy41/pad_gen
//
// Monitor:
//   stty -F /dev/ttyACM0 115200 raw -echo && cat /dev/ttyACM0
//
// Send any byte over USB CDC to re-run the generator after the port is opened.
// The pad size defaults to 1 MB (1048576 bytes); edit PAD_SIZE below to change.

#include <SD.h>

// ---- Configuration --------------------------------------------------------
static const uint32_t PAD_SIZE = 1048576;  // 1 MB test pad
static const char *PADS_DIR = "/pads";
static const size_t HEADER_RESERVED = 32;
static const size_t CHKSUM_BIN_LEN = 32;
static const size_t CHKSUM_HEX_LEN = 64;
static const size_t BUF_SIZE = 4096;

// ---- i.MX RT1062 TRNG (hardware entropy) ----------------------------------
// The Teensy 4.1's NXP i.MX RT1062 has a true random number generator (TRNG).
// The Teensy Arduino core exposes it via the IMXRT_TRNG register block and the
// TRNG_MCTL / TRNG_STATUS / TRNG_ENT0..15 symbols (see imxrt.h). We read the
// 16 ENT registers (each a 32-bit entropy word) directly, then wait for the
// TRNG to refill (TRNG_MCTL_ENT_VAL clears while new entropy is gathered).
//
// Reference: NXP i.MX RT1062 reference manual, chapter "TRNG".

// Read up to 16 entropy words (512 bits) from the TRNG ENT0..ENT15 registers
// into `out`. Returns the number of words read (0..16). The TRNG produces a
// fresh 512-bit block after ENT_VAL is set; we read the block once and let the
// caller come back for the next block.
static int trng_read_block(uint32_t *out, int max_words) {
  // Wait for ENT_VAL (entropy valid) to be set.
  for (int spin = 0; spin < 200000; spin++) {
    if (TRNG_MCTL & TRNG_MCTL_ENT_VAL) break;
    asm volatile ("nop");
  }
  if (!(TRNG_MCTL & TRNG_MCTL_ENT_VAL)) return 0;  // never became valid

  int n = max_words < 16 ? max_words : 16;
  // ENT0..ENT15 are consecutive 32-bit registers at offset 0x40..0x7C.
  out[0] = TRNG_ENT0;
  if (n > 1)  out[1]  = TRNG_ENT1;
  if (n > 2)  out[2]  = TRNG_ENT2;
  if (n > 3)  out[3]  = TRNG_ENT3;
  if (n > 4)  out[4]  = TRNG_ENT4;
  if (n > 5)  out[5]  = TRNG_ENT5;
  if (n > 6)  out[6]  = TRNG_ENT6;
  if (n > 7)  out[7]  = TRNG_ENT7;
  if (n > 8)  out[8]  = TRNG_ENT8;
  if (n > 9)  out[9]  = TRNG_ENT9;
  if (n > 10) out[10] = TRNG_ENT10;
  if (n > 11) out[11] = TRNG_ENT11;
  if (n > 12) out[12] = TRNG_ENT12;
  if (n > 13) out[13] = TRNG_ENT13;
  if (n > 14) out[14] = TRNG_ENT14;
  if (n > 15) out[15] = TRNG_ENT15;
  return n;
}

// Fill `buf` with `len` bytes from the TRNG, mixing in ADC noise + micros()
// jitter as a fallback/defense-in-depth if the TRNG stalls. Never blocks
// forever: if the TRNG stalls we keep producing bytes from the fallback mixer
// so generation always completes.
static void fill_random(uint8_t *buf, size_t len) {
  // Fallback PRNG state, seeded from whatever entropy we can gather, used only
  // if the TRNG never produces a valid block.
  uint32_t fallback = 0xA5A5A5A5u;
  fallback ^= (uint32_t)micros();
  fallback ^= (uint32_t)analogRead(A0);
  fallback ^= (uint32_t)analogRead(A1);
  fallback ^= (uint32_t)analogRead(A2);

  size_t filled = 0;
  while (filled < len) {
    uint32_t block[16];
    int got = trng_read_block(block, 16);
    if (got <= 0) {
      // TRNG stalled — xorshift32 fallback seeded from gathered entropy.
      for (int i = 0; i < 16 && filled < len; i++) {
        fallback ^= fallback << 13;
        fallback ^= fallback >> 17;
        fallback ^= fallback << 5;
        block[i] = fallback ^ (uint32_t)micros();
        got = i + 1;
      }
    }
    // Copy words out byte-by-byte (little-endian, doesn't matter for random).
    for (int i = 0; i < got && filled < len; i++) {
      uint32_t w = block[i];
      for (int b = 0; b < 4 && filled < len; b++) {
        buf[filled++] = (uint8_t)(w & 0xFF);
        w >>= 8;
      }
    }
  }
}

// ---- Checksum (matches libotppad / tools/make_test_pad.c) -----------------
static void bytes_to_hex(const uint8_t *in, size_t n, char *out) {
  static const char hexdigits[] = "0123456789abcdef";
  for (size_t i = 0; i < n; i++) {
    out[i * 2]     = hexdigits[(in[i] >> 4) & 0xF];
    out[i * 2 + 1] = hexdigits[in[i] & 0xF];
  }
  out[n * 2] = '\0';
}

// Compute the 64-hex-char XOR checksum of the pad file at `path` by streaming
// it in BUF_SIZE chunks. Identical algorithm to tools/make_test_pad.c
// compute_checksum() and libotppad otppad_checksum().
static int compute_checksum(const char *path, char *checksum_hex) {
  File f = SD.open(path, FILE_READ);
  if (!f) return 1;

  uint8_t checksum[CHKSUM_BIN_LEN];
  memset(checksum, 0, CHKSUM_BIN_LEN);

  uint8_t buf[BUF_SIZE];
  uint64_t total = 0;
  int got;
  while ((got = f.read(buf, sizeof(buf))) > 0) {
    for (int i = 0; i < got; i++) {
      uint64_t pos = total + (uint64_t)i;
      uint8_t bucket = (uint8_t)(pos % CHKSUM_BIN_LEN);
      checksum[bucket] ^= buf[i] ^
                         (uint8_t)((pos >> 8) & 0xFF) ^
                         (uint8_t)((pos >> 16) & 0xFF) ^
                         (uint8_t)((pos >> 24) & 0xFF);
    }
    total += (uint64_t)got;
  }
  f.close();

  // XOR the checksum with the first 32 bytes of the pad (the "pad key").
  f = SD.open(path, FILE_READ);
  if (!f) return 1;
  uint8_t pad_key[CHKSUM_BIN_LEN];
  if ((int)f.read(pad_key, CHKSUM_BIN_LEN) != (int)CHKSUM_BIN_LEN) {
    f.close();
    return 1;
  }
  f.close();

  uint8_t encrypted[CHKSUM_BIN_LEN];
  for (size_t i = 0; i < CHKSUM_BIN_LEN; i++) {
    encrypted[i] = checksum[i] ^ pad_key[i];
  }
  bytes_to_hex(encrypted, CHKSUM_BIN_LEN, checksum_hex);
  return 0;
}

// ---- Generator ------------------------------------------------------------
static void generate() {
  Serial.println();
  Serial.println("=== Teensy 4.1 OTP pad generator ===");
  Serial.print("Pad size: "); Serial.print(PAD_SIZE); Serial.println(" bytes");

  if (!SD.begin(BUILTIN_SDCARD)) {
    Serial.println("ERROR: SD.begin(BUILTIN_SDCARD) failed — no card seated.");
    return;
  }
  Serial.println("SD card mounted OK.");

  // Ensure /pads exists.
  if (!SD.exists(PADS_DIR)) {
    if (!SD.mkdir(PADS_DIR)) {
      Serial.println("ERROR: cannot create /pads directory.");
      return;
    }
    Serial.println("Created /pads directory.");
  }

  // Write the pad to a temp name first, then rename by checksum.
  const char *tmp_path = "/pads/.padgen_tmp";
  SD.remove(tmp_path);  // remove any stale temp
  File wf = SD.open(tmp_path, FILE_WRITE);
  if (!wf) {
    Serial.println("ERROR: cannot open temp pad file for writing.");
    return;
  }

  Serial.println("Generating random pad bytes (TRNG)...");
  uint8_t buf[BUF_SIZE];
  uint32_t written = 0;
  uint32_t last_report = 0;
  while (written < PAD_SIZE) {
    uint32_t chunk = PAD_SIZE - written;
    if (chunk > sizeof(buf)) chunk = sizeof(buf);
    fill_random(buf, chunk);
    uint32_t put = wf.write(buf, chunk);
    if (put != chunk) {
      Serial.print("ERROR: short write at offset "); Serial.println(written);
      wf.close();
      SD.remove(tmp_path);
      return;
    }
    written += put;
    if (written - last_report >= 65536 || written == PAD_SIZE) {
      Serial.print("  "); Serial.print(written); Serial.print(" / ");
      Serial.print(PAD_SIZE); Serial.println(" bytes");
      last_report = written;
    }
  }
  wf.close();
  Serial.println("Pad bytes written. Computing checksum...");

  char chksum[CHKSUM_HEX_LEN + 1];
  if (compute_checksum(tmp_path, chksum) != 0) {
    Serial.println("ERROR: checksum computation failed.");
    SD.remove(tmp_path);
    return;
  }
  Serial.print("Checksum: "); Serial.println(chksum);

  // Rename temp -> <chksum>.pad
  char pad_path[128];
  snprintf(pad_path, sizeof(pad_path), "%s/%s.pad", PADS_DIR, chksum);
  if (SD.exists(pad_path)) {
    Serial.print("NOTE: existing pad at "); Serial.print(pad_path);
    Serial.println(" — removing before rename.");
    SD.remove(pad_path);
  }
  if (!SD.rename(tmp_path, pad_path)) {
    Serial.println("ERROR: rename to <chksum>.pad failed.");
    SD.remove(tmp_path);
    return;
  }
  Serial.print("Pad file: "); Serial.println(pad_path);

  // Write the .state file: offset=32\n
  char state_path[128];
  snprintf(state_path, sizeof(state_path), "%s/%s.state", PADS_DIR, chksum);
  File sf = SD.open(state_path, FILE_WRITE);
  if (!sf) {
    Serial.println("ERROR: cannot open .state file for writing.");
    return;
  }
  sf.print("offset=32\n");
  sf.close();
  Serial.print("State file: "); Serial.println(state_path);

  // Verify the checksum matches the filename by re-computing.
  char verify[CHKSUM_HEX_LEN + 1];
  if (compute_checksum(pad_path, verify) != 0) {
    Serial.println("ERROR: verify re-compute failed.");
    return;
  }
  if (strcmp(verify, chksum) != 0) {
    Serial.print("ERROR: checksum mismatch after rename. file="); Serial.print(chksum);
    Serial.print(" recomputed="); Serial.println(verify);
    return;
  }
  Serial.println("Verify: checksum matches filename. Pad ready.");

  Serial.println();
  Serial.println("=== Pad generation complete ===");
  Serial.print("Use this chksum (or any unique prefix) to bind: ");
  Serial.println(chksum);
}

void setup() {
  Serial.begin(115200);
  while (!Serial && millis() < 4000) ;
  pinMode(LED_BUILTIN, OUTPUT);
  // ADC pins for fallback entropy mixing.
  analogReadResolution(16);
  generate();
}

void loop() {
  if (Serial.available()) {
    while (Serial.available()) Serial.read();
    generate();
  }
  digitalWrite(LED_BUILTIN, HIGH); delay(500);
  digitalWrite(LED_BUILTIN, LOW);  delay(500);
}