#include <gtest/gtest.h>
#include <array>
#include <vector>
#include <cstring>

#include "helpers/ota/MotaContainer.h"
#include "helpers/ota/MerkleTree.h"
#include "helpers/ota/BlockBitmap.h"
#include "helpers/ota/Multihash.h"
#include "helpers/ota/FirmwareInfo.h"
#include "helpers/ota/SignerAllowlist.h"
#include "helpers/ota/OtaStore.h"
#include "helpers/ota/OtaProtocol.h"
#include "helpers/ota/OtaManager.h"
#include "mota_vectors.h"   // auto-generated by tools/mota/gen_vectors.py

extern "C" {
  #include "helpers/ota/detools/detools.h"   // vendored detools 0.53.0 embeddable decoder
}

using namespace mesh::ota;

// Build a flashed-image layout (body || fixed 56-byte EndF) the way the host packager / build hook do:
// marker(4) body_len(4) body_hash8(8) fw_version(4) target_id(4) hw_id(32). Identity is always present
// (zero/"" = unknown).
static std::vector<uint8_t> make_image_id(const std::vector<uint8_t>& body, uint32_t fw_version,
                                          uint32_t target_id, const char* hw_id) {
  std::vector<uint8_t> img = body;
  img.insert(img.end(), ENDF_MAGIC, ENDF_MAGIC + 4);
  uint32_t n = (uint32_t)body.size();
  for (int i = 0; i < 4; i++) img.push_back((uint8_t)(n >> (8 * i)));
  uint8_t h[8]; mh8(h, body.data(), body.size());
  img.insert(img.end(), h, h + 8);
  for (int i = 0; i < 4; i++) img.push_back((uint8_t)(fw_version >> (8 * i)));
  for (int i = 0; i < 4; i++) img.push_back((uint8_t)(target_id  >> (8 * i)));
  uint8_t hw[32] = {0}; size_t k = hw_id ? strlen(hw_id) : 0; if (k > 32) k = 32; if (k) memcpy(hw, hw_id, k);
  img.insert(img.end(), hw, hw + 32);                      // -> fixed 56-byte trailer
  return img;
}
static std::vector<uint8_t> make_image(const std::vector<uint8_t>& body) {
  return make_image_id(body, 0, 0, "");                    // zero identity (still a full 56-byte trailer)
}

// --- cross-check the C++ parser/merkle against the Python reference vectors ----------------

TEST(OtaParse, ParsesReferenceContainer) {
  MotaManifest m;
  ASSERT_TRUE(mota_parse(MOTA_VEC, MOTA_VEC_LEN, m));
  EXPECT_EQ(m.format_ver, MOTA_FORMAT_VER);
  EXPECT_TRUE(m.is_full());
  EXPECT_FALSE(m.is_signed());
  EXPECT_EQ(m.target_id, EXP_TARGET_ID);
  EXPECT_EQ(m.fw_version, EXP_FW_VERSION);
  EXPECT_EQ(m.image_size, EXP_IMAGE_SIZE);
  EXPECT_EQ(m.payload_size, EXP_PAYLOAD_SIZE);
  EXPECT_EQ(m.block_count, EXP_BLOCK_COUNT);
  EXPECT_EQ(m.block_size_log2, EXP_BLOCK_SIZE_LOG2);
  EXPECT_EQ(m.codec_id, EXP_CODEC_ID);
  EXPECT_EQ(0, memcmp(m.merkle_root, EXP_MERKLE_ROOT, 4));
  EXPECT_EQ(0, memcmp(m.image_hash, EXP_IMAGE_HASH, 32));
  ASSERT_NE(m.hw_id, nullptr);
  EXPECT_EQ(0, memcmp(m.hw_id, EXP_HW_ID, 32));        // v2 hardware tag ("TESTHW" NUL-padded)
  EXPECT_EQ(0, memcmp(m.approval, APPROVAL_NOT, 4));   // distributed = not approved
  EXPECT_FALSE(m.is_approved());
}

TEST(OtaParse, RejectsTampering) {
  MotaManifest m;
  // bad magic
  std::vector<uint8_t> b(MOTA_VEC, MOTA_VEC + MOTA_VEC_LEN);
  b[0] ^= 0xFF;
  EXPECT_FALSE(mota_parse(b.data(), b.size(), m));
  // bad trailer
  b.assign(MOTA_VEC, MOTA_VEC + MOTA_VEC_LEN);
  b[b.size() - 1] ^= 0xFF;
  EXPECT_FALSE(mota_parse(b.data(), b.size(), m));
  // wrong total-size field
  b.assign(MOTA_VEC, MOTA_VEC + MOTA_VEC_LEN);
  b[4] ^= 0x01;
  EXPECT_FALSE(mota_parse(b.data(), b.size(), m));
}

// block_idx is a uint16 on the wire, so a manifest needing > 65535 blocks can't be addressed and must be
// rejected at parse (this also keeps block_count*4 from overflowing the leaves-length computation).
TEST(OtaParse, RejectsTooManyBlocks) {
  auto manifest = [](uint32_t payload_size, uint8_t bsl) {           // fixed-layout unsigned-full manifest
    std::vector<uint8_t> m(MOTA_MFL, 0);
    m[0] = MOTA_FORMAT_VER; m[1] = MFLAG_FULL; m[2] = 0x12;
    m[15] = payload_size; m[16] = payload_size >> 8; m[17] = payload_size >> 16; m[18] = payload_size >> 24;
    m[19] = bsl;
    return m;
  };
  MotaManifest mm;
  auto over = manifest(65536u * 1024u, 10);                          // 65536 blocks -> rejected
  EXPECT_FALSE(mota_parse_manifest(over.data(), over.size(), mm));
  auto ok = manifest(65535u * 1024u, 10);                            // 65535 blocks -> allowed
  EXPECT_TRUE(mota_parse_manifest(ok.data(), ok.size(), mm));
  EXPECT_EQ(mm.block_count, 65535u);
}

TEST(OtaMerkle, RootMatchesVectorAndLeaves) {
  MotaManifest m;
  ASSERT_TRUE(mota_parse(MOTA_VEC, MOTA_VEC_LEN, m));
  // root recomputed from stored leaves[] == merkle_root field == Python's EXP_MERKLE_ROOT
  uint8_t root[4];
  merkle_root(root, m.leaves, m.block_count);
  EXPECT_EQ(0, memcmp(root, EXP_MERKLE_ROOT, 4));
  EXPECT_TRUE(mota_check_root(m));
  // recompute each leaf from the payload block and compare to the stored leaf
  uint32_t bs = m.block_size();
  for (uint32_t i = 0; i < m.block_count; i++) {
    uint32_t off = i * bs;
    uint32_t len = (off + bs <= m.payload_size) ? bs : (m.payload_size - off);
    uint8_t leaf[4];
    merkle_leaf(leaf, m.payload + off, len);
    EXPECT_EQ(0, memcmp(leaf, m.leaves + i * 4, 4)) << "leaf " << i;
  }
}

TEST(OtaMerkle, FullImageHashMatches) {
  MotaManifest m;
  ASSERT_TRUE(mota_parse(MOTA_VEC, MOTA_VEC_LEN, m));
  EXPECT_TRUE(mota_check_image_hash_full(m));
}

TEST(OtaMerkle, ProofFromReferenceVerifies) {
  MotaManifest m;
  ASSERT_TRUE(mota_parse(MOTA_VEC, MOTA_VEC_LEN, m));
  uint32_t bs = m.block_size();
  uint32_t off = PROOF_INDEX * bs;
  uint32_t len = (off + bs <= m.payload_size) ? bs : (m.payload_size - off);

  EXPECT_TRUE(merkle_verify(m.payload + off, len, PROOF_INDEX,
                            PROOF_SIBLINGS, PROOF_NSIB, EXP_MERKLE_ROOT, m.block_count));

  // tampered block -> fails
  std::vector<uint8_t> blk(m.payload + off, m.payload + off + len);
  blk[0] ^= 0xFF;
  EXPECT_FALSE(merkle_verify(blk.data(), len, PROOF_INDEX,
                             PROOF_SIBLINGS, PROOF_NSIB, EXP_MERKLE_ROOT, m.block_count));

  // wrong index with the same proof -> fails
  EXPECT_FALSE(merkle_verify(m.payload + off, len, PROOF_INDEX + 1,
                             PROOF_SIBLINGS, PROOF_NSIB, EXP_MERKLE_ROOT, m.block_count));
}

// --- validate the O(log n) binary-counter root vs a plain level-by-level reference ----------

static void ref_root(uint8_t out[4], std::vector<std::array<uint8_t, 4>> level) {
  while (level.size() > 1) {
    std::vector<std::array<uint8_t, 4>> nxt;
    for (size_t i = 0; i < level.size(); i += 2) {
      if (i + 1 < level.size()) {
        std::array<uint8_t, 4> p;
        merkle_combine(p.data(), level[i].data(), level[i + 1].data());
        nxt.push_back(p);
      } else {
        nxt.push_back(level[i]);  // promote lone last node
      }
    }
    level.swap(nxt);
  }
  std::memcpy(out, level[0].data(), 4);
}

TEST(OtaMerkle, BinaryCounterMatchesLevelByLevel) {
  uint32_t state = 0x12345678;
  auto rnd = [&]() { state = state * 1103515245u + 12345u; return (uint8_t)(state >> 16); };
  // O(log n) root must equal the plain level-by-level root for every count
  for (uint32_t count = 1; count <= 600; count++) {
    std::vector<uint8_t> leaves(count * 4);
    std::vector<std::array<uint8_t, 4>> ref(count);
    for (uint32_t i = 0; i < count; i++)
      for (int j = 0; j < 4; j++) { uint8_t v = rnd(); leaves[i * 4 + j] = v; ref[i][j] = v; }
    uint8_t a[4], b[4];
    merkle_root(a, leaves.data(), count);
    ref_root(b, ref);
    ASSERT_EQ(0, std::memcmp(a, b, 4)) << "root mismatch count=" << count;
  }
}

// Verify every block's proof for several tricky counts, using proofs generated by the Python
// reference (the oracle) — covers deep promotion chains (100, 255, 256, ...).
TEST(OtaMerkle, ReferenceProofsAllIndices) {
  for (int c = 0; c < N_PROOF_CASES; c++) {
    const ProofCase& pc = PROOF_CASES[c];
    uint8_t root[4];
    merkle_root(root, pc.leaves, pc.count);
    EXPECT_EQ(0, std::memcmp(root, pc.root, 4)) << "root mismatch count=" << pc.count;
    for (uint32_t i = 0; i < pc.count; i++) {
      EXPECT_TRUE(merkle_verify_from_leaf(pc.leaves + i * 4, i,
                                          pc.pblob + pc.poff[i], pc.pnsib[i], pc.root, pc.count))
          << "count=" << pc.count << " idx=" << i;
    }
    // a wrong sibling for index 0 must fail
    if (pc.pnsib[0] > 0) {
      std::vector<uint8_t> bad(pc.pblob + pc.poff[0], pc.pblob + pc.poff[0] + pc.pnsib[0] * 4);
      bad[0] ^= 0xFF;
      EXPECT_FALSE(merkle_verify_from_leaf(pc.leaves, 0, bad.data(), pc.pnsib[0], pc.root, pc.count));
    }
  }
}

// --- availability bitmap (derived from leaves[]) -------------------------------------------

TEST(OtaBitmap, AllPresentForCompleteContainer) {
  MotaManifest m;
  ASSERT_TRUE(mota_parse(MOTA_VEC, MOTA_VEC_LEN, m));
  EXPECT_TRUE(all_present(m.leaves, m.block_count));
  EXPECT_EQ(count_present(m.leaves, m.block_count), m.block_count);

  // a leaf slot of all-FF (erased) means "missing"; bitmap round-trips
  std::vector<uint8_t> leaves(m.leaves, m.leaves + m.block_count * 4);
  std::memset(&leaves[4], 0xFF, 4);  // mark block 1 missing
  EXPECT_FALSE(leaf_present(leaves.data(), 1));
  EXPECT_FALSE(all_present(leaves.data(), m.block_count));
  EXPECT_EQ(count_present(leaves.data(), m.block_count), m.block_count - 1);

  std::vector<uint8_t> bm(bitmap_bytes(m.block_count));
  leaves_to_bitmap(leaves.data(), m.block_count, bm.data());
  EXPECT_FALSE(bitmap_get(bm.data(), 1));
  EXPECT_TRUE(bitmap_get(bm.data(), 0));
}

// --- EndF self-firmware scan (P2) -----------------------------------------------------------

TEST(OtaFirmwareInfo, FindsEndFInImage) {
  std::vector<uint8_t> body(4321);
  for (size_t i = 0; i < body.size(); i++) body[i] = (uint8_t)(i * 37 + 11);
  std::vector<uint8_t> img = make_image(body);

  // simulate a flash region: image, then erased 0xFF up to the partition end
  std::vector<uint8_t> region = img;
  region.resize(img.size() + 4096, 0xFF);

  SelfFwInfo fi;
  ASSERT_TRUE(find_self_firmware(region.data(), (uint32_t)region.size(), fi, /*verify_body=*/true));
  EXPECT_EQ(fi.body_len, body.size());
  EXPECT_EQ(fi.image_len, img.size());
  EXPECT_EQ(fi.endf_offset, body.size());
  uint8_t h[8]; mh8(h, body.data(), body.size());
  EXPECT_EQ(0, std::memcmp(fi.body_hash, h, 8));
}

// The self-describing identity lives at fixed offsets in the 56-byte EndF and is always parsed; a
// zero-identity trailer reports zero/"" (unknown), still at the fixed 56-byte size.
TEST(OtaFirmwareInfo, ParsesIdentity) {
  std::vector<uint8_t> body(2000);
  for (size_t i = 0; i < body.size(); i++) body[i] = (uint8_t)(i * 13 + 5);

  auto img = make_image_id(body, 0x01100000u, 0x04d413fdu, "RAK4631");
  std::vector<uint8_t> region = img; region.resize(img.size() + 4096, 0xFF);
  SelfFwInfo fi;
  ASSERT_TRUE(find_self_firmware(region.data(), (uint32_t)region.size(), fi, /*verify_body=*/true));
  EXPECT_EQ(fi.body_len, body.size());
  EXPECT_EQ(fi.image_len, body.size() + 56);               // fixed trailer length
  EXPECT_EQ(fi.fw_version, 0x01100000u);
  EXPECT_EQ(fi.target_id, 0x04d413fdu);
  EXPECT_STREQ(fi.hw_id, "RAK4631");

  auto img1 = make_image(body);                            // zero-identity trailer (still 56 bytes)
  std::vector<uint8_t> r1 = img1; r1.resize(img1.size() + 64, 0xFF);
  SelfFwInfo fi1;
  ASSERT_TRUE(find_self_firmware(r1.data(), (uint32_t)r1.size(), fi1, true));
  EXPECT_EQ(fi1.fw_version, 0u);
  EXPECT_EQ(fi1.target_id, 0u);
  EXPECT_STREQ(fi1.hw_id, "");
  EXPECT_EQ(fi1.image_len, body.size() + 56);
}

TEST(OtaFirmwareInfo, IgnoresStagedMotaHigherInRegion) {
  // The firmware's own EndF must win even when a staged .mota (which embeds its own EndF) sits
  // above it in the same region — the body_len == offset check disambiguates.
  std::vector<uint8_t> body(2000);
  for (size_t i = 0; i < body.size(); i++) body[i] = (uint8_t)(i ^ 0x5A);
  std::vector<uint8_t> img = make_image(body);

  std::vector<uint8_t> region = img;
  region.resize(8192, 0xFF);                                  // gap
  // drop the reference .mota (which contains an embedded EndF in its payload) higher up
  region.insert(region.end(), MOTA_VEC, MOTA_VEC + MOTA_VEC_LEN);

  SelfFwInfo fi;
  ASSERT_TRUE(find_self_firmware(region.data(), (uint32_t)region.size(), fi, true));
  EXPECT_EQ(fi.endf_offset, body.size());                     // found OUR firmware, not the .mota's
  EXPECT_EQ(fi.body_len, body.size());
}

TEST(OtaFirmwareInfo, NoMarkerReturnsFalse) {
  std::vector<uint8_t> region(1000, 0xAB);
  SelfFwInfo fi;
  EXPECT_FALSE(find_self_firmware(region.data(), (uint32_t)region.size(), fi));
}

// --- signer allowlist (P3) ------------------------------------------------------------------

TEST(OtaAllowlist, AddContainsRemoveSerialize) {
  SignerAllowlist a;
  uint8_t k1[32], k2[32], k3[32];
  memset(k1, 0x11, 32); memset(k2, 0x22, 32); memset(k3, 0x33, 32);
  EXPECT_FALSE(a.contains(k1));
  EXPECT_TRUE(a.add(k1));
  EXPECT_TRUE(a.add(k2));
  EXPECT_TRUE(a.add(k1));            // idempotent
  EXPECT_EQ(a.count(), 2);
  EXPECT_TRUE(a.contains(k1));
  EXPECT_FALSE(a.contains(k3));

  uint8_t buf[1 + MAX_OTA_SIGNERS * 32];
  uint32_t n = a.serialize(buf, sizeof(buf));
  EXPECT_EQ(n, 1u + 2 * 32);
  SignerAllowlist b;
  EXPECT_TRUE(b.deserialize(buf, n));
  EXPECT_EQ(b.count(), 2);
  EXPECT_TRUE(b.contains(k1) && b.contains(k2));

  EXPECT_TRUE(a.remove(k1));
  EXPECT_EQ(a.count(), 1);
  EXPECT_FALSE(a.contains(k1));
  EXPECT_TRUE(a.contains(k2));
}

// --- RAM store: out-of-order writes + availability via leaves[] --------------------------------

TEST(OtaStoreRamTest, RandomAccessAndErasedSentinel) {
  OtaStoreRam<4096> s;
  ASSERT_TRUE(s.begin(1000));
  EXPECT_EQ(s.staged_size(), 1000u);
  uint8_t blk[8] = {1,2,3,4,5,6,7,8};
  EXPECT_TRUE(s.write(500, blk, 8));   // out-of-order offset
  EXPECT_TRUE(s.write(0, blk, 8));
  EXPECT_FALSE(s.write(998, blk, 8));  // out of range
  uint8_t rd[8];
  EXPECT_TRUE(s.read(500, rd, 8));
  EXPECT_EQ(0, memcmp(rd, blk, 8));
  // untouched region reads as erased 0xFF (so an unfilled leaf slot is "missing")
  EXPECT_TRUE(s.read(100, rd, 8));
  for (int i = 0; i < 8; i++) EXPECT_EQ(rd[i], 0xFF);
}

// --- merkle proof GENERATION (server side) matches the Python oracle ---------------------------

TEST(OtaMerkle, GenProofMatchesPythonAndVerifies) {
  for (int c = 0; c < N_PROOF_CASES; c++) {
    const ProofCase& pc = PROOF_CASES[c];
    std::vector<uint8_t> scratch(pc.count * 4);
    uint8_t out[32 * 4];
    for (uint32_t i = 0; i < pc.count; i++) {
      uint8_t n = merkle_gen_proof(pc.leaves, pc.count, i, scratch.data(), out);
      ASSERT_EQ(n, pc.pnsib[i]) << "count=" << pc.count << " idx=" << i;
      EXPECT_EQ(0, std::memcmp(out, pc.pblob + pc.poff[i], (size_t)n * 4))
          << "gen_proof != python count=" << pc.count << " idx=" << i;
      EXPECT_TRUE(merkle_verify_from_leaf(pc.leaves + i * 4, i, out, n, pc.root, pc.count));
    }
  }
}

// --- protocol codec round-trips ---------------------------------------------------------------

TEST(OtaProtocol, CodecRoundTrips) {
  uint8_t buf[200];

  // OTA_ADV is now a tiny per-node beacon: seeder_id + n_motas + set_digest
  AdvMsg adv{{0x29,0x17,0xe4,0xf7}, 7, {0xde,0xad,0xbe,0xef}};
  uint16_t n = encode_adv(buf, sizeof(buf), adv);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_ADV); EXPECT_EQ(n, 10);
  AdvMsg a2; ASSERT_TRUE(decode_adv(buf, n, a2));
  EXPECT_EQ(0, memcmp(a2.seeder_id, adv.seeder_id, 4));
  EXPECT_EQ(a2.n_motas, 7);
  EXPECT_EQ(0, memcmp(a2.set_digest, adv.set_digest, 4));

  // OTA_QUERY: ask a source (by seeder_id) for the offering set_digest, optionally filtered to a target
  QueryMsg qy{{0x29,0x17,0xe4,0xf7}, {0xd1,0xd2,0xd3,0xd4}, 0x11223344};
  n = encode_query(buf, sizeof(buf), qy);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_QUERY);
  QueryMsg q2; ASSERT_TRUE(decode_query(buf, n, q2));
  EXPECT_EQ(0, memcmp(q2.seeder_id, qy.seeder_id, 4));
  EXPECT_EQ(0, memcmp(q2.set_digest, qy.set_digest, 4));
  EXPECT_EQ(q2.filter_target, 0x11223344u);

  // OTA_HAVE: a 2-row catalog (mid, target, fwver, codec, flags per row) tagged with the offering digest
  uint8_t rows[2 * OTA_HAVE_ROW_BYTES];
  for (int i = 0; i < 2 * OTA_HAVE_ROW_BYTES; i++) rows[i] = (uint8_t)(i + 1);
  HaveMsg hv{{0x29,0x17,0xe4,0xf7}, {0xd1,0xd2,0xd3,0xd4}, 0, 1, 2, rows};
  n = encode_have(buf, sizeof(buf), hv);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_HAVE);
  HaveMsg h2; ASSERT_TRUE(decode_have(buf, n, h2));
  EXPECT_EQ(0, memcmp(h2.seeder_id, hv.seeder_id, 4));
  EXPECT_EQ(0, memcmp(h2.set_digest, hv.set_digest, 4));
  EXPECT_EQ(h2.frag_total, 1); EXPECT_EQ(h2.n_rows, 2);
  EXPECT_EQ(0, memcmp(h2.rows, rows, 2 * OTA_HAVE_ROW_BYTES));

  GetManifestMsg gm{{1,2,3,4}, 0x0002};   // want only manifest fragment 1 (recovery)
  n = encode_get_manifest(buf, sizeof(buf), gm);
  GetManifestMsg g2; ASSERT_TRUE(decode_get_manifest(buf, n, g2));
  EXPECT_EQ(0, memcmp(g2.manifest_id, gm.manifest_id, 4));
  EXPECT_EQ(g2.want_mask, 0x0002);

  uint8_t mbytes[40]; for (int i = 0; i < 40; i++) mbytes[i] = (uint8_t)(i + 1);
  ManifestMsg mm{{9,8,7,6}, 0, 1, mbytes, 40};
  n = encode_manifest(buf, sizeof(buf), mm);
  ManifestMsg m2; ASSERT_TRUE(decode_manifest(buf, n, m2));
  EXPECT_EQ(m2.frag_idx, 0); EXPECT_EQ(m2.frag_total, 1); EXPECT_EQ(m2.len, 40);
  EXPECT_EQ(0, memcmp(m2.bytes, mbytes, 40));

  ReqMsg rq{{4,3,2,1}, 7, 0x005f};        // block 7, want fragments {0,1,2,3,4,6} (a recovery mask)
  n = encode_req(buf, sizeof(buf), rq);
  ReqMsg r2; ASSERT_TRUE(decode_req(buf, n, r2));
  EXPECT_EQ(r2.block_idx, 7); EXPECT_EQ(r2.want_mask, 0x005f);

  // GET_LEAVES: bulk-fetch the target leaves[] with a fragment want_mask (motatool warm-start)
  GetLeavesMsg gl{{5,6,7,8}, 0x0007};     // want leaves fragments {0,1,2}
  n = encode_get_leaves(buf, sizeof(buf), gl);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_GET_LEAVES);
  GetLeavesMsg gl2; ASSERT_TRUE(decode_get_leaves(buf, n, gl2));
  EXPECT_EQ(0, memcmp(gl2.manifest_id, gl.manifest_id, 4)); EXPECT_EQ(gl2.want_mask, 0x0007);

  // LEAVES: one fragment of the leaves[] array
  uint8_t lbytes[80]; for (int i = 0; i < 80; i++) lbytes[i] = (uint8_t)(200 - i);
  LeavesMsg lm{{5,6,7,8}, 1, 4, lbytes, 80};
  n = encode_leaves(buf, sizeof(buf), lm);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_LEAVES);
  LeavesMsg lm2; ASSERT_TRUE(decode_leaves(buf, n, lm2));
  EXPECT_EQ(lm2.frag_idx, 1); EXPECT_EQ(lm2.frag_total, 4); EXPECT_EQ(lm2.len, 80);
  EXPECT_EQ(0, memcmp(lm2.bytes, lbytes, 80));

  // DATA is one self-describing fragment of a block (frag_off places it; proof is fetched separately)
  uint8_t data[100]; for (int i = 0; i < 100; i++) data[i] = (uint8_t)(i * 3);
  DataMsg dm{{0,1,2,3}, 42, 0, data, 100};                  // block 42, fragment at offset 0
  n = encode_data(buf, sizeof(buf), dm);
  DataMsg d2; ASSERT_TRUE(decode_data(buf, n, d2));
  EXPECT_EQ(d2.block_idx, 42); EXPECT_EQ(d2.frag_off, 0);
  EXPECT_EQ(d2.data_len, 100); EXPECT_EQ(0, memcmp(d2.data, data, 100));

  // a later slice of the same block (non-zero frag_off)
  DataMsg dm2{{0,1,2,3}, 42, 160, data, 50};
  n = encode_data(buf, sizeof(buf), dm2);
  DataMsg d3; ASSERT_TRUE(decode_data(buf, n, d3));
  EXPECT_EQ(d3.block_idx, 42); EXPECT_EQ(d3.frag_off, 160); EXPECT_EQ(d3.data_len, 50);

  // REQ_PROOF: request the merkle proof for one (reassembled) block
  ReqProofMsg rp{{7,7,8,8}, 13};
  n = encode_req_proof(buf, sizeof(buf), rp);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_REQ_PROOF);
  ReqProofMsg rp2; ASSERT_TRUE(decode_req_proof(buf, n, rp2));
  EXPECT_EQ(0, memcmp(rp2.manifest_id, rp.manifest_id, 4)); EXPECT_EQ(rp2.block_idx, 13);

  // PROOF: ordered sibling digests for one block
  uint8_t proof[12]; for (int i = 0; i < 12; i++) proof[i] = (uint8_t)(0xA0 + i);
  ProofMsg pm{{7,7,8,8}, 13, 3, proof};
  n = encode_proof(buf, sizeof(buf), pm);
  ASSERT_GT(n, 0); EXPECT_EQ(ota_msg_type(buf, n), OTA_PROOF);
  ProofMsg pm2; ASSERT_TRUE(decode_proof(buf, n, pm2));
  EXPECT_EQ(0, memcmp(pm2.manifest_id, pm.manifest_id, 4));
  EXPECT_EQ(pm2.block_idx, 13); EXPECT_EQ(pm2.n_proof, 3);
  EXPECT_EQ(0, memcmp(pm2.proof, proof, 12));
}

// --- full transfer simulation between two OtaManagers (P4b) ------------------------------------

namespace {
struct SimMsg { OtaManager* dest; std::vector<uint8_t> bytes; };
static std::vector<SimMsg> g_q;
struct SendTo { OtaManager* dest; };
static void sim_send(void* ctx, const uint8_t* msg, uint16_t len, bool /*flood*/) {
  g_q.push_back({((SendTo*)ctx)->dest, std::vector<uint8_t>(msg, msg + len)});
}
// Drive the bus to quiescence: deliver queued messages; when idle, advance the client's clock (monotonic
// across calls, so a jittered query scheduled in a prior pump still comes due) and call loop() (fires the
// scheduled catalog query / block re-requests). Two idle ticks in a row = quiescent.
static uint32_t g_clk = 0;
static void pump(OtaManager& client, int guard_max = 200000) {
  int idle = 0, guard = 0;
  while (guard++ < guard_max) {
    if (!g_q.empty()) {
      SimMsg m = std::move(g_q.front()); g_q.erase(g_q.begin());
      m.dest->on_message(m.bytes.data(), (uint16_t)m.bytes.size());
      idle = 0;
    } else {
      g_clk += 5000; client.set_clock(g_clk); client.loop();
      if (!g_q.empty()) { idle = 0; continue; }
      if (++idle >= 2) break;
    }
  }
}

// A test MotaSource backing an external "folder" with one or more complete `.mota` images held in RAM —
// the simplest concrete transport (a real device uses serial/BLE/WiFi/FS, same interface). describe()
// parses each container for the catalog + region offsets; read() is a bounds-checked memcpy.
class RamMotaSource : public mesh::ota::MotaSource {
public:
  void add(const uint8_t* buf, uint32_t len) { if (_n < 8) { _buf[_n] = buf; _len[_n] = len; _n++; } }
  uint8_t count() override { return _n; }
  bool describe(uint8_t idx, mesh::ota::MotaDesc& d) override {
    if (idx >= _n) return false;
    MotaManifest m;
    if (!mota_parse(_buf[idx], _len[idx], m)) return false;
    std::memcpy(d.mid, m.merkle_root, 4);
    d.target_id = m.target_id; d.fw_version = m.fw_version;
    d.codec_id = m.codec_id; d.flags = m.flags;
    d.total_size = _len[idx];
    d.leaves_off = (uint32_t)(m.leaves - _buf[idx]);
    d.block_count = m.block_count;
    d.payload_off = (uint32_t)(m.payload - _buf[idx]);
    d.payload_size = m.payload_size;
    return true;
  }
  bool read(uint8_t idx, uint32_t off, uint8_t* out, uint32_t len) override {
    if (idx >= _n || (uint64_t)off + len > _len[idx]) return false;
    std::memcpy(out, _buf[idx] + off, len);
    return true;
  }
private:
  const uint8_t* _buf[8] = {nullptr}; uint32_t _len[8] = {0}; uint8_t _n = 0;
};
}

TEST(OtaTransfer, TwoManagersFullTransfer) {
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};

  server.begin(/*server's own target irrelevant for serving*/ 0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);   // tests exercise fetch-on-advert; policy default is OFF

  ASSERT_TRUE(server.serve(SIM_MOTA, SIM_MOTA_LEN));
  server.announce();   // -> client hears the beacon, queries, catalogs, then fetches

  pump(client);        // beacon -> query -> have -> startFetch -> full transfer

  EXPECT_EQ(client.fetchState(), OtaManager::COMPLETE);
  EXPECT_EQ(client.blocksHave(), client.blocksTotal());
  EXPECT_GT(client.blocksTotal(), 1u);

  // the client's reassembled container must be byte-identical to the original .mota...
  ASSERT_EQ(store.staged_size(), SIM_MOTA_LEN);
  EXPECT_EQ(0, std::memcmp(store.data(), SIM_MOTA, SIM_MOTA_LEN));

  // ...and independently re-verify it parses with a matching root + image_hash
  MotaManifest m;
  ASSERT_TRUE(mota_parse(store.data(), store.staged_size(), m));
  EXPECT_TRUE(mota_check_root(m));
  EXPECT_TRUE(mota_check_image_hash_full(m));
}

static void deliver_manifest_fragment(OtaManager& client, const uint8_t mid[4], uint8_t frag_idx,
                                      const uint8_t* bytes, uint16_t len) {
  uint8_t wire[MAX_PACKET_PAYLOAD];
  ManifestMsg msg;
  memcpy(msg.manifest_id, mid, 4);
  msg.frag_idx = frag_idx;
  msg.frag_total = (uint8_t)((MOTA_MFL + OTA_MF_FRAG - 1) / OTA_MF_FRAG);
  msg.bytes = bytes;
  msg.len = len;
  uint16_t wire_len = encode_manifest(wire, sizeof(wire), msg);
  ASSERT_GT(wire_len, 0);
  client.on_message(wire, wire_len);
}

TEST(OtaTransfer, RejectsShortNonFinalManifestFragment) {
  g_q.clear();
  MotaManifest manifest;
  ASSERT_TRUE(mota_parse(SIM_MOTA, SIM_MOTA_LEN, manifest));

  OtaManager client;
  OtaStoreRam<4096> store;
  SendTo to_none{&client};
  client.begin(SIM_TARGET_ID, sim_send, &to_none);
  client.set_fetch_store(&store);
  client.pull(manifest.merkle_root, manifest.target_id);
  g_q.clear();

  const uint8_t* bytes = manifest.manifest_start;
  const uint16_t final_len = (uint16_t)(MOTA_MFL - OTA_MF_FRAG);
  deliver_manifest_fragment(client, manifest.merkle_root, 1, bytes + OTA_MF_FRAG, final_len);
  deliver_manifest_fragment(client, manifest.merkle_root, 0, bytes, OTA_MF_FRAG - 1);
  EXPECT_EQ(client.fetchState(), OtaManager::WANT_MANIFEST);

  deliver_manifest_fragment(client, manifest.merkle_root, 0, bytes, OTA_MF_FRAG);
  EXPECT_EQ(client.fetchState(), OtaManager::FETCHING);
}

TEST(OtaTransfer, RejectsManifestBlockExponentBeforeShift) {
  g_q.clear();
  MotaManifest manifest;
  ASSERT_TRUE(mota_parse(SIM_MOTA, SIM_MOTA_LEN, manifest));
  std::array<uint8_t, MOTA_MFL> bytes;
  memcpy(bytes.data(), manifest.manifest_start, bytes.size());
  bytes[19] = 32;

  OtaManager client;
  OtaStoreRam<4096> store;
  SendTo to_none{&client};
  client.begin(SIM_TARGET_ID, sim_send, &to_none);
  client.set_fetch_store(&store);
  client.pull(manifest.merkle_root, manifest.target_id);
  g_q.clear();

  deliver_manifest_fragment(client, manifest.merkle_root, 0, bytes.data(), OTA_MF_FRAG);
  deliver_manifest_fragment(client, manifest.merkle_root, 1, bytes.data() + OTA_MF_FRAG,
                            (uint16_t)(MOTA_MFL - OTA_MF_FRAG));
  EXPECT_EQ(client.fetchState(), OtaManager::FAILED);
}

// Same end-to-end transfer, but with 1 KB logical blocks: each block is delivered as several
// self-describing DATA fragments (frag_off), reassembled by the client, then its merkle PROOF is
// requested + verified separately before the block is committed. Exercises the multi-fragment path.
TEST(OtaTransfer, MultiFragmentBlocks) {
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};

  server.begin(0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);

  ASSERT_TRUE(server.serve(SIM_MOTA_1K, SIM_MOTA_1K_LEN));
  server.announce();

  pump(client);

  EXPECT_EQ(client.fetchState(), OtaManager::COMPLETE);
  EXPECT_EQ(client.blocksTotal(), SIM_MOTA_1K_BLOCKS);   // 1 KB blocks => fewer, larger blocks
  EXPECT_EQ(client.blocksHave(), client.blocksTotal());

  ASSERT_EQ(store.staged_size(), SIM_MOTA_1K_LEN);
  EXPECT_EQ(0, std::memcmp(store.data(), SIM_MOTA_1K, SIM_MOTA_1K_LEN));
  MotaManifest m;
  ASSERT_TRUE(mota_parse(store.data(), store.staged_size(), m));
  EXPECT_TRUE(mota_check_root(m));
  EXPECT_TRUE(mota_check_image_hash_full(m));
}

// Multi-mota folder serve: a node serves its OWN fw (view0) PLUS an external folder (RamMotaSource) of
// other `.mota`. Peers discover BOTH via the tiny beacon -> query -> broadcast HAVE catalog, then fetch an
// external mota end-to-end. The relaying node never holds the folder image in RAM — it streams the
// manifest/leaves/blocks from the source on demand (loadSource + srcReadTramp + proof-gen from read
// leaves). The fetched bytes must equal the original `.mota` (proves the trustless relay is byte-exact).
TEST(OtaFolder, ServesSelfPlusFolderAndFetchesExternal) {
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};

  server.begin(/*own target irrelevant for serving*/ 0, sim_send, &to_client);
  uint8_t srv_id[4] = {0xAB, 0xCD, 0xEF, 0x01}; server.set_seeder_id(srv_id);
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);

  MotaManifest mSelf, mExt;
  ASSERT_TRUE(mota_parse(SIM_MOTA,    SIM_MOTA_LEN,    mSelf));   // served as our own fw (view0)
  ASSERT_TRUE(mota_parse(SIM_MOTA_1K, SIM_MOTA_1K_LEN, mExt));    // served from the external folder

  ASSERT_TRUE(server.serve(SIM_MOTA, SIM_MOTA_LEN));             // entry 0 = self
  static RamMotaSource folder;
  folder.add(SIM_MOTA_1K, SIM_MOTA_1K_LEN);                      // an external image (different mid)
  folder.add(SIM_MOTA, SIM_MOTA_LEN);                            // same as self -> must be DEDUPED
  ASSERT_TRUE(server.add_source(&folder));
  EXPECT_EQ(server.servedCount(), 2);                           // self + 1 distinct folder mota (dedup)

  // discovery: beacon -> the client catalogs the source, queries it, and the broadcast HAVE fills the
  // catalog with BOTH served mids.
  server.announce();
  pump(client);
  client.queryAll();
  pump(client);
  EXPECT_EQ(client.catalogCount(), 2);

  // fetch the EXTERNAL (folder) mota by mid -> served via the source, relayed block-by-block.
  client.pull(mExt.merkle_root, mExt.target_id);
  pump(client);

  EXPECT_EQ(client.fetchState(), OtaManager::COMPLETE);
  EXPECT_EQ(client.blocksHave(), client.blocksTotal());
  ASSERT_EQ(store.staged_size(), SIM_MOTA_1K_LEN);
  EXPECT_EQ(0, std::memcmp(store.data(), SIM_MOTA_1K, SIM_MOTA_1K_LEN));   // byte-exact relay
  MotaManifest got;
  ASSERT_TRUE(mota_parse(store.data(), store.staged_size(), got));
  EXPECT_TRUE(mota_check_root(got));
  EXPECT_TRUE(mota_check_image_hash_full(got));
}

// Fetch-resume across a reboot: a client commits some blocks, "reboots" (a fresh OtaManager on the SAME
// persisted store), and resumeStaged() re-adopts the partial container and finishes the remaining blocks —
// without re-fetching the manifest or the blocks already present.
TEST(OtaTransfer, ResumeAfterReboot) {
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};

  server.begin(0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);

  ASSERT_TRUE(server.serve(SIM_MOTA_1K, SIM_MOTA_1K_LEN));
  server.announce();

  // drive only until the first block commits, then "crash"
  int idle = 0, guard = 0;
  while (guard++ < 100000) {
    if (!g_q.empty()) {
      SimMsg msg = std::move(g_q.front()); g_q.erase(g_q.begin());
      msg.dest->on_message(msg.bytes.data(), (uint16_t)msg.bytes.size());
      idle = 0;
    } else {
      g_clk += 5000; client.set_clock(g_clk); client.loop();
      if (!g_q.empty()) { idle = 0; } else if (++idle >= 2) break;
    }
    if (client.blocksHave() >= 1) break;
  }
  ASSERT_GE(client.blocksHave(), 1u);
  ASSERT_LT(client.blocksHave(), client.blocksTotal());     // genuinely partial
  uint32_t had = client.blocksHave();
  g_q.clear();                                              // in-flight packets are lost in the "reboot"

  // "reboot": a brand-new manager on the SAME store (its bytes survived) resumes the partial
  OtaManager client2;
  to_client.dest = &client2;                                // server now replies to the rebooted client
  SendTo to_server2{&server};
  client2.begin(SIM_TARGET_ID, sim_send, &to_server2);
  client2.set_fetch_store(&store);
  ASSERT_TRUE(client2.resumeStaged(nullptr));               // adopt whatever is staged
  EXPECT_EQ(client2.blocksHave(), had);                     // resumed exactly where we left off
  EXPECT_EQ(client2.fetchState(), OtaManager::FETCHING);
  EXPECT_EQ(client2.blocksTotal(), SIM_MOTA_1K_BLOCKS);

  pump(client2);
  EXPECT_EQ(client2.fetchState(), OtaManager::COMPLETE);
  ASSERT_EQ(store.staged_size(), SIM_MOTA_1K_LEN);
  EXPECT_EQ(0, std::memcmp(store.data(), SIM_MOTA_1K, SIM_MOTA_1K_LEN));   // byte-identical to the original
}

TEST(OtaTransfer, ClientRejectsWrongTarget) {
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};
  server.begin(0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID ^ 0x1u, sim_send, &to_server);   // different target -> not interested
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);   // tests exercise fetch-on-advert; policy default is OFF
  ASSERT_TRUE(server.serve(SIM_MOTA, SIM_MOTA_LEN));
  server.announce();
  pump(client);   // catalogs the row but wantRow rejects it (wrong target) -> never fetches
  EXPECT_EQ(client.fetchState(), OtaManager::IDLE);   // never started
}

TEST(OtaTransfer, ManualCrossTargetFetch) {
  // A node whose own target differs from the served firmware normally won't fetch (role-switch case:
  // e.g. companion wanting repeater firmware). An explicit want() override lets it fetch deliberately.
  g_q.clear();
  OtaManager server, client;
  OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};
  server.begin(0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID ^ 0xABCDu, sim_send, &to_server);   // DIFFERENT own target
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);   // tests exercise fetch-on-advert; policy default is OFF
  ASSERT_TRUE(server.serve(SIM_MOTA, SIM_MOTA_LEN));

  // without the override: catalogs the row but won't fetch (wrong target)
  server.announce();
  pump(client);
  EXPECT_EQ(client.fetchState(), OtaManager::IDLE);

  // with want(): deliberately fetch the different-target firmware to completion
  client.want(SIM_TARGET_ID);
  server.announce();
  pump(client);
  EXPECT_EQ(client.fetchState(), OtaManager::COMPLETE);
  ASSERT_EQ(store.staged_size(), SIM_MOTA_LEN);
  EXPECT_EQ(0, std::memcmp(store.data(), SIM_MOTA, SIM_MOTA_LEN));
}

// Encode a 1-row OTA_HAVE catalog (the discovery reply a peer acts on).
static uint16_t make_have1(uint8_t* buf, uint16_t cap, const uint8_t mid[4],
                           uint32_t target, uint32_t fwver, uint8_t codec, uint8_t flags) {
  uint8_t row[OTA_HAVE_ROW_BYTES];
  memcpy(row, mid, 4);
  row[4]=target; row[5]=target>>8; row[6]=target>>16; row[7]=target>>24;
  row[8]=fwver; row[9]=fwver>>8; row[10]=fwver>>16; row[11]=fwver>>24;
  row[12]=codec; row[13]=flags;
  row[14]=0; row[15]=0;   // have_count (unused in this 1-row discovery test)
  HaveMsg hv{{0xAA,0xBB,0xCC,0xDD}, {0,0,0,0}, 0, 1, 1, row};
  return encode_have(buf, cap, hv);
}

// A node must not fetch firmware it can't apply: a catalog row whose codec the platform can't decode is
// not fetched. FULL + the platform's delta codec(s) are accepted.
TEST(OtaTransfer, RejectsIncompatibleCodec) {
  g_q.clear();
  OtaManager client; OtaStoreRam<4096> store;
  SendTo to_server{&client};                       // dest unused (we only check client state)
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);
  client.set_apply_codec(CODEC_DETOOLS_INPLACE);   // nRF52-style: accepts only full + in-place
  uint8_t b[64];

  // a SEQUENTIAL delta for our target -> incompatible -> not fetched (stays IDLE)
  uint8_t midA[4] = {1,2,3,4};
  client.on_message(b, make_have1(b, sizeof(b), midA, SIM_TARGET_ID, 0x01000000, CODEC_DETOOLS_SEQUENTIAL, 0));
  EXPECT_EQ(client.fetchState(), OtaManager::IDLE);

  // an IN-PLACE delta for our target -> compatible -> begins fetching (requests the manifest)
  uint8_t midB[4] = {5,6,7,8};
  client.on_message(b, make_have1(b, sizeof(b), midB, SIM_TARGET_ID, 0x01000000, CODEC_DETOOLS_INPLACE, 0));
  EXPECT_EQ(client.fetchState(), OtaManager::WANT_MANIFEST);
  g_q.clear();
}

// Encode a 1-row OTA_HAVE from a specific seeder, carrying have_count (Phase-2 awareness).
static uint16_t make_have_row(uint8_t* buf, uint16_t cap, const uint8_t mid[4], uint32_t target,
                              uint32_t fwver, uint8_t codec, uint8_t flags,
                              const uint8_t seeder[4], uint16_t have_count) {
  uint8_t row[OTA_HAVE_ROW_BYTES];
  memcpy(row, mid, 4);
  row[4]=target; row[5]=target>>8; row[6]=target>>16; row[7]=target>>24;
  row[8]=fwver; row[9]=fwver>>8; row[10]=fwver>>16; row[11]=fwver>>24;
  row[12]=codec; row[13]=flags;
  row[14]=(uint8_t)(have_count & 0xFF); row[15]=(uint8_t)(have_count >> 8);
  HaveMsg hv; memcpy(hv.seeder_id, seeder, 4); memset(hv.set_digest, 0, 4);
  hv.frag_idx=0; hv.frag_total=1; hv.n_rows=1; hv.rows=row;
  return encode_have(buf, cap, hv);
}

// Catalog accounting: "N nodes have it" must count DISTINCT seeders (a repeated HAVE from one node must
// not inflate it), and have_max tracks the best progress any source reported.
TEST(OtaCatalog, DistinctSeederCountAndHaveCount) {
  OtaManager m; SendTo none{&m}; m.begin(SIM_TARGET_ID, sim_send, &none);
  uint8_t b[64]; uint8_t mid[4]={9,9,9,9};
  uint8_t s1[4]={1,0,0,0}, s2[4]={2,0,0,0};
  m.on_message(b, make_have_row(b, sizeof b, mid, SIM_TARGET_ID, 0x01020300, CODEC_FULL, 0, s1, 5));
  m.on_message(b, make_have_row(b, sizeof b, mid, SIM_TARGET_ID, 0x01020300, CODEC_FULL, 0, s1, 7));  // same seeder
  ASSERT_EQ(m.catalogCount(), 1);
  EXPECT_EQ(m.catalogRow(0)->n_seeders, 1);          // counted once despite two HAVEs
  EXPECT_EQ(m.catalogRow(0)->have_max, 7u);          // max progress seen
  m.on_message(b, make_have_row(b, sizeof b, mid, SIM_TARGET_ID, 0x01020300, CODEC_FULL, 0, s2, 3));  // new seeder
  EXPECT_EQ(m.catalogRow(0)->n_seeders, 2);
  EXPECT_EQ(m.catalogRow(0)->have_max, 7u);          // still the max, not overwritten by the lower one
  g_q.clear();
}

// An unanswered GET_MANIFEST must not pin the fetch slot forever: after OTA_MANIFEST_MAX_RETRY ticks with
// no manifest, the session gives up (FAILED) so a new pull can take the slot. (Lowest-priority + bounded.)
TEST(OtaTransfer, ManifestGiveUpAfterRetries) {
  g_q.clear();
  OtaManager client; OtaStoreRam<4096> store; SendTo to_none{&client};
  client.begin(SIM_TARGET_ID, sim_send, &to_none);
  client.set_fetch_store(&store);
  uint8_t mid[4]={7,7,7,7};
  client.pull(mid, SIM_TARGET_ID);                   // no server -> stuck WANT_MANIFEST
  EXPECT_EQ(client.fetchState(), OtaManager::WANT_MANIFEST);
  for (int i = 0; i < OTA_MANIFEST_MAX_RETRY + 2; i++) { g_clk += 5000; client.set_clock(g_clk); client.loop(); g_q.clear(); }
  EXPECT_EQ(client.fetchState(), OtaManager::FAILED);
}

// A receiver never becomes a source, either while fetching or after completion.
TEST(OtaTransfer, ReceiverDoesNotReSeed) {
  g_q.clear();
  OtaManager server, client; OtaStoreRam<4096> store;
  SendTo to_client{&client}, to_server{&server};
  server.begin(0, sim_send, &to_client);
  client.begin(SIM_TARGET_ID, sim_send, &to_server);
  client.set_fetch_store(&store);
  client.set_autofetch(OtaManager::AUTOFETCH_ANY);
  ASSERT_TRUE(server.serve(SIM_MOTA, SIM_MOTA_LEN));
  server.announce();
  pump(client);
  ASSERT_EQ(client.fetchState(), OtaManager::COMPLETE);
  EXPECT_EQ(client.servedCount(), 0);
  client.reset_session();
  EXPECT_EQ(client.servedCount(), 0);
}

// --- detools delta decode (vendored detools C decoder, CRLE-only build) ----------------------
// Mirrors the device apply path (src/helpers/ota/OtaApply.cpp): base read via from_read/from_seek,
// patch streamed via patch_read, output written via to_write. Proves the on-device delta apply uses
// detools 0.53.0's own decoder and reproduces the exact target the host packager targeted.
namespace {
struct DTMem {
  const uint8_t* base; long base_len; long base_pos;
  const uint8_t* patch; long patch_len; long patch_pos;
  std::vector<uint8_t> out;
};
int dt_from_read(void* a, uint8_t* b, size_t n) {
  DTMem* c = (DTMem*)a;
  if (c->base_pos < 0 || c->base_pos + (long)n > c->base_len) return -DETOOLS_IO_FAILED;
  std::memcpy(b, c->base + c->base_pos, n); c->base_pos += (long)n; return DETOOLS_OK;
}
int dt_from_seek(void* a, int off) {
  DTMem* c = (DTMem*)a; c->base_pos += off;
  if (c->base_pos < 0 || c->base_pos > c->base_len) return -DETOOLS_IO_FAILED;
  return DETOOLS_OK;
}
int dt_patch_read(void* a, uint8_t* b, size_t n) {
  DTMem* c = (DTMem*)a;
  if (c->patch_pos + (long)n > c->patch_len) return -DETOOLS_IO_FAILED;
  std::memcpy(b, c->patch + c->patch_pos, n); c->patch_pos += (long)n; return DETOOLS_OK;
}
int dt_to_write(void* a, const uint8_t* b, size_t n) {
  DTMem* c = (DTMem*)a; c->out.insert(c->out.end(), b, b + n); return DETOOLS_OK;
}

// In-place apply over a flat memory region (models the nRF52 app workspace / the bootloader's flash).
struct DTInPlace {
  std::vector<uint8_t> mem;             // [0,memory_size): base in, target out
  const uint8_t* patch; long plen, ppos; int step;
};
int ip_mem_read(void* a, void* dst, uintptr_t src, size_t n) {
  DTInPlace* c = (DTInPlace*)a; if (src + n > c->mem.size()) return -DETOOLS_IO_FAILED;
  std::memcpy(dst, c->mem.data() + src, n); return DETOOLS_OK;
}
int ip_mem_write(void* a, uintptr_t dst, void* src, size_t n) {
  DTInPlace* c = (DTInPlace*)a; if (dst + n > c->mem.size()) return -DETOOLS_IO_FAILED;
  std::memcpy(c->mem.data() + dst, src, n); return DETOOLS_OK;
}
int ip_mem_erase(void* a, uintptr_t addr, size_t n) {
  DTInPlace* c = (DTInPlace*)a; if (addr + n > c->mem.size()) return -DETOOLS_IO_FAILED;
  std::memset(c->mem.data() + addr, 0xFF, n); return DETOOLS_OK;
}
int ip_step_set(void* a, int s) { ((DTInPlace*)a)->step = s; return DETOOLS_OK; }
int ip_step_get(void* a, int* s) { *s = ((DTInPlace*)a)->step; return DETOOLS_OK; }
int ip_patch_read(void* a, uint8_t* b, size_t n) {
  DTInPlace* c = (DTInPlace*)a; if (c->ppos + (long)n > c->plen) return -DETOOLS_IO_FAILED;
  std::memcpy(b, c->patch + c->ppos, n); c->ppos += (long)n; return DETOOLS_OK;
}
}  // namespace

TEST(Detools, SequentialCrlePatchReproducesTarget) {
  DTMem c{DT_BASE, (long)DT_BASE_LEN, 0, DT_PATCH, (long)DT_PATCH_LEN, 0, {}};
  int r = detools_apply_patch_callbacks(dt_from_read, dt_from_seek, dt_patch_read,
                                        (size_t)DT_PATCH_LEN, dt_to_write, &c);
  ASSERT_EQ(r, (int)DT_TARGET_LEN);                 // returns to-size on success
  ASSERT_EQ(c.out.size(), (size_t)DT_TARGET_LEN);
  EXPECT_EQ(0, std::memcmp(c.out.data(), DT_TARGET, DT_TARGET_LEN));
}

TEST(Detools, WrongBaseDoesNotReproduceTarget) {
  // a base that differs from the one the patch was built against must NOT yield the target
  std::vector<uint8_t> bad(DT_BASE, DT_BASE + DT_BASE_LEN);
  for (size_t i = 0; i < bad.size(); i += 7) bad[i] ^= 0xFF;
  DTMem c{bad.data(), (long)bad.size(), 0, DT_PATCH, (long)DT_PATCH_LEN, 0, {}};
  int r = detools_apply_patch_callbacks(dt_from_read, dt_from_seek, dt_patch_read,
                                        (size_t)DT_PATCH_LEN, dt_to_write, &c);
  bool reproduced = (r == (int)DT_TARGET_LEN && c.out.size() == (size_t)DT_TARGET_LEN &&
                     std::memcmp(c.out.data(), DT_TARGET, DT_TARGET_LEN) == 0);
  EXPECT_FALSE(reproduced);   // wrong base -> wrong/short output (the device then fails image_hash)
}

TEST(Detools, TruncatedPatchFails) {
  DTMem c{DT_BASE, (long)DT_BASE_LEN, 0, DT_PATCH, (long)(DT_PATCH_LEN / 2), 0, {}};
  int r = detools_apply_patch_callbacks(dt_from_read, dt_from_seek, dt_patch_read,
                                        (size_t)(DT_PATCH_LEN / 2), dt_to_write, &c);
  EXPECT_TRUE(r < 0 || c.out.size() != (size_t)DT_TARGET_LEN);
}

// nRF52 path: the bootloader applies an in-place patch over the single app slot. Model the app
// region as a DT_IP_MEM buffer holding the base; after apply, region[0:to_size] must equal the target.
TEST(Detools, InPlaceCrlePatchReproducesTarget) {
  DTInPlace c; c.mem.assign(DT_IP_MEM, 0xFF);
  std::memcpy(c.mem.data(), DT_IP_BASE, DT_IP_BASE_LEN);   // base loaded at offset 0
  c.patch = DT_IP_PATCH; c.plen = DT_IP_PATCH_LEN; c.ppos = 0; c.step = 0;
  int r = detools_apply_patch_in_place_callbacks(ip_mem_read, ip_mem_write, ip_mem_erase,
                                                 ip_step_set, ip_step_get, ip_patch_read,
                                                 (size_t)DT_IP_PATCH_LEN, &c);
  ASSERT_EQ(r, (int)DT_IP_TARGET_LEN);                     // returns to-size on success
  EXPECT_EQ(0, std::memcmp(c.mem.data(), DT_IP_TARGET, DT_IP_TARGET_LEN));
}

// --- leaf-diff warm-start core (motatool folder-capture): the device fetches the target leaves[], recomputes
// the root to authenticate them, then keeps every seed block whose leaf matches and refetches only the rest.
// This exercises that logic (leaf authentication + per-block diff) with no store/fetch machinery. ----------
TEST(OtaWarmStart, LeafDiffAuthenticatesAndFindsDifferingBlocks) {
  const uint32_t BS = 16, BC = 5;
  std::vector<uint8_t> target(BS * BC), seed(BS * BC);
  for (uint32_t i = 0; i < target.size(); i++) target[i] = seed[i] = (uint8_t)(i * 7 + 3);
  seed[1 * BS + 5] ^= 0xFF;      // blocks 1 and 3 differ in the seed (a non-deterministic-rebuild style diff)
  seed[3 * BS + 0] ^= 0x01;

  // target leaves + root (what the device receives over OTA_LEAVES + the manifest merkle_root)
  uint8_t tleaves[BC * 4], troot[4];
  for (uint32_t i = 0; i < BC; i++) merkle_leaf(tleaves + i * 4, target.data() + i * BS, BS);
  merkle_root(troot, tleaves, BC);

  // authenticate the fetched leaves: recomputing the root from them must equal the manifest root
  uint8_t chk[4]; merkle_root(chk, tleaves, BC);
  EXPECT_EQ(0, memcmp(chk, troot, 4));

  // diff: a seed block is kept iff its leaf equals the (authenticated) target leaf
  int nmiss = 0; bool miss[BC] = {false};
  for (uint32_t i = 0; i < BC; i++) {
    uint8_t sl[4]; merkle_leaf(sl, seed.data() + i * BS, BS);
    if (memcmp(sl, tleaves + i * 4, 4) != 0) { miss[i] = true; nmiss++; }
  }
  EXPECT_EQ(nmiss, 2);
  EXPECT_TRUE(miss[1]);  EXPECT_TRUE(miss[3]);
  EXPECT_FALSE(miss[0]); EXPECT_FALSE(miss[2]); EXPECT_FALSE(miss[4]);
}

int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}