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

#include "helpers/ota/OtaFlashLayout_nrf52.h"
#include "helpers/ota/OtaFlashLayout_esp32.h"
#include "helpers/ota/OtaStoreQspiNrf52.h"

using namespace mesh::ota;

// These lock down the nRF52 single-slot staging geometry that OtaStoreFlashNrf52::begin()/
// reopen() rely on. A received `.mota` is placed bottom-aligned below the filesystem region - ExtraFS
// (0xD4000) / InternalFS (0xED000), where the node's user preferences live - and above the running image.
// The prefs region is assumed IMMUTABLE (its bytes are outside the served/hashed self-image), so staging
// or an in-place apply must never reach into it. If a layout constant or the placement math is edited
// inconsistently, these fail here instead of silently corrupting prefs / the app on real hardware.

static constexpr uint32_t APP_V6 = MOTA_NRF52_APP_BASE_S140_V6;
static constexpr uint32_t APP_V7 = MOTA_NRF52_APP_BASE_S140_V7;
static constexpr uint32_t LEGACY = MOTA_NRF52_STAGE_CEILING_LEGACY;
static constexpr uint32_t EXPANDED = MOTA_NRF52_STAGE_CEILING_EXPANDED;
static constexpr uint32_t SHARED_BOOT_START = MOTA_NRF52_SHARED_BOOT_STAGE_START;
static constexpr uint32_t APP_IMAGE_SIZE = 520u * 1024u;
static constexpr uint32_t APP_END_V6 = APP_V6 + APP_IMAGE_SIZE;
static constexpr uint32_t APP_END_V7 = APP_V7 + APP_IMAGE_SIZE;
static constexpr uint32_t CAP_V6 = LEGACY - APP_END_V6;
static constexpr uint32_t CAP_V7 = LEGACY - APP_END_V7;
static constexpr uint32_t CAP_V6_EXPANDED = EXPANDED - APP_END_V6;
static constexpr uint32_t CAP_V7_EXPANDED = EXPANDED - APP_END_V7;

TEST(OtaQspiTiming, PreservesDeepPowerDownEntryAndWakeGuards) {
  // MX25R1635F requires 10 us to enter DPD plus 30 us before another command;
  // its release latency can reach 45 us. These constants are consumed by the
  // real HAL path, so a future power-saving edit cannot restore the live
  // plan_layout()->begin() race without failing the native suite.
  EXPECT_GE(MOTA_QSPI_DPD_ENTRY_GUARD_US, 50u);
  EXPECT_GE(MOTA_QSPI_DPD_WAKE_GUARD_US, 45u);
}

TEST(OtaQspiWake, ShiftsReleaseCommandMostSignificantBitFirst) {
  // This byte is emitted over GPIO before TASKS_ACTIVATE. If activation is
  // attempted first, a sleeping NOR ignores it and READY never arrives.
  const bool expected[] = {true, false, true, false, true, false, true, true};
  EXPECT_EQ(MOTA_QSPI_RELEASE_FROM_DPD_OPCODE, 0xABu);
  for (uint8_t bit = 0; bit < 8u; bit++) {
    EXPECT_EQ(mota_qspi_release_from_dpd_bit(bit), expected[bit]);
  }
  EXPECT_FALSE(mota_qspi_release_from_dpd_bit(8u));
}

TEST(OtaQspiStatus, TreatsOnlyTheNorWriteInProgressBitAsBusy) {
  // Nordic READY is not flash completion. The hardware path polls RDSR until
  // this predicate clears after every program and erase operation.
  EXPECT_FALSE(mota_qspi_status_busy(0x00));
  EXPECT_FALSE(mota_qspi_status_busy(0xFC));
  EXPECT_TRUE(mota_qspi_status_busy(0x01));
  EXPECT_TRUE(mota_qspi_status_busy(0xFF));
}

TEST(OtaQspiDiagnostics, ExposesStableFailureStageNames) {
  EXPECT_STREQ(mota_qspi_stage_name(OtaQspiStage::PROGRAM), "program");
  EXPECT_STREQ(mota_qspi_stage_name(OtaQspiStage::PROGRAM_BUSY), "program-busy");
  EXPECT_STREQ(mota_qspi_stage_name(OtaQspiStage::ERASE_BUSY), "erase-busy");
  EXPECT_STREQ(mota_qspi_stage_name(OtaQspiStage::INVALIDATE_VERIFY), "invalidate-verify");
}

TEST(OtaFlashPlan, SelectsCeilingFromLinkedLayoutAndStorage) {
  // Actual internal secondary storage is authoritative regardless of linker selection.
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(EXPANDED, true), LEGACY);
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(LEGACY, true), LEGACY);
  // Standard or ExtraFS linker without an internal secondary filesystem reclaims the unused 100 KiB.
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(EXPANDED, false), EXPANDED);
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(LEGACY, false), EXPANDED);
  // The explicit XIAO boot-update linker leaves 40 KiB scratch at 0xE0000 and still stages externally.
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(MOTA_NRF52_BOOT_SCRATCH_START, false), EXPANDED);
  // An unrecognized linker region is never permission to erase a larger window.
  EXPECT_EQ(mota_nrf52_stage_ceiling_for_layout(0xE1000u, false), LEGACY);
}

TEST(OtaFlashPlan, ValidatesExternalInplacePatchGeometryBeforeHandoff) {
  const uint32_t workspace = EXPANDED - APP_V7;
  const uint32_t running = 620000u;
  const uint32_t target = 633984u;
  ASSERT_TRUE(mota_nrf52_external_patch_geometry_valid(
      workspace, MOTA_NRF52_FLASH_PAGE, 0, running, target,
      workspace, running, target));

  EXPECT_FALSE(mota_nrf52_external_patch_geometry_valid(
      workspace + 1, MOTA_NRF52_FLASH_PAGE, 0, running, target,
      workspace, running, target));
  EXPECT_FALSE(mota_nrf52_external_patch_geometry_valid(
      workspace, 2048, 0, running, target, workspace, running, target));
  EXPECT_FALSE(mota_nrf52_external_patch_geometry_valid(
      workspace, MOTA_NRF52_FLASH_PAGE, 1, running, target,
      workspace, running, target));
  EXPECT_FALSE(mota_nrf52_external_patch_geometry_valid(
      workspace, MOTA_NRF52_FLASH_PAGE, 0, running - 1, target,
      workspace, running, target));
  EXPECT_FALSE(mota_nrf52_external_patch_geometry_valid(
      workspace, MOTA_NRF52_FLASH_PAGE, 0, running, target + 1,
      workspace, running, target));
}

// A typical running image (~520 KB) leaves room; the container lands strictly above it and below ExtraFS.
TEST(OtaFlashPlan, StagesBelowFilesystemAndAboveApp) {
  uint32_t start = 0xDEADBEEF;
  ASSERT_TRUE(mota_nrf52_stage_plan(64u * 1024u, APP_V6, APP_END_V6, LEGACY, start));
  EXPECT_GE(start, APP_END_V6);                         // never overlaps the running image
  EXPECT_LE(start + 64u * 1024u, LEGACY);              // never reaches into ExtraFS/prefs
  EXPECT_EQ(start % MOTA_NRF52_FLASH_PAGE, 0u);         // page-aligned (the flash erase unit)
}

// Bottom-aligned: start is the page-aligned FS_START - total_size, so the trailer sits within the
// highest page below the ceiling where the bootloader's downward scan finds it.
TEST(OtaFlashPlan, BottomAlignedBelowCeiling) {
  uint32_t start = 0;
  uint32_t total = 60000;
  ASSERT_TRUE(mota_nrf52_stage_plan(total, APP_V6, APP_V6, LEGACY, start));
  EXPECT_EQ(start, (LEGACY - total) & ~(MOTA_NRF52_FLASH_PAGE - 1));
  EXPECT_LE(start + total, LEGACY);
  EXPECT_GT(start + total, LEGACY - MOTA_NRF52_FLASH_PAGE);   // within one page of the ceiling
}

// An exactly-capacity container fills the page-aligned space above the app; one byte more never fits.
TEST(OtaFlashPlan, RejectsOversizedContainer) {
  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(CAP_V6, APP_V6, APP_END_V6, LEGACY, start));
  EXPECT_EQ(start, APP_END_V6);
  EXPECT_EQ(start + CAP_V6, LEGACY);
  EXPECT_FALSE(mota_nrf52_stage_plan(CAP_V6 + 1, APP_V6, APP_END_V6, LEGACY, start));
}

// The package carries its own checked memory_size, so a valid app may exceed the old 608 KiB fallback.
TEST(OtaFlashPlan, AcceptsAppLargerThanFallbackWhenPackageFits) {
  const uint32_t app_end = APP_V6 + MOTA_NRF52_FALLBACK_INPLACE_MEMORY + 32u * 1024u;
  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(16u * 1024u, APP_V6, app_end, LEGACY, start));
  EXPECT_GE(start, app_end);
}

// Minimum container is header(8)+trailer(5)=13 bytes; anything smaller is not a container.
TEST(OtaFlashPlan, RejectsUndersizedContainer) {
  uint32_t start = 0;
  EXPECT_FALSE(mota_nrf52_stage_plan(12, APP_V6, APP_V6, LEGACY, start));
  EXPECT_TRUE(mota_nrf52_stage_plan(13, APP_V6, APP_V6, LEGACY, start));
}

// The user-preferences filesystems (ExtraFS @ 0xD4000, InternalFS @ 0xED000) are entirely ABOVE any
// staged container. Hard-code the FS addresses here (independent of the layout header) so a future edit
// that drifts into a filesystem is caught.
TEST(OtaFlashPlan, PrefsRegionNeverStaged) {
  const uint32_t EXTRAFS_START    = 0xD4000u;   // companion ExtraFS (CustomLFS(0xD4000, 0x19000))
  const uint32_t INTERNALFS_START = 0xED000u;   // primary LittleFS (holds /com_prefs)
  EXPECT_EQ(LEGACY, EXTRAFS_START);
  EXPECT_EQ(EXPANDED, INTERNALFS_START);
  EXPECT_LT(EXTRAFS_START, INTERNALFS_START);
  // the largest possible staged container still ends at the ceiling, never into a filesystem
  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(CAP_V6, APP_V6, APP_END_V6, LEGACY, start));
  EXPECT_LE(start + CAP_V6, EXTRAFS_START);
}

// S140 v7 moves the app start by one page. Runtime linker-base discovery must leave a correspondingly
// smaller but still safe staging region rather than scanning the v6 address and missing EndF.
TEST(OtaFlashPlan, SupportsS140V7RuntimeBase) {
  EXPECT_TRUE(mota_nrf52_layout_valid(APP_V7, LEGACY));
  EXPECT_EQ(mota_nrf52_stage_capacity(APP_V7, APP_END_V7, LEGACY), CAP_V7);
  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(CAP_V7, APP_V7, APP_END_V7, LEGACY, start));
  EXPECT_EQ(start, APP_END_V7);
  EXPECT_EQ(start + CAP_V7, LEGACY);
}

TEST(OtaFlashPlan, ExpandedCeilingAddsExactly100KiBForV6AndV7) {
  EXPECT_EQ(EXPANDED - LEGACY, 100u * 1024u);
  EXPECT_EQ(CAP_V6_EXPANDED - CAP_V6, 100u * 1024u);
  EXPECT_EQ(CAP_V7_EXPANDED - CAP_V7, 100u * 1024u);
  EXPECT_EQ(mota_nrf52_stage_capacity(APP_V6, APP_END_V6, EXPANDED), CAP_V6_EXPANDED);
  EXPECT_EQ(mota_nrf52_stage_capacity(APP_V7, APP_END_V7, EXPANDED), CAP_V7_EXPANDED);

  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(CAP_V7_EXPANDED, APP_V7, APP_END_V7, EXPANDED, start));
  EXPECT_EQ(start, APP_END_V7);
  EXPECT_EQ(start + CAP_V7_EXPANDED, EXPANDED);
}

TEST(OtaFlashPlan, HybridUsesTheFrozenFlashChargeAndRamSuffix) {
  uint32_t start = 0, flash = 0, ram = 0;
  const uint32_t page = MOTA_NRF52_FLASH_PAGE;
  ASSERT_TRUE(mota_nrf52_hybrid_stage_plan(
      page + 1u, APP_V6, APP_END_V6, EXPANDED, start, flash, ram));
  EXPECT_EQ(start, EXPANDED - page);
  EXPECT_EQ(flash, page);
  EXPECT_EQ(ram, 1u);

  ASSERT_TRUE(mota_nrf52_hybrid_stage_plan(
      17u * page, APP_V6, APP_END_V6, EXPANDED, start, flash, ram));
  EXPECT_EQ(start, EXPANDED - page);
  EXPECT_EQ(flash, page);
  EXPECT_EQ(ram, MOTA_NRF52_HYBRID_RAM_SIZE);

  ASSERT_TRUE(mota_nrf52_hybrid_stage_plan(
      17u * page + 1u, APP_V6, APP_END_V6,
      EXPANDED, start, flash, ram));
  EXPECT_EQ(start, EXPANDED - 2u * page);
  EXPECT_EQ(flash, 2u * page);
  EXPECT_EQ(ram, 15u * page + 1u);

  ASSERT_TRUE(mota_nrf52_hybrid_stage_plan(
      17u * page, APP_V7, APP_END_V7,
      EXPANDED, start, flash, ram));
  EXPECT_EQ(start, EXPANDED - page);
  EXPECT_EQ(flash, page);
  EXPECT_EQ(ram, MOTA_NRF52_HYBRID_RAM_SIZE);
}

TEST(OtaFlashPlan, HybridFailsClosedOutsideItsExactProfile) {
  uint32_t start = 0xAAAAAAAAu, flash = 0xBBBBBBBBu, ram = 0xCCCCCCCCu;
  const uint32_t page = MOTA_NRF52_FLASH_PAGE;
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      MOTA_NRF52_CONTAINER_MIN_SIZE, APP_V6, APP_END_V6,
      EXPANDED, start, flash, ram));
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      page, APP_V6, APP_END_V6, EXPANDED, start, flash, ram));
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      17u * page, APP_V6, APP_END_V6, LEGACY, start, flash, ram));
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      17u * page + 1u, APP_V6, EXPANDED - page,
      EXPANDED, start, flash, ram));
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      17u * page, 0x00028000u, 0x00028000u,
      EXPANDED, start, flash, ram));
  EXPECT_FALSE(mota_nrf52_hybrid_stage_plan(
      UINT32_MAX, APP_V6, APP_END_V6,
      EXPANDED, start, flash, ram));
  EXPECT_EQ(start, 0xAAAAAAAAu);
  EXPECT_EQ(flash, 0xBBBBBBBBu);
  EXPECT_EQ(ram, 0xCCCCCCCCu);
}

TEST(OtaFlashPlan, RejectsAppOutsideSelectedRegion) {
  uint32_t start = 0x1234ABCD;
  EXPECT_FALSE(mota_nrf52_stage_plan(4096, APP_V6, LEGACY + 1, LEGACY, start));
  EXPECT_EQ(start, 0x1234ABCDu);
}

TEST(OtaFlashPlan, RejectsUnknownCeiling) {
  uint32_t start = 0x1234ABCD;
  EXPECT_FALSE(mota_nrf52_stage_plan(4096, APP_V6, APP_V6, 0xE1000u, start));
  EXPECT_EQ(start, 0x1234ABCDu);
}

TEST(OtaFlashPlan, BootPackageUsesOneDynamicSharedInternalSlot) {
  // v3 = header + fixed manifest + forty leaves + 40 KiB image + trailer.
  const uint32_t total = 8u + 197u + 40u * 4u + 40u * 1024u + 5u;
  ASSERT_EQ(total, MOTA_NRF52_BOOT_CONTAINER_SIZE);
  uint32_t start = 0;
  ASSERT_TRUE(mota_nrf52_shared_boot_stage_plan(
      total, APP_V6, true, SHARED_BOOT_START, start));
  EXPECT_EQ(start, SHARED_BOOT_START);
  EXPECT_EQ(start, (EXPANDED - total) & ~(MOTA_NRF52_FLASH_PAGE - 1u));

  // OTAFIX compacts payload offset +365 forward inside these same eleven
  // pages; no independent scratch bank participates in the internal path.
  const uint32_t payload_offset = 8u + 197u + 40u * 4u;
  EXPECT_EQ(payload_offset, 365u);
  EXPECT_LE(start + payload_offset + 40u * 1024u, EXPANDED);
  EXPECT_LE(start + 40u * 1024u, EXPANDED);

  // Runtime headroom, not a special linker, is authoritative.
  EXPECT_FALSE(mota_nrf52_shared_boot_stage_plan(
      total, APP_V6, true, SHARED_BOOT_START + 1u, start));
  EXPECT_FALSE(mota_nrf52_shared_boot_stage_plan(
      total - 1u, APP_V6, true, APP_V6, start));
  EXPECT_FALSE(mota_nrf52_shared_boot_stage_plan(
      total + 1u, APP_V6, true, APP_V6, start));
}

TEST(OtaFlashPlan, SharedInternalStoreKeepsNormalApplicationCeilingAndHandoff) {
  EXPECT_TRUE(mota_nrf52_target_image_fits(
      APP_V6, EXPANDED - APP_V6, EXPANDED));
  EXPECT_FALSE(mota_nrf52_target_image_fits(
      APP_V6, EXPANDED - APP_V6 + 1u, EXPANDED));
  EXPECT_EQ(mota_nrf52_flash_stage_handoff(EXPANDED),
            GPREGRET2_OTA_STAGE_EXPANDED);
  EXPECT_EQ(mota_nrf52_flash_stage_handoff(LEGACY),
            GPREGRET2_OTA_STAGE_LEGACY);
}

TEST(OtaFlashPlan, InternalOrdinaryDeltaUsesTheSameEd000Store) {
  const uint32_t running_end = APP_V6 + 500u * 1024u;
  const uint32_t delta_container = 28u * 1024u;
  uint32_t staged = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(delta_container, APP_V6, running_end,
                                    EXPANDED, staged));
  EXPECT_GE(staged, running_end);
  EXPECT_LE(staged + delta_container, EXPANDED);
  EXPECT_EQ(mota_nrf52_flash_stage_handoff(EXPANDED), 0xEDu);
}

TEST(OtaFlashPlan, LargeInternalOrdinaryDeltaMayStageBelowSharedBootStart) {
  // An ordinary delta container can exceed the boot package's eleven pages.
  // It bottom-aligns below E2000 and remains valid when detools'
  // encoded workspace ends before the actual container start.
  const uint32_t running_end = APP_V6 + 600u * 1024u;
  const uint32_t delta_container = 80u * 1024u;
  uint32_t staged = 0;
  ASSERT_TRUE(mota_nrf52_stage_plan(delta_container, APP_V6, running_end,
                                    EXPANDED, staged));
  EXPECT_LT(staged, SHARED_BOOT_START);
  EXPECT_GE(staged, running_end);

  const uint32_t workspace = staged - APP_V6;
  const uint32_t target_size = workspace;
  EXPECT_TRUE(mota_nrf52_internal_patch_workspace_valid(
      workspace, APP_V6, staged, target_size, EXPANDED));
  EXPECT_FALSE(mota_nrf52_internal_patch_workspace_valid(
      workspace + 1u, APP_V6, staged, target_size, EXPANDED));
  EXPECT_FALSE(mota_nrf52_internal_patch_workspace_valid(
      workspace, APP_V6, staged, EXPANDED - APP_V6 + 1u, EXPANDED));
}

TEST(OtaFlashPlan, MissingEndfRejectsEverySharedInternalPackageBeforeErase) {
  uint32_t protected_end = 0;
  ASSERT_TRUE(mota_nrf52_protected_app_end(
      APP_V6, EXPANDED, false, 0, false, protected_end));
  EXPECT_EQ(protected_end, APP_V6 + MOTA_NRF52_FALLBACK_INPLACE_MEMORY);

  // A shared-internal-update build is linked through ED000, so the old 608 KiB
  // estimate is not a safe erase floor for either package kind. Missing EndF
  // disables all internal staging before begin()/reopen can erase a page.
  EXPECT_FALSE(mota_nrf52_protected_app_end(
      APP_V6, EXPANDED, false, 0, true, protected_end));

  // The old rescue estimate would also admit this ordinary 80 KiB delta at
  // D9000 even though a legal ED000-linked image can have live bytes there.
  // A privileged shared-internal build therefore never calls stage_plan with
  // that estimate: protected_app_end() above returns false/capacity zero.
  const uint32_t ordinary_container = 80u * 1024u;
  uint32_t unsafe_start = 0;
  const uint32_t legacy_estimate =
      APP_V6 + MOTA_NRF52_FALLBACK_INPLACE_MEMORY;
  ASSERT_TRUE(mota_nrf52_stage_plan(
      ordinary_container, APP_V6, legacy_estimate, EXPANDED, unsafe_start));
  const uint32_t legal_live_tail = unsafe_start + MOTA_NRF52_FLASH_PAGE;
  ASSERT_GT(legal_live_tail, legacy_estimate);
  ASSERT_LT(legal_live_tail, EXPANDED);
  EXPECT_FALSE(mota_nrf52_stage_plan(
      ordinary_container, APP_V6, legal_live_tail, EXPANDED, protected_end));

  // Model a legal linked image whose tail extends beyond the old 608 KiB
  // estimate. The privileged planner rejects missing EndF even if the generic
  // rescue fallback would appear to leave enough room.
  const uint32_t real_tail = APP_V6 + MOTA_NRF52_FALLBACK_INPLACE_MEMORY + 0x8000u;
  ASSERT_LT(real_tail, SHARED_BOOT_START);
  uint32_t staged = 0;
  EXPECT_FALSE(mota_nrf52_shared_boot_stage_plan(
      MOTA_NRF52_BOOT_CONTAINER_SIZE, APP_V6, false, protected_end, staged));
  ASSERT_TRUE(mota_nrf52_shared_boot_stage_plan(
      MOTA_NRF52_BOOT_CONTAINER_SIZE, APP_V6, true, real_tail, staged));
  EXPECT_EQ(staged, SHARED_BOOT_START);
}

TEST(OtaFlashPlan, ReopenBoundsUntrustedTotalBeforeManifestRead) {
  const uint32_t start = SHARED_BOOT_START;
  EXPECT_TRUE(mota_nrf52_container_span_valid(
      start, EXPANDED, MOTA_NRF52_BOOT_CONTAINER_SIZE, 8u + 197u + 5u));
  EXPECT_FALSE(mota_nrf52_container_span_valid(
      start, EXPANDED, UINT32_MAX, 8u + 197u + 5u));
  EXPECT_FALSE(mota_nrf52_container_span_valid(
      start, EXPANDED, 8u + 197u + 4u, 8u + 197u + 5u));
  EXPECT_FALSE(mota_nrf52_container_span_valid(
      EXPANDED + 1u, EXPANDED, MOTA_NRF52_BOOT_CONTAINER_SIZE,
      8u + 197u + 5u));
}

namespace {

struct FakeStagedHeader {
  uint32_t address;
  uint32_t total;
  bool present;
  bool invalidate_ok;
  bool invalidated;
};

struct FakeStagedFlash {
  FakeStagedHeader* headers;
  size_t count;
  uint32_t invalidate_calls;
};

bool fake_read_staged_header(void* context, uint32_t address,
                             uint32_t& total) {
  FakeStagedFlash* flash = static_cast<FakeStagedFlash*>(context);
  for (size_t i = 0; i < flash->count; ++i) {
    FakeStagedHeader& header = flash->headers[i];
    if (header.address == address && header.present) {
      total = header.total;
      return true;
    }
  }
  return false;
}

bool fake_invalidate_staged_header(void* context, uint32_t address) {
  FakeStagedFlash* flash = static_cast<FakeStagedFlash*>(context);
  ++flash->invalidate_calls;
  for (size_t i = 0; i < flash->count; ++i) {
    FakeStagedHeader& header = flash->headers[i];
    if (header.address != address || !header.present) continue;
    if (!header.invalidate_ok) return false;
    header.present = false;
    header.invalidated = true;
    return true;
  }
  return false;
}

}  // namespace

TEST(OtaFlashDiscard, FreshIdleObjectStillInvalidatesPersistedHeader) {
  // The callback context models flash only: there is deliberately no live
  // OtaStore session or staged_size. Cancellation after reboot must still
  // find the bottom-aligned persistent header.
  FakeStagedHeader headers[] = {
      {0xE2000u, EXPANDED - 0xE2000u, true, true, false},
  };
  FakeStagedFlash flash{headers, 1u, 0u};
  uint32_t invalidated = 99u;

  EXPECT_TRUE(mota_nrf52_discard_staged_headers(
      APP_V6, APP_END_V6, EXPANDED, &flash, fake_read_staged_header,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(invalidated, 1u);
  EXPECT_EQ(flash.invalidate_calls, 1u);
  EXPECT_TRUE(headers[0].invalidated);

  // Durable discard is idempotent: a new scan cannot reopen or re-invalidate
  // the header that the first call consumed.
  invalidated = 99u;
  EXPECT_TRUE(mota_nrf52_discard_staged_headers(
      APP_V6, APP_END_V6, EXPANDED, &flash, fake_read_staged_header,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(invalidated, 0u);
  EXPECT_EQ(flash.invalidate_calls, 1u);
}

TEST(OtaFlashDiscard, InvalidatesEveryCurrentAndLegacyLayoutHeader) {
  FakeStagedHeader headers[] = {
      // Two expanded-ceiling containers from transfers of different sizes.
      {0xE2000u, EXPANDED - 0xE2000u, true, true, false},
      {0xD8000u, EXPANDED - 0xD8000u, true, true, false},
      // A retained container from the older D4000 ceiling geometry.
      {0xC0000u, LEGACY - 0xC0000u, true, true, false},
      // Magic at a page whose total does not bottom-align there is not a
      // reopenable store header and must remain untouched.
      {0xD9000u, 0x8000u, true, true, false},
  };
  FakeStagedFlash flash{headers, 4u, 0u};
  uint32_t invalidated = 0;

  EXPECT_TRUE(mota_nrf52_discard_staged_headers(
      APP_V6, APP_END_V6, EXPANDED, &flash, fake_read_staged_header,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(invalidated, 3u);
  EXPECT_EQ(flash.invalidate_calls, 3u);
  EXPECT_TRUE(headers[0].invalidated);
  EXPECT_TRUE(headers[1].invalidated);
  EXPECT_TRUE(headers[2].invalidated);
  EXPECT_TRUE(headers[3].present);
  EXPECT_FALSE(headers[3].invalidated);
}

TEST(OtaFlashDiscard, ReportsFailureButContinuesInvalidatingOtherHeaders) {
  FakeStagedHeader headers[] = {
      {0xE2000u, EXPANDED - 0xE2000u, true, false, false},
      {0xD8000u, EXPANDED - 0xD8000u, true, true, false},
  };
  FakeStagedFlash flash{headers, 2u, 0u};
  uint32_t invalidated = 0;

  EXPECT_FALSE(mota_nrf52_discard_staged_headers(
      APP_V6, APP_END_V6, EXPANDED, &flash, fake_read_staged_header,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(invalidated, 1u);
  EXPECT_EQ(flash.invalidate_calls, 2u);
  EXPECT_TRUE(headers[0].present);
  EXPECT_TRUE(headers[1].invalidated);
}

TEST(OtaFlashDiscard, UnsafeBoundsFailBeforeAnyFlashCallback) {
  FakeStagedHeader headers[] = {
      {0xE2000u, EXPANDED - 0xE2000u, true, true, false},
  };
  FakeStagedFlash flash{headers, 1u, 0u};
  uint32_t invalidated = 99u;

  EXPECT_FALSE(mota_nrf52_discard_staged_headers(
      APP_V6, EXPANDED + 1u, EXPANDED, &flash, fake_read_staged_header,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(invalidated, 0u);
  EXPECT_EQ(flash.invalidate_calls, 0u);
  EXPECT_TRUE(headers[0].present);

  EXPECT_FALSE(mota_nrf52_discard_staged_headers(
      APP_V6, APP_END_V6, EXPANDED, &flash, nullptr,
      fake_invalidate_staged_header, &invalidated));
  EXPECT_EQ(flash.invalidate_calls, 0u);
}

TEST(OtaEsp32FlashPlan, FullAndDeltaPlacementStaySectorAligned) {
  constexpr uint32_t partition = 2u * 1024u * 1024u;
  constexpr uint32_t sector = 4096u;
  constexpr uint32_t meta_capacity = 65536u;
  MotaEsp32StageLayout full;
  ASSERT_TRUE(mota_esp32_stage_layout(
      partition, sector, meta_capacity, true, 1024u * 1024u, 40000u,
      1024u * 1024u, full));
  EXPECT_EQ(full.total, 40000u + 1024u * 1024u + 5u);
  EXPECT_EQ(full.meta_flush, 40960u);
  EXPECT_EQ(full.meta_part, partition - full.meta_flush);
  EXPECT_EQ(full.meta_part % sector, 0u);
  EXPECT_EQ(full.pay_part0, 0u);

  MotaEsp32StageLayout delta;
  ASSERT_TRUE(mota_esp32_stage_layout(
      partition, sector, meta_capacity, false, 900000u, 1234u,
      100000u, delta));
  EXPECT_EQ(delta.total, 101239u);
  EXPECT_EQ(delta.meta_span, sector);
  EXPECT_EQ(delta.meta_part, delta.write_start);
  EXPECT_EQ(delta.write_start % sector, 0u);
  EXPECT_EQ(delta.pay_part0, delta.write_start + sector);
}

TEST(OtaEsp32FlashPlan, RejectsOverflowAndUnfittableMetadata) {
  MotaEsp32StageLayout layout;
  EXPECT_FALSE(mota_esp32_stage_layout(
      4097u, 4096u, 65536u, true, 1u, 200u, 1u, layout));
  EXPECT_FALSE(mota_esp32_stage_layout(
      2u * 1024u * 1024u, 4096u, 65536u, true, 1u, UINT32_MAX,
      1u, layout));
  EXPECT_FALSE(mota_esp32_stage_layout(
      32768u, 4096u, 65536u, true, 1u, 40000u, 1u, layout));
  EXPECT_FALSE(mota_esp32_stage_layout(
      2u * 1024u * 1024u, 4096u, 65536u, false,
      2u * 1024u * 1024u, 1000u, 1000u, layout));
}

TEST(OtaEsp32FlashPlan, RejectsMalformedOrOversizedFullPayloadBeforeErase) {
  constexpr uint32_t partition = 2u * 1024u * 1024u;
  constexpr uint32_t sector = 4096u;
  constexpr uint32_t meta_capacity = 65536u;
  MotaEsp32StageLayout layout;

  // FULL means that the payload is exactly the final application image.
  EXPECT_FALSE(mota_esp32_stage_layout(
      partition, sector, meta_capacity, true, 100000u, 2000u,
      99999u, layout));
  EXPECT_FALSE(mota_esp32_stage_layout(
      partition, sector, meta_capacity, true, 100000u, 2000u,
      100001u, layout));

  // The logical container and rounded metadata/payload regions must both fit
  // wholly inside the inactive slot before OtaStoreFlashEsp32::begin().
  EXPECT_FALSE(mota_esp32_stage_layout(
      partition, sector, meta_capacity, true, partition, 2000u,
      partition, layout));
  EXPECT_FALSE(mota_esp32_stage_layout(
      8192u, sector, meta_capacity, true, 4097u, 1u, 4097u,
      layout));
}

namespace {

struct FakeEspStagedHeader {
  uint32_t offset;
  bool reopenable;
  bool probe_ok;
  bool invalidate_ok;
  bool invalidated;
};

struct FakeEspPartition {
  FakeEspStagedHeader* headers;
  size_t count;
  uint32_t probe_calls;
  uint32_t invalidate_calls;
};

bool fake_probe_esp_header(void* context, uint32_t offset,
                           bool& reopenable) {
  FakeEspPartition* partition = static_cast<FakeEspPartition*>(context);
  ++partition->probe_calls;
  reopenable = false;
  for (size_t i = 0; i < partition->count; ++i) {
    FakeEspStagedHeader& header = partition->headers[i];
    if (header.offset != offset || header.invalidated) continue;
    reopenable = header.reopenable;
    return header.probe_ok;
  }
  return true;
}

bool fake_invalidate_esp_header(void* context, uint32_t offset) {
  FakeEspPartition* partition = static_cast<FakeEspPartition*>(context);
  ++partition->invalidate_calls;
  for (size_t i = 0; i < partition->count; ++i) {
    FakeEspStagedHeader& header = partition->headers[i];
    if (header.offset != offset || header.invalidated) continue;
    if (!header.invalidate_ok) return false;
    header.invalidated = true;
    return true;
  }
  return false;
}

}  // namespace

TEST(OtaEsp32FlashDiscard, FreshIdleScanInvalidatesEveryReopenableHeader) {
  constexpr uint32_t sector = 4096u;
  constexpr uint32_t partition_size = 6u * sector;
  FakeEspStagedHeader headers[] = {
      {5u * sector, true, true, true, false},
      {3u * sector, false, true, true, false},
      {1u * sector, true, true, true, false},
  };
  FakeEspPartition partition{headers, 3u, 0u, 0u};
  uint32_t invalidated = 99u;

  EXPECT_TRUE(mota_esp32_discard_staged_headers(
      partition_size, sector, &partition, fake_probe_esp_header,
      fake_invalidate_esp_header, &invalidated));
  EXPECT_EQ(partition.probe_calls, 6u);
  EXPECT_EQ(partition.invalidate_calls, 2u);
  EXPECT_EQ(invalidated, 2u);
  EXPECT_TRUE(headers[0].invalidated);
  EXPECT_FALSE(headers[1].invalidated);
  EXPECT_TRUE(headers[2].invalidated);
}

TEST(OtaEsp32FlashDiscard, IoAndInvalidateFailuresAreReportedAfterFullScan) {
  constexpr uint32_t sector = 4096u;
  constexpr uint32_t partition_size = 6u * sector;
  FakeEspStagedHeader headers[] = {
      {4u * sector, false, false, true, false},
      {3u * sector, true, true, false, false},
      {1u * sector, true, true, true, false},
  };
  FakeEspPartition partition{headers, 3u, 0u, 0u};
  uint32_t invalidated = 0u;

  EXPECT_FALSE(mota_esp32_discard_staged_headers(
      partition_size, sector, &partition, fake_probe_esp_header,
      fake_invalidate_esp_header, &invalidated));
  EXPECT_EQ(partition.probe_calls, 6u);
  EXPECT_EQ(partition.invalidate_calls, 2u);
  EXPECT_EQ(invalidated, 1u);
  EXPECT_FALSE(headers[1].invalidated);
  EXPECT_TRUE(headers[2].invalidated);
}

TEST(OtaEsp32FlashDiscard, InvalidScanArgumentsFailWithoutCallbacks) {
  FakeEspPartition partition{nullptr, 0u, 0u, 0u};
  uint32_t invalidated = 99u;
  EXPECT_FALSE(mota_esp32_discard_staged_headers(
      4095u, 4096u, &partition, fake_probe_esp_header,
      fake_invalidate_esp_header, &invalidated));
  EXPECT_FALSE(mota_esp32_discard_staged_headers(
      4097u, 4096u, &partition, fake_probe_esp_header,
      fake_invalidate_esp_header, &invalidated));
  EXPECT_EQ(invalidated, 0u);
  EXPECT_EQ(partition.probe_calls, 0u);
  EXPECT_EQ(partition.invalidate_calls, 0u);
}

// out_start is only written on success - a rejected plan must not clobber the caller's variable.
TEST(OtaFlashPlan, LeavesOutputUntouchedOnReject) {
  uint32_t start = 0x1234ABCD;
  EXPECT_FALSE(mota_nrf52_stage_plan(CAP_V6 + 1, APP_V6, APP_END_V6, LEGACY, start));
  EXPECT_EQ(start, 0x1234ABCDu);
}