// ============================================================================
// c2_proxy_tunnel.cpp — BLE↔LoRa C2 API tunnel implementation
//
// See c2_proxy_tunnel.h for full architecture and wire format. This file
// contains:
//   * Two reassembly slot rings (REQ-from-phone, RSP-from-T2)
//   * BLE-write input queue + LoRa-tx output queue (drained from tick())
//   * BLE-notify output queue (drained from tick())
//   * Per-chunk + full-payload CRC32 verification (BZIP2 poly, same as
//     payload_buffer.cpp — chosen so the on-wire bytes are byte-identical
//     to MSG_TYPE_PAYLOAD_CHUNK chunks for any tooling that already groks
//     the existing format)
// ============================================================================

#include "c2_proxy_tunnel.h"

#include <cstring>
#include <cstdlib>

#ifndef UNIT_TEST
#  include <Arduino.h>            // millis()
#  include "esp_heap_caps.h"
#else
#  include <cstdint>
#  include <chrono>
#  ifndef MALLOC_CAP_SPIRAM
#    define MALLOC_CAP_SPIRAM 0x00000008u
#  endif
static inline void* heap_caps_malloc(size_t n, uint32_t /*caps*/) { return malloc(n); }
static inline void  heap_caps_free(void* p) { free(p); }
static inline uint32_t millis() {
    using namespace std::chrono;
    return (uint32_t)duration_cast<milliseconds>(
        steady_clock::now().time_since_epoch()).count();
}
#endif

// Forward decl from ble_companion (weak hook).
extern "C" bool ble_companion_is_authenticated(uint32_t now_ms) __attribute__((weak));

// ---------------------------------------------------------------------------
// CRC32/BZIP2 — identical to payload_buffer.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;
}

// ---------------------------------------------------------------------------
// Reassembly slot
// ---------------------------------------------------------------------------
struct ProxySlot {
    bool      in_use;
    uint32_t  request_id;
    uint32_t  total_size;
    uint16_t  total_chunks;
    uint16_t  received_chunks;
    uint8_t*  buffer;
    uint8_t*  rx_bitmap;
    uint32_t  first_seen_ms;
    bool      complete;
};

static ProxySlot s_req_slots[C2_PROXY_CONCURRENT];
static ProxySlot s_rsp_slots[C2_PROXY_CONCURRENT];

// ---------------------------------------------------------------------------
// Outbound queues. We use simple ring buffers of full wire packets so the
// tick() loop can do non-blocking flushes regardless of LoRa airtime.
// ---------------------------------------------------------------------------
struct WireFrame {
    uint8_t  data[C2_PROXY_HDR_LEN + C2_PROXY_MAX_DATA];
    uint16_t len;
    bool     used;
};

static WireFrame s_lora_tx_queue[C2_PROXY_BLE_TX_QUEUE];
static WireFrame s_ble_tx_queue [C2_PROXY_BLE_TX_QUEUE];

// Forward decl — definition below.
static bool enqueue_fragments_passthrough_pending(WireFrame* ring,
                                                  const uint8_t* buf, size_t len);

// ---------------------------------------------------------------------------
// Stats
// ---------------------------------------------------------------------------
static C2ProxyStats s_stats;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
static inline void bm_set(uint8_t* bm, uint16_t idx) {
    bm[idx >> 3] |= (uint8_t)(1u << (idx & 7u));
}
static inline bool bm_test(const uint8_t* bm, uint16_t idx) {
    return (bm[idx >> 3] & (uint8_t)(1u << (idx & 7u))) != 0;
}

static int find_slot(ProxySlot* ring, uint32_t request_id) {
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++) {
        if (ring[i].in_use && ring[i].request_id == request_id) return i;
    }
    return -1;
}
static int alloc_slot(ProxySlot* ring) {
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++) {
        if (!ring[i].in_use) return i;
    }
    return -1;
}
static void release_slot(ProxySlot& s) {
    if (s.buffer)    { heap_caps_free(s.buffer);    s.buffer    = nullptr; }
    if (s.rx_bitmap) { heap_caps_free(s.rx_bitmap); s.rx_bitmap = nullptr; }
    memset(&s, 0, sizeof(s));
}

// ---------------------------------------------------------------------------
// Generic chunk handler — used for both REQ-from-phone and RSP-from-T2.
// Returns C2P_OK + sets `*completed=true` when the assembled buffer is ready
// (the slot remains in_use; caller releases after consuming).
// ---------------------------------------------------------------------------
static C2ProxyStatus handle_chunk_generic(ProxySlot* ring,
                                          const uint8_t* buf, size_t len,
                                          uint8_t expected_type,
                                          int* out_slot_idx,
                                          bool* completed) {
    if (out_slot_idx) *out_slot_idx = -1;
    if (completed)    *completed    = false;
    // BUG FIX: previously only enforced a LOWER bound on len. The "patched"
    // path in c2_proxy_handle_ble_write (buf[0] != expected type on entry)
    // separately caps len at C2_PROXY_HDR_LEN + C2_PROXY_MAX_DATA before
    // ever reaching here, but the "direct" path (buf[0] already correct)
    // did not apply the same cap -- an oversized single "chunk" (violating
    // the documented <=200-byte-per-chunk wire format) could still pass
    // through and, depending on total_size/chunk_idx math, write past
    // where a well-formed chunk would end. Enforce the same upper bound
    // here so both callers get it uniformly, regardless of which path led
    // here (defense-in-depth: BLE ATT MTU limits make this hard to trigger
    // in practice, but the check should not depend on the caller remembering).
    if (!buf || len < C2_PROXY_HDR_LEN || len > C2_PROXY_HDR_LEN + C2_PROXY_MAX_DATA) {
        return C2P_BAD_HDR;
    }
    if (buf[0] != expected_type)        return C2P_BAD_HDR;

    uint32_t request_id;
    uint16_t total_chunks;
    uint16_t chunk_idx;
    uint32_t total_size;
    uint32_t chunk_crc32;

    memcpy(&request_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);

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

    if (total_chunks == 0 || chunk_idx >= total_chunks) return C2P_BAD_HDR;
    if (total_size == 0 || total_size > C2_PROXY_MAX_PAYLOAD) return C2P_TOO_LARGE;

    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) { s_stats.dropped_bad_crc++; return C2P_BAD_CRC; }
    }

    int sidx = find_slot(ring, request_id);
    if (sidx < 0) {
        sidx = alloc_slot(ring);
        if (sidx < 0) { s_stats.dropped_no_slot++; return C2P_FULL; }
        uint8_t* pbuf = (uint8_t*)heap_caps_malloc(total_size, MALLOC_CAP_SPIRAM);
        if (!pbuf) { s_stats.dropped_no_slot++; return C2P_FULL; }
        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); s_stats.dropped_no_slot++; return C2P_FULL; }
        memset(pbuf, 0, total_size);
        memset(bm,   0, bm_bytes);
        ring[sidx].in_use         = true;
        ring[sidx].request_id     = request_id;
        ring[sidx].total_size     = total_size;
        ring[sidx].total_chunks   = total_chunks;
        ring[sidx].received_chunks = 0;
        ring[sidx].buffer         = pbuf;
        ring[sidx].rx_bitmap      = bm;
        ring[sidx].first_seen_ms  = millis();
        ring[sidx].complete       = false;
    }

    ProxySlot& slot = ring[sidx];
    if (bm_test(slot.rx_bitmap, chunk_idx)) return C2P_DUPLICATE;

    uint32_t write_offset = (uint32_t)chunk_idx * C2_PROXY_MAX_DATA;
    if (write_offset + chunk_len > slot.total_size) return C2P_BAD_HDR;
    memcpy(slot.buffer + write_offset, chunk_data, chunk_len);
    bm_set(slot.rx_bitmap, chunk_idx);
    slot.received_chunks++;

    if (slot.received_chunks < slot.total_chunks) {
        if (out_slot_idx) *out_slot_idx = sidx;
        return C2P_OK;
    }

    // Final chunk → verify full-payload CRC stored in last chunk's crc field
    if (is_final_chunk) {
        uint32_t full_crc = crc32_bzip2(slot.buffer, slot.total_size);
        if (full_crc != chunk_crc32) {
            release_slot(slot);
            s_stats.dropped_bad_crc++;
            return C2P_BAD_CRC;
        }
    }

    slot.complete = true;
    if (completed)    *completed = true;
    if (out_slot_idx) *out_slot_idx = sidx;
    return C2P_OK;
}

// ---------------------------------------------------------------------------
// Weak default hooks (overridable by main.cpp)
//
// c2_proxy_lora_tx/c2_proxy_ble_notify's weak DEFINITIONS (not just
// declarations) live in this same translation unit -- fine for the real
// firmware build (main.cpp's strong override in a separate .o wins at
// link time, standard weak-symbol resolution), but a native Unity test
// that pulls this whole .cpp in via a single-TU #include cannot ALSO
// define a strong override of the same name without a same-TU redefinition
// error (weak-vs-strong resolution only happens across separate object
// files, not within one). Guarded out under UNIT_TEST so host tests always
// supply their own hooks -- consistent with how every other stubbed
// dependency in this codebase's tests works.
// ---------------------------------------------------------------------------
#ifndef UNIT_TEST
__attribute__((weak)) bool c2_proxy_lora_tx(const uint8_t* /*buf*/, size_t /*len*/) {
    // Default: not wired. main.cpp must override and route to RadioLib.
    return false;
}
__attribute__((weak)) bool c2_proxy_ble_notify(const uint8_t* /*buf*/, size_t /*len*/) {
    // Default: not wired. ble_proxy_chars.cpp overrides this.
    return false;
}
#endif // !UNIT_TEST
__attribute__((weak)) bool c2_proxy_is_authenticated() {
    if (ble_companion_is_authenticated) {
        return ble_companion_is_authenticated(millis());
    }
    return false;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
void c2_proxy_tunnel_init() {
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++) {
        release_slot(s_req_slots[i]);
        release_slot(s_rsp_slots[i]);
    }
    memset(s_lora_tx_queue, 0, sizeof(s_lora_tx_queue));
    memset(s_ble_tx_queue,  0, sizeof(s_ble_tx_queue));
    memset(&s_stats, 0, sizeof(s_stats));
}

C2ProxyStatus c2_proxy_handle_ble_write(const uint8_t* buf, size_t len) {
    if (!c2_proxy_is_authenticated()) {
        s_stats.dropped_not_auth++;
        return C2P_NOT_AUTH;
    }
    if (!buf || len < C2_PROXY_HDR_LEN) return C2P_BAD_HDR;

    // The phone may send the chunk with msg_type=0x24 already in buf[0], or
    // the BLE characteristic write callback may strip/replace it. We accept
    // both — if buf[0] != 0x24 we patch a temporary header copy.
    s_stats.ble_in_chunks++;
    int sidx = -1;
    bool completed = false;
    if (buf[0] != MSG_TYPE_C2_PROXY_REQ) {
        // Patch-and-handle: copy first few bytes, swap type, handle.
        uint8_t hdr[C2_PROXY_HDR_LEN + C2_PROXY_MAX_DATA];
        if (len > sizeof(hdr)) return C2P_BAD_HDR;
        memcpy(hdr, buf, len);
        hdr[0] = MSG_TYPE_C2_PROXY_REQ;
        C2ProxyStatus rc = handle_chunk_generic(s_req_slots, hdr, len,
                                                MSG_TYPE_C2_PROXY_REQ,
                                                &sidx, &completed);
        // Even partial OK: enqueue the LoRa-side TX of THIS exact chunk.
        if (rc == C2P_OK) {
            if (!enqueue_fragments_passthrough_pending(s_lora_tx_queue, hdr, len)) {
                // Fall through; the slot is still alive for future retries
                // when the BLE writer resends. Not ideal, but sane.
            }
        }
        if (completed) {
            s_stats.ble_in_envelopes++;
            // The generic chunk-passthrough already queued LoRa TX above; the
            // T2 receiver does its own reassembly. We can release our local
            // tracking slot.
            release_slot(s_req_slots[sidx]);
        }
        s_stats.in_flight_req = 0;
        for (int i = 0; i < C2_PROXY_CONCURRENT; i++)
            if (s_req_slots[i].in_use) s_stats.in_flight_req++;
        return rc;
    }

    C2ProxyStatus rc = handle_chunk_generic(s_req_slots, buf, len,
                                            MSG_TYPE_C2_PROXY_REQ,
                                            &sidx, &completed);
    if (rc == C2P_OK) {
        // Forward THIS chunk to T2 unchanged — T2 reassembles end-to-end.
        if (!enqueue_fragments_passthrough_pending(s_lora_tx_queue, buf, len)) {
            // Best-effort: not fatal; phone will retry envelope on timeout.
        }
    }
    if (completed) {
        s_stats.ble_in_envelopes++;
        release_slot(s_req_slots[sidx]);
    }
    s_stats.in_flight_req = 0;
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++)
        if (s_req_slots[i].in_use) s_stats.in_flight_req++;
    return rc;
}

// Helper used above — declared earlier via forward declaration.
static bool enqueue_fragments_passthrough_pending(WireFrame* ring,
                                                  const uint8_t* buf, size_t len) {
    if (len > C2_PROXY_HDR_LEN + C2_PROXY_MAX_DATA) return false;
    for (int i = 0; i < C2_PROXY_BLE_TX_QUEUE; i++) {
        if (!ring[i].used) {
            memset(ring[i].data, 0, sizeof(ring[i].data));
            memcpy(ring[i].data, buf, len);
            ring[i].len  = (uint16_t)len;
            ring[i].used = true;
            return true;
        }
    }
    return false;
}

C2ProxyStatus c2_proxy_handle_lora_rsp(const uint8_t* buf, size_t len) {
    s_stats.lora_in_chunks++;
    int sidx = -1;
    bool completed = false;
    C2ProxyStatus rc = handle_chunk_generic(s_rsp_slots, buf, len,
                                            MSG_TYPE_C2_PROXY_RSP,
                                            &sidx, &completed);
    if (rc != C2P_OK) return rc;

    // For each RSP chunk we receive, also queue a BLE notify of the SAME
    // chunk so the phone reassembles. The phone's reassembler is identical
    // to the T2-side reassembler (same 24-byte header pattern).
    if (!enqueue_fragments_passthrough_pending(s_ble_tx_queue, buf, len)) {
        // Could not queue — the phone will retry the envelope after timeout.
    }

    if (completed) {
        s_stats.lora_in_envelopes++;
        release_slot(s_rsp_slots[sidx]);
    }
    s_stats.in_flight_rsp = 0;
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++)
        if (s_rsp_slots[i].in_use) s_stats.in_flight_rsp++;
    return C2P_OK;
}

C2ProxyStatus c2_proxy_handle_lora_req(const uint8_t* buf, size_t len) {
    // Today: this T1 only forwards REQs from its own bonded phone. If a
    // neighboring T1 sees a REQ it isn't the originator of, it ignores it.
    // The slot rings on neighbors stay clean. This hook exists so future
    // multi-master mesh fan-out can plug in without an ABI bump.
    (void)buf; (void)len;
    return C2P_OK;
}

void c2_proxy_tunnel_tick() {
    // Drain LoRa-tx queue
    for (int i = 0; i < C2_PROXY_BLE_TX_QUEUE; i++) {
        if (s_lora_tx_queue[i].used) {
            if (c2_proxy_lora_tx(s_lora_tx_queue[i].data, s_lora_tx_queue[i].len)) {
                s_stats.lora_out_chunks++;
                s_lora_tx_queue[i].used = false;
            } else {
                // Keep frame for next tick; don't spam retries within one tick
                break;
            }
        }
    }
    // Drain BLE-notify queue
    for (int i = 0; i < C2_PROXY_BLE_TX_QUEUE; i++) {
        if (s_ble_tx_queue[i].used) {
            if (c2_proxy_ble_notify(s_ble_tx_queue[i].data, s_ble_tx_queue[i].len)) {
                s_stats.ble_out_chunks++;
                s_ble_tx_queue[i].used = false;
            } else {
                break;
            }
        }
    }
    // Reap stale slots (>30s with no progress)
    uint32_t now = millis();
    for (int i = 0; i < C2_PROXY_CONCURRENT; i++) {
        if (s_req_slots[i].in_use &&
            (now - s_req_slots[i].first_seen_ms) > 30000u) {
            release_slot(s_req_slots[i]);
        }
        if (s_rsp_slots[i].in_use &&
            (now - s_rsp_slots[i].first_seen_ms) > 30000u) {
            release_slot(s_rsp_slots[i]);
        }
    }
}

void c2_proxy_get_stats(C2ProxyStats* out) {
    if (!out) return;
    *out = s_stats;
}

#ifdef UNIT_TEST
void c2_proxy_tunnel_test_reset() { c2_proxy_tunnel_init(); }
void c2_proxy_tunnel_test_drain_lora_queue() {
    for (int i = 0; i < C2_PROXY_BLE_TX_QUEUE; i++) {
        if (s_lora_tx_queue[i].used) {
            (void)c2_proxy_lora_tx(s_lora_tx_queue[i].data, s_lora_tx_queue[i].len);
            s_lora_tx_queue[i].used = false;
            s_stats.lora_out_chunks++;
        }
    }
}
void c2_proxy_tunnel_test_drain_ble_queue() {
    for (int i = 0; i < C2_PROXY_BLE_TX_QUEUE; i++) {
        if (s_ble_tx_queue[i].used) {
            (void)c2_proxy_ble_notify(s_ble_tx_queue[i].data, s_ble_tx_queue[i].len);
            s_ble_tx_queue[i].used = false;
            s_stats.ble_out_chunks++;
        }
    }
}
#endif
