#include "OtaStoreFlashEsp32.h"

#if defined(ESP32_PLATFORM) && defined(OTA_FLASH_STORE)

#include "OtaDebug.h"
#include "OtaByteIO.h"          // rd_u32le (persisted header total)
#include "MotaContainer.h"      // mota_parse_manifest (reopen: rebuild geometry from the staged manifest)
#include <string.h>
#include <stdlib.h>             // malloc/free (the meta buffer is sized per fetch)
#include "esp_ota_ops.h"        // esp_ota_get_next_update_partition (the inactive A/B slot)

namespace mesh {
namespace ota {

OtaStoreFlashEsp32::~OtaStoreFlashEsp32() { free(_meta); }

bool OtaStoreFlashEsp32::acquire() {
  if (!_part) {
    _part = esp_ota_get_next_update_partition(nullptr);
    _psize = _part ? _part->size : 0;
  }
  return _part != nullptr;
}

void OtaStoreFlashEsp32::set_layout(const MotaEsp32StageLayout& layout) {
  _total = layout.total;
  _meta_span = layout.meta_span;
  _meta_flush = layout.meta_flush;
  _write_start = layout.write_start;
  _meta_part = layout.meta_part;
  _pay_log0 = layout.pay_log0;
  _pay_part0 = layout.pay_part0;
  _trailer_part = _full ? _meta_part + _meta_bytes
                         : _write_start + _total - 5u;
}

// Compute the slot placement from _full/_image_size/_meta_bytes/_pay_size (already set). Returns false if
// it won't fit. Shared by plan_layout (fresh fetch), reopen (resume), and discard validation so all three
// derive identical geometry.
bool OtaStoreFlashEsp32::layout() {
  MotaEsp32StageLayout planned;
  if (!mota_esp32_stage_layout(
          _psize, SEC, OTA_ESP32_META_CAP, _full, _image_size,
          _meta_bytes, _pay_size, planned)) {
    return false;
  }
  set_layout(planned);
  return true;
}

// Choose placement from the parsed manifest and refuse anything that won't fit, BEFORE any block is
// staged. (See the header for the delta-vs-full layout rationale.)
bool OtaStoreFlashEsp32::plan_layout(bool is_full, uint32_t image_size,
                                     uint32_t payload_off, uint32_t payload_size,
                                     bool is_bootloader) {
  if (is_bootloader) return false;
  if (!acquire()) return false;
  _full = is_full; _image_size = image_size; _meta_bytes = payload_off; _pay_size = payload_size;
  bool ok = layout();
  OTA_DBG("OTA esp32: plan %s total=%u image=%u meta=%u write_start=%u meta_part=%u ok=%d\n",
          is_full ? "FULL" : "DELTA", (unsigned)_total, (unsigned)image_size, (unsigned)_meta_bytes,
          (unsigned)_write_start, (unsigned)_meta_part, (int)ok);
  return ok;
}

bool OtaStoreFlashEsp32::begin(uint32_t total_size) {
  if (!_part || _total == 0 || total_size != _total) return false;   // plan_layout must have run + agreed
  free(_meta);
  _meta = (uint8_t*)malloc(_meta_flush);                             // header+manifest+leaves(+full trailer)
  if (!_meta) { _total = 0; return false; }
  memset(_meta, 0xFF, _meta_flush);
  memset(_trailer, 0xFF, sizeof(_trailer));
  _pay_open = false; _pay_sec = 0; _pay_seen_end = 0; _flushed = false;
  _io_ok = true;
  return true;
}

void OtaStoreFlashEsp32::clear() {
  free(_meta); _meta = nullptr;
  _total = 0; _pay_open = false; _pay_seen_end = 0; _flushed = false;   // _part kept (re-acquire is fine)
}

OtaStoreFlashEsp32::CandidateProbe OtaStoreFlashEsp32::probe_candidate(
    uint32_t offset, StagedCandidate& candidate) const {
  if (!_part || _psize < SEC || offset > _psize - SEC) {
    return CandidateProbe::INVALID;
  }

  uint8_t header[8];
  if (esp_partition_read(_part, offset, header, sizeof(header)) != ESP_OK) {
    return CandidateProbe::IO_ERROR;
  }
  if (memcmp(header, MOTA_MAGIC, sizeof(MOTA_MAGIC)) != 0) {
    return CandidateProbe::INVALID;
  }

  const uint32_t total = rd_u32le(header + sizeof(MOTA_MAGIC));
  if (total < 8u + MOTA_MFL + 5u || total > _psize) {
    return CandidateProbe::INVALID;
  }

  uint8_t manifest_bytes[256];
  uint32_t manifest_read = total - 8u;
  if (manifest_read > sizeof(manifest_bytes)) {
    manifest_read = sizeof(manifest_bytes);
  }
  if (esp_partition_read(_part, offset + 8u, manifest_bytes,
                         manifest_read) != ESP_OK) {
    return CandidateProbe::IO_ERROR;
  }

  MotaManifest manifest;
  if (!mota_parse_manifest(manifest_bytes, manifest_read, manifest)) {
    return CandidateProbe::INVALID;
  }
  const uint32_t manifest_len =
      (uint32_t)(manifest.approval - manifest.manifest_start) +
      sizeof(APPROVAL_YES);
  const uint64_t payload_off64 =
      8u + (uint64_t)manifest_len + (uint64_t)manifest.block_count * 4u;
  if (payload_off64 > UINT32_MAX ||
      payload_off64 + (uint64_t)manifest.payload_size + 5u != total) {
    return CandidateProbe::INVALID;
  }

  MotaEsp32StageLayout planned;
  if (!mota_esp32_stage_layout(
          _psize, SEC, OTA_ESP32_META_CAP, manifest.is_full(),
          manifest.image_size, (uint32_t)payload_off64,
          manifest.payload_size, planned) ||
      planned.meta_part != offset || planned.total != total) {
    return CandidateProbe::INVALID;
  }

  candidate.layout = planned;
  candidate.full = manifest.is_full();
  candidate.image_size = manifest.image_size;
  candidate.meta_bytes = (uint32_t)payload_off64;
  candidate.pay_size = manifest.payload_size;
  memcpy(candidate.mid, manifest.merkle_root, sizeof(candidate.mid));
  candidate.target = manifest.target_id;
  return CandidateProbe::VALID;
}

bool OtaStoreFlashEsp32::probe_staged_header(
    void* context, uint32_t offset, bool& reopenable) {
  OtaStoreFlashEsp32* store = static_cast<OtaStoreFlashEsp32*>(context);
  StagedCandidate candidate;
  const CandidateProbe result = store->probe_candidate(offset, candidate);
  reopenable = result == CandidateProbe::VALID;
  return result != CandidateProbe::IO_ERROR;
}

bool OtaStoreFlashEsp32::invalidate_staged_header(
    void* context, uint32_t offset) {
  OtaStoreFlashEsp32* store = static_cast<OtaStoreFlashEsp32*>(context);
  if (!store->_part || store->_psize < SEC || offset > store->_psize - SEC ||
      esp_partition_read(store->_part, offset, store->_pay, SEC) != ESP_OK) {
    return false;
  }

  // Both markers are inside the already-validated first metadata sector.
  // One sector RMW avoids erasing a delta's adjacent metadata/payload twice.
  memset(store->_pay, 0, sizeof(MOTA_MAGIC));
  memset(store->_pay + 8u + MOTA_OFF_APPROVAL, 0,
         sizeof(APPROVAL_YES));
  if (esp_partition_erase_range(store->_part, offset, SEC) != ESP_OK ||
      esp_partition_write(store->_part, offset, store->_pay, SEC) != ESP_OK) {
    return false;
  }

  uint8_t check[sizeof(MOTA_MAGIC)];
  if (esp_partition_read(store->_part, offset, check, sizeof(check)) != ESP_OK ||
      memcmp(check, store->_pay, sizeof(check)) != 0 ||
      esp_partition_read(store->_part, offset + 8u + MOTA_OFF_APPROVAL,
                         check, sizeof(check)) != ESP_OK ||
      memcmp(check, store->_pay + 8u + MOTA_OFF_APPROVAL,
             sizeof(check)) != 0) {
    return false;
  }
  return true;
}

bool OtaStoreFlashEsp32::discard() {
  bool ok = acquire();
  uint32_t invalidated = 0;
  if (ok) {
    ok = mota_esp32_discard_staged_headers(
        _psize, SEC, this, probe_staged_header,
        invalidate_staged_header, &invalidated);
  }
  OTA_DBG("OTA esp32: discard invalidated=%u ok=%d\n",
          (unsigned)invalidated, (int)ok);
  clear();
  _io_ok = ok;
  return ok;
}

uint8_t* OtaStoreFlashEsp32::meta_slot(uint32_t L) {
  if (in_trailer(L)) return _full ? (_meta + _meta_bytes + (L - (_total - 5)))
                                  : (_trailer + (L - (_total - 5)));
  if (L < _meta_span) return _meta + L;
  return nullptr;                                    // payload -> sliding sector / flash
}
const uint8_t* OtaStoreFlashEsp32::meta_slot_c(uint32_t L) const {
  if (in_trailer(L)) return _full ? (_meta + _meta_bytes + (L - (_total - 5)))
                                  : (_trailer + (L - (_total - 5)));
  if (L < _meta_span) return _meta + L;
  return nullptr;
}

// Bytes from `pos` that stay in one region, and (for payload) one flash sector.
uint32_t OtaStoreFlashEsp32::run(uint32_t pos, uint32_t remain) const {
  if (pos >= _total - 5) return remain;                          // trailer (<=5, one buffer)
  if (pos < _meta_span) {
    uint32_t end = _meta_span < _total - 5u ? _meta_span : _total - 5u;
    uint32_t c = end - pos;
    return remain < c ? remain : c;                  // never consume the separate trailer through metadata
  }
  uint32_t poff = pay_part(pos);
  uint32_t to_sec = SEC - (poff % SEC);
  uint32_t to_end = (_total - 5) - pos;                          // don't cross into the trailer
  uint32_t c = remain;
  if (to_sec < c) c = to_sec;
  if (to_end < c) c = to_end;
  return c;
}

void OtaStoreFlashEsp32::flush_sector(uint32_t slot_off, const uint8_t* buf, uint32_t n) {
  if (!_io_ok) return;
  if (esp_partition_erase_range(_part, slot_off, n) != ESP_OK) { _io_ok = false; return; }
  if (esp_partition_write(_part, slot_off, buf, n) != ESP_OK)   { _io_ok = false; return; }
  OTA_DBG("OTA esp32: flushed %u B @ slot+%u\n", (unsigned)n, (unsigned)slot_off);
}

void OtaStoreFlashEsp32::flush_pay() {
  if (_pay_open) { flush_sector((uint32_t)_pay_sec * SEC, _pay, SEC); _pay_open = false; }
}

void OtaStoreFlashEsp32::open_pay(uint32_t sec) {
  if (_pay_open) flush_pay();
  if (sec < _pay_seen_end) {
    // Any sector up to and INCLUDING the highest opened sector may contain
    // acknowledged blocks. Preserve it even after visiting a lower sector
    // (e.g. full-image blocks 4,3,5 visit sectors 1,0,1). An exclusive bound
    // also distinguishes a genuinely fresh sector 0 from a revisit to it.
    // Some gaps below the bound may be unvisited, but their leaves are absent;
    // conservatively keeping those bytes cannot erase a committed block.
    if (esp_partition_read(_part, (size_t)sec * SEC, _pay, SEC) != ESP_OK) _io_ok = false;
  } else {
    memset(_pay, 0xFF, SEC);                          // fresh sector
    _pay_seen_end = sec + 1u;
  }
  _pay_sec = sec; _pay_open = true;
}

bool OtaStoreFlashEsp32::write(uint32_t offset, const uint8_t* d, uint32_t len) {
  if ((uint64_t)offset + len > _total || !_io_ok) return false;
  if (len != 0) _flushed = false;   // a later finalize/checkpoint must persist continued writes
  for (uint32_t pos = offset, end = offset + len; pos < end; ) {
    uint32_t n = run(pos, end - pos);
    if (uint8_t* dst = meta_slot(pos)) {              // meta / leaves / trailer -> pinned RAM
      memcpy(dst, d, n);
    } else {                                          // payload -> sliding sector
      uint32_t poff = pay_part(pos);
      uint32_t sec  = poff / SEC;
      if (!_pay_open || sec != _pay_sec) open_pay(sec);
      memcpy(_pay + (poff % SEC), d, n);
    }
    pos += n; d += n;
    if (!_io_ok) return false;
  }
  return true;
}

bool OtaStoreFlashEsp32::read(uint32_t offset, uint8_t* buf, uint32_t len) const {
  if ((uint64_t)offset + len > _total) return false;
  for (uint32_t pos = offset, end = offset + len; pos < end; ) {
    uint32_t n = run(pos, end - pos);
    if (const uint8_t* src = meta_slot_c(pos)) {
      memcpy(buf, src, n);
    } else {
      uint32_t poff = pay_part(pos);
      if (_pay_open && poff / SEC == _pay_sec) memcpy(buf, _pay + (poff % SEC), n);   // still in RAM
      else if (esp_partition_read(_part, poff, buf, n) != ESP_OK) return false;       // flushed -> flash
    }
    pos += n; buf += n;
  }
  return true;
}

bool OtaStoreFlashEsp32::finalize() {
  if (_flushed) return _io_ok;
  if (_total == 0 || !_io_ok) return false;
  if (!_full) {
    // delta: drop the 5-byte trailer into the sliding payload sector(s) so it flushes with them (the
    // trailer sits right after the payload; this also covers the rare case it spills into a fresh sector).
    uint32_t tpoff = _write_start + (_total - 5);
    for (uint32_t off = 0; off < 5; ) {
      const uint32_t logical = _total - 5u + off;
      if (logical < _meta_span) {
        // A small delta's tail can share the pinned metadata sector. Put those
        // bytes in the buffer that will be flushed last, including a split tail.
        uint32_t n = _meta_span - logical;
        if (n > 5u - off) n = 5u - off;
        memcpy(_meta + logical, _trailer + off, n);
        off += n;
        continue;
      }
      uint32_t sec = (tpoff + off) / SEC;
      if (!_pay_open || sec != _pay_sec) open_pay(sec);
      uint32_t in = SEC - ((tpoff + off) % SEC); if (in > 5 - off) in = 5 - off;
      memcpy(_pay + ((tpoff + off) % SEC), _trailer + off, in);
      off += in;
    }
  }
  flush_pay();                                        // last payload sector(s) (+ delta trailer)
  flush_sector(_meta_part, _meta, _meta_flush);       // meta (+ trailer for full)
  _flushed = _io_ok;
  OTA_DBG("OTA esp32: finalize %s io_ok=%d\n", _full ? "FULL" : "DELTA", (int)_io_ok);
  return _io_ok;
}

// Persist mid-transfer progress so a reboot can resume. Flush the open payload sector (KEEP it buffered so
// continued in-order writes aren't lost), then the meta region (leaves). Payload-before-leaves keeps it
// consistent: every block whose leaf is now in flash also has its payload in flash. Infrequent.
void OtaStoreFlashEsp32::checkpoint() {
  if (_total == 0 || _flushed || !_io_ok) return;
  if (_pay_open) flush_sector((uint32_t)_pay_sec * SEC, _pay, SEC);   // flush but leave _pay/_pay_open intact
  flush_sector(_meta_part, _meta, _meta_flush);                       // header + manifest + leaves (+full trailer)
}

// Re-attach to a container already staged in the slot (after a reboot), WITHOUT erasing. The header
// (MOTA_MAGIC + total) sits at a sector boundary (meta_part for full, write_start for delta - both
// sector-aligned), so scan sector starts from the bottom up. A candidate is accepted only if: magic +
// plausible total, the manifest parses, the container geometry is self-consistent with the total, AND the
// recomputed placement lands the meta exactly where we found the magic. Scan all candidates before
// adoption: a smaller, older container can have a higher address than a newer checkpoint.
bool OtaStoreFlashEsp32::reopen() {
  return reopenFor(nullptr, 0);
}

bool OtaStoreFlashEsp32::reopenFor(const uint8_t* want_mid, uint32_t expected_target) {
  if (!acquire() || _psize < SEC) return false;
  bool found = false;
  StagedCandidate selected;
  // The meta/container is staged at the bottom of the slot. Large delta containers can begin well above
  // the final 512 KB, so scan the complete partition rather than silently losing resumability after reboot.
  for (uint32_t o = mota_esp32_align_down(_psize - SEC, SEC); ; o -= SEC) {
    StagedCandidate candidate;
    const CandidateProbe probe = probe_candidate(o, candidate);
    if (probe == CandidateProbe::IO_ERROR) return false;
    if (probe == CandidateProbe::VALID) {
      // No generation marker exists. Unqualified boot resume must not infer
      // recency from scan order, target, version, or apparent block progress.
      const bool matches = !want_mid ||
          (memcmp(candidate.mid, want_mid, sizeof(candidate.mid)) == 0 &&
           (expected_target == 0 || candidate.target == expected_target));
      if (matches) {
        if (found) return false;
        selected = candidate;
        found = true;
      }
    }
    if (o == 0) break;
  }
  if (!found || (expected_target != 0 && selected.target != expected_target)) return false;
  _full = selected.full;
  _image_size = selected.image_size;
  _meta_bytes = selected.meta_bytes;
  _pay_size = selected.pay_size;
  set_layout(selected.layout);
  free(_meta); _meta = (uint8_t*)malloc(_meta_flush);
  if (!_meta) { _total = 0; return false; }
  if (esp_partition_read(_part, _meta_part, _meta, _meta_flush) != ESP_OK) {
    free(_meta); _meta = nullptr; _total = 0; return false;
  }
  // Raw staging defers the fixed trailer; an unvisited sector may still
  // contain old firmware bytes after an early checkpoint. Restore framing
  // from the validated layout, not that stale tail. The manager separately
  // revalidates the selected image and every present payload block.
  memcpy(_trailer, MOTA_TRAILER, sizeof(_trailer));
  if (_full) memcpy(_meta + _meta_bytes, MOTA_TRAILER, sizeof(_trailer));
  _pay_open = false; _pay_sec = 0; _flushed = false; _io_ok = true;
  // No volatile visit history survives reboot. Treat the entire partition
  // as potentially written so every resumed sector preserves committed data.
  _pay_seen_end = _psize / SEC;
  OTA_DBG("OTA esp32: reopen %s total=%u meta_part=%u\n", _full ? "FULL" : "DELTA", (unsigned)_total, (unsigned)_meta_part);
  return true;
}

} // namespace ota
} // namespace mesh

#endif