// ============================================================================
// payload_buffer.cpp — Multi-chunk LoRa→PSRAM payload staging
//
// See payload_buffer.h for full API and wire-format documentation.
// ============================================================================

#include "payload_buffer.h"

#ifndef MSG_TYPE_PAYLOAD_CHUNK
#  ifdef UNIT_TEST
#    define MSG_TYPE_PAYLOAD_CHUNK 0x14
#  else
#    include "config.h"
#  endif
#endif

#ifndef UNIT_TEST
#  include "esp_heap_caps.h"   // heap_caps_malloc / free
#else
// Host-side stubs: map heap_caps_malloc to plain malloc so tests compile.
#  include <cstdlib>
#  include <cstring>
#  ifndef MALLOC_CAP_SPIRAM
#    define MALLOC_CAP_SPIRAM  0x00000008u
#  endif
static inline void* heap_caps_malloc(size_t size, uint32_t /*caps*/) {
    return malloc(size);
}
static inline void heap_caps_free(void* ptr) { free(ptr); }
#endif

#include <cstring>
#include <cstdlib>

// ---------------------------------------------------------------------------
// CRC32/BZIP2 (same poly as command_handler.cpp)
// ---------------------------------------------------------------------------
static uint32_t crc32_bzip2(const uint8_t* data, size_t len) {
    uint32_t crc = 0xFFFFFFFFu;
    for (size_t i = 0; i < len; i++) {
        crc ^= ((uint32_t)data[i]) << 24;
        for (int b = 0; b < 8; b++) {
            if (crc & 0x80000000u) crc = (crc << 1) ^ 0x04C11DB7u;
            else                   crc <<= 1;
        }
    }
    return crc ^ 0xFFFFFFFFu;
}

// ---------------------------------------------------------------------------
// Internal slot bookkeeping
// ---------------------------------------------------------------------------
struct PayloadSlot {
    uint32_t  payload_id;
    uint32_t  total_size;
    uint16_t  total_chunks;
    uint16_t  received_chunks;
    uint8_t*  buffer;         // PSRAM allocation, NULL if slot free
    bool      complete;       // CRC verified, ready for delivery
    // Received-chunk bitmap: up to 65535 chunks but we cap at 256KB/1B = 256K
    // chunks (unrealistically large). Use a heap bitmap sized at allocation.
    uint8_t*  rx_bitmap;      // one bit per chunk index, stored in PSRAM
};

static PayloadSlot s_slots[PAYLOAD_CONCURRENT_MAX];

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static int find_slot(uint32_t payload_id) {
    for (int i = 0; i < PAYLOAD_CONCURRENT_MAX; i++) {
        if (s_slots[i].buffer && s_slots[i].payload_id == payload_id) return i;
    }
    return -1;
}

static int free_slot_index() {
    for (int i = 0; i < PAYLOAD_CONCURRENT_MAX; i++) {
        if (!s_slots[i].buffer) return i;
    }
    return -1;
}

static void release_slot(int idx) {
    if (s_slots[idx].buffer) {
        heap_caps_free(s_slots[idx].buffer);
        s_slots[idx].buffer = nullptr;
    }
    if (s_slots[idx].rx_bitmap) {
        heap_caps_free(s_slots[idx].rx_bitmap);
        s_slots[idx].rx_bitmap = nullptr;
    }
    memset(&s_slots[idx], 0, sizeof(PayloadSlot));
}

// Bitmap helpers
static inline void bitmap_set(uint8_t* bm, uint16_t idx) {
    bm[idx >> 3] |= (uint8_t)(1u << (idx & 7u));
}
static inline bool bitmap_test(const uint8_t* bm, uint16_t idx) {
    return (bm[idx >> 3] & (uint8_t)(1u << (idx & 7u))) != 0;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

void payload_buffer_init() {
    for (int i = 0; i < PAYLOAD_CONCURRENT_MAX; i++) {
        release_slot(i);
    }
}

PayloadBufStatus payload_buffer_handle_chunk(const uint8_t* buf, size_t len) {
    // Minimum: 1B type + 23B rest of header = 24B, plus at least 0B data
    if (!buf || len < PAYLOAD_CHUNK_HDR_LEN) return PB_BAD_HDR;
    if (buf[0] != MSG_TYPE_PAYLOAD_CHUNK)     return PB_BAD_HDR;

    // Parse header fields (all LE)
    uint32_t payload_id;
    uint16_t total_chunks;
    uint16_t chunk_idx;
    uint32_t total_size;
    uint32_t chunk_crc32;

    memcpy(&payload_id,   buf + 1,  4);
    memcpy(&total_chunks, buf + 5,  2);
    memcpy(&chunk_idx,    buf + 7,  2);
    memcpy(&total_size,   buf + 9,  4);
    memcpy(&chunk_crc32,  buf + 13, 4);
    // bytes [17..23] = reserved, ignored

    const uint8_t* chunk_data = buf + PAYLOAD_CHUNK_HDR_LEN;
    size_t         chunk_len  = len - PAYLOAD_CHUNK_HDR_LEN;

    // Sanity: chunk_idx must be < total_chunks
    if (total_chunks == 0 || chunk_idx >= total_chunks) return PB_BAD_HDR;
    // Payload size cap
    if (total_size == 0 || total_size > PAYLOAD_MAX_SIZE) return PB_TOO_LARGE;

    // Per-chunk CRC32 verification.
    // Exception: the final chunk (chunk_idx == total_chunks-1) carries the
    // FULL-PAYLOAD CRC32 in the chunk_crc32 field instead of a per-chunk CRC.
    // Skip the per-chunk check for the final chunk; verify full-payload CRC
    // after all chunks have been assembled (see below).
    bool is_final_chunk = (chunk_idx == total_chunks - 1);
    if (!is_final_chunk) {
        uint32_t computed = crc32_bzip2(chunk_data, chunk_len);
        if (computed != chunk_crc32) return PB_BAD_CHUNK_CRC;
    }

    // Locate or allocate a slot
    int sidx = find_slot(payload_id);
    if (sidx < 0) {
        sidx = free_slot_index();
        if (sidx < 0) return PB_FULL;

        // Allocate PSRAM for the full payload buffer
        uint8_t* pbuf = (uint8_t*)heap_caps_malloc(total_size, MALLOC_CAP_SPIRAM);
        if (!pbuf) return PB_ALLOC_FAIL;

        // Allocate bitmap: ceil(total_chunks / 8) bytes
        size_t   bm_bytes = ((size_t)total_chunks + 7u) / 8u;
        uint8_t* bm = (uint8_t*)heap_caps_malloc(bm_bytes, MALLOC_CAP_SPIRAM);
        if (!bm) { heap_caps_free(pbuf); return PB_ALLOC_FAIL; }
        memset(bm, 0, bm_bytes);
        memset(pbuf, 0, total_size);

        s_slots[sidx].payload_id       = payload_id;
        s_slots[sidx].total_size       = total_size;
        s_slots[sidx].total_chunks     = total_chunks;
        s_slots[sidx].received_chunks  = 0;
        s_slots[sidx].buffer           = pbuf;
        s_slots[sidx].rx_bitmap        = bm;
        s_slots[sidx].complete         = false;
    }

    PayloadSlot& slot = s_slots[sidx];

    // Duplicate chunk guard
    if (bitmap_test(slot.rx_bitmap, chunk_idx)) return PB_DUPLICATE;

    // Compute write offset: equal to chunk_idx * (PAYLOAD_CHUNK_MAX_DATA) is
    // WRONG for variable-length last chunks. Instead we use a simple scheme:
    // sender must pack chunks contiguously; chunk_idx * chunk_len won't work
    // for variable last chunk. The correct approach is that the C2 sends
    // equal-sized chunks (up to PAYLOAD_CHUNK_MAX_DATA bytes each) except
    // possibly the last. Write offset = chunk_idx * PAYLOAD_CHUNK_MAX_DATA.
    uint32_t write_offset = (uint32_t)chunk_idx * PAYLOAD_CHUNK_MAX_DATA;
    if (write_offset + chunk_len > slot.total_size) {
        // Guard overflow — shouldn't happen with a well-formed sender
        return PB_BAD_HDR;
    }

    memcpy(slot.buffer + write_offset, chunk_data, chunk_len);
    bitmap_set(slot.rx_bitmap, chunk_idx);
    slot.received_chunks++;

    if (slot.received_chunks < slot.total_chunks) return PB_OK;

    // All chunks received — verify full-payload CRC32
    // Protocol: sender sends the full-payload CRC32 in the LAST chunk's
    // chunk_crc32 field (chunk_idx == total_chunks - 1).  Per-chunk CRCs
    // cover individual chunk payloads; the final chunk's CRC32 field is
    // reused to carry the aggregate CRC32 of the entire assembled buffer.
    //
    // If the final chunk IS the one we just wrote (chunk_idx == total-1):
    //   verify assembled[0..total_size-1] against the stored final_crc value.
    // This means the integrator on the T2 side MUST place the aggregate CRC
    // in the final chunk's chunk_crc32 field.
    if (chunk_idx == slot.total_chunks - 1) {
        uint32_t full_crc = crc32_bzip2(slot.buffer, slot.total_size);
        // We stored chunk_crc32 from this (final) chunk; reuse it as the
        // full-payload reference. Because we already verified per-chunk CRC
        // above against chunk_crc32, and here we want the full CRC, the
        // protocol requires the sender to put the full-buffer CRC into the
        // final chunk's header field. We accept if they match.
        if (full_crc != chunk_crc32) {
            release_slot(sidx);
            return PB_BAD_FINAL_CRC;
        }
    }

    slot.complete = true;
    return PB_COMPLETE;
}

bool payload_buffer_get(uint32_t payload_id,
                        const uint8_t** data_out,
                        size_t*         len_out) {
    if (!data_out || !len_out) return false;
    int sidx = find_slot(payload_id);
    if (sidx < 0) return false;
    if (!s_slots[sidx].complete) return false;
    *data_out = s_slots[sidx].buffer;
    *len_out  = s_slots[sidx].total_size;
    return true;
}

void payload_buffer_release(uint32_t payload_id) {
    int sidx = find_slot(payload_id);
    if (sidx >= 0) release_slot(sidx);
}
