// ============================================================================
// mesh_routing.cpp — MRSH implementation
//
// Compile-out: wrap the bodies in #ifdef PDL_ENABLE_MESH_ROUTING so stock
// builds incur zero cost. All public symbols exist unconditionally.
// ============================================================================

#include "mesh_routing.h"

#include <cstring>
#include <cstdio>

// Use Arduino millis() when available; otherwise fall back to a monotonic
// counter for unit tests / host builds.
#if defined(ARDUINO) || defined(ESP_PLATFORM)
    #include <Arduino.h>
    static inline uint32_t mrsh_now_ms() { return millis(); }
#else
    #include <chrono>
    static inline uint32_t mrsh_now_ms() {
        using namespace std::chrono;
        static auto start = steady_clock::now();
        auto dt = duration_cast<milliseconds>(steady_clock::now() - start).count();
        return static_cast<uint32_t>(dt);
    }
#endif

// ── Forward declare the LoRa TX that main.cpp owns. We declare it weakly so
// unit tests can stub it. When built against real firmware, this resolves
// to the existing mesh_publish.cpp / LoRa TX path.
extern "C" __attribute__((weak))
int mesh_lora_send(const uint8_t* data, size_t len);

// Also declare a weak reference to the UAV relay so we can auto-deliver
// mesh-routed UAV payloads.
extern "C" __attribute__((weak))
void uav_relay_handle(const uint8_t* data, size_t len);

// ── Internal state ──────────────────────────────────────────────────────
static NeighborEntry s_neighbors[MRSH_NEIGHBOR_TABLE_SIZE];
static char s_own_id[17];
static uint8_t s_own_radios_bitmap = 0;
static uint8_t s_own_hop_to_gateway = 1;

static uint16_t s_next_packet_id = 1;
static uint32_t s_last_discover_ms = 0;
static uint32_t s_last_stats_ms = 0;
static uint32_t s_last_store_retry_ms = 0;

// Rolling 60 s stat counters (reset each stats cycle)
static uint32_t s_pkts_forwarded = 0;
static uint32_t s_pkts_dropped = 0;
static uint32_t s_pkts_delivered = 0;

// Store-and-forward queue
struct StoreSlot {
    uint8_t active;
    uint8_t pkt[MRSH_ROUTE_MAX_BYTES];
    size_t  pkt_len;
    uint32_t enqueued_ms;
    uint32_t last_retry_ms;
    char     dest_id[17];
};
static StoreSlot s_queue[MRSH_STORE_FWD_SLOTS];

// Delivery callback (default: uav_relay_handle if linked, else nothing)
static MrshDeliverCb s_deliver_cb = nullptr;

static void default_deliver(const uint8_t* payload, size_t len) {
    if (uav_relay_handle) uav_relay_handle(payload, len);
}

void mesh_set_deliver_callback(MrshDeliverCb cb) { s_deliver_cb = cb; }

// ────────────────────────────────────────────────────────────────────────
//  Hash + CRC
// ────────────────────────────────────────────────────────────────────────
void mesh_node_hash(const char* node_id, uint8_t out[6]) {
    // FNV-1a 64-bit → take lower 6 bytes
    uint64_t h = 1469598103934665603ULL;
    for (const char* p = node_id; p && *p; ++p) {
        h ^= static_cast<uint8_t>(*p);
        h *= 1099511628211ULL;
    }
    for (int i = 0; i < 6; i++) out[i] = static_cast<uint8_t>((h >> (i * 8)) & 0xFF);
}

uint32_t mesh_crc32(const uint8_t* data, size_t len) {
    // CRC-32/ISO-HDLC (poly 0xEDB88320), standard Ethernet CRC.
    uint32_t crc = 0xFFFFFFFFu;
    for (size_t i = 0; i < len; ++i) {
        crc ^= data[i];
        for (int b = 0; b < 8; b++) {
            crc = (crc >> 1) ^ (0xEDB88320u & (-(int32_t)(crc & 1)));
        }
    }
    return ~crc;
}

static inline void write_u16_le(uint8_t* p, uint16_t v) { p[0] = v; p[1] = v >> 8; }
static inline void write_u32_le(uint8_t* p, uint32_t v) {
    p[0] = v; p[1] = v >> 8; p[2] = v >> 16; p[3] = v >> 24;
}
static inline uint16_t read_u16_le(const uint8_t* p) { return p[0] | (p[1] << 8); }
static inline uint32_t read_u32_le(const uint8_t* p) {
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

// ────────────────────────────────────────────────────────────────────────
//  Framing
// ────────────────────────────────────────────────────────────────────────
size_t mesh_build_route_packet(
    uint8_t* out, size_t out_max,
    uint16_t packet_id,
    const char* src_id, const char* dest_id,
    uint8_t ttl, uint8_t flags, uint8_t next_hop_hint,
    const uint8_t* payload, size_t payload_len)
{
    if (!out || payload_len > MRSH_MAX_PAYLOAD) return 0;
    const size_t total = MRSH_ROUTE_HEADER_BYTES + payload_len + MRSH_ROUTE_FOOTER_BYTES;
    if (out_max < total) return 0;

    out[0] = MSG_TYPE_MESH_ROUTE;
    write_u16_le(out + 1, packet_id);
    mesh_node_hash(src_id, out + 3);            // 6 bytes
    if (dest_id) mesh_node_hash(dest_id, out + 9);
    else memset(out + 9, 0xFF, 6);              // broadcast
    out[15] = ttl;
    out[16] = flags;
    out[17] = next_hop_hint;
    out[18] = static_cast<uint8_t>(payload_len);
    if (payload && payload_len) memcpy(out + MRSH_ROUTE_HEADER_BYTES, payload, payload_len);
    // Wire layout: header (MRSH_ROUTE_HEADER_BYTES) + payload (payload_len) + CRC (4).
    const size_t crc_offset = MRSH_ROUTE_HEADER_BYTES + payload_len;
    const uint32_t crc = mesh_crc32(out, crc_offset);
    write_u32_le(out + crc_offset, crc);
    return crc_offset + 4;
}

bool mesh_parse_route_packet(const uint8_t* data, size_t len, MrshRouteHeader* out) {
    if (!data || !out) return false;
    if (len < MRSH_ROUTE_HEADER_BYTES + 4) return false;
    if (data[0] != MSG_TYPE_MESH_ROUTE) return false;

    out->packet_id = read_u16_le(data + 1);
    memcpy(out->src_hash,  data + 3, 6);
    memcpy(out->dest_hash, data + 9, 6);
    out->ttl = data[15];
    out->flags = data[16];
    out->next_hop_hint = data[17];
    out->payload_len = data[18];
    if (out->payload_len > MRSH_MAX_PAYLOAD) return false;

    const size_t crc_offset = MRSH_ROUTE_HEADER_BYTES + out->payload_len;
    if (crc_offset + 4 > len) return false;
    out->payload = data + MRSH_ROUTE_HEADER_BYTES;
    out->crc32_claimed  = read_u32_le(data + crc_offset);
    out->crc32_computed = mesh_crc32(data, crc_offset);
    return out->crc32_claimed == out->crc32_computed;
}

size_t mesh_build_ack_packet(
    uint8_t* out, size_t out_max,
    uint16_t packet_id,
    const char* acknowledger_id,
    uint8_t status, uint8_t hop_count_actual,
    int8_t rssi_at_dest)
{
    if (!out || out_max < MRSH_ACK_BYTES) return 0;
    out[0] = MSG_TYPE_MESH_ACK;
    write_u16_le(out + 1, packet_id);
    mesh_node_hash(acknowledger_id, out + 3);
    out[9] = status;
    out[10] = hop_count_actual;
    out[11] = static_cast<uint8_t>(rssi_at_dest);
    write_u32_le(out + 12, mesh_crc32(out, 12));
    return MRSH_ACK_BYTES;
}

size_t mesh_build_discover_packet(
    uint8_t* out, size_t out_max,
    const char* own_id, uint8_t radios_bitmap,
    uint8_t hop_to_gateway, uint8_t own_link_quality)
{
    if (!out || out_max < MRSH_DISCOVER_BYTES) return 0;
    memset(out, 0, MRSH_DISCOVER_BYTES);
    out[0] = MSG_TYPE_MESH_DISCOVER;
    mesh_node_hash(own_id, out + 1);            // 6 bytes
    out[7] = radios_bitmap;
    out[8] = hop_to_gateway;
    out[9] = own_link_quality;
    // 10..15: printable id prefix (truncated) for human debug
    size_t copy = strnlen(own_id, 16);
    if (copy > 6) copy = 6;
    memcpy(out + 10, own_id, copy);
    write_u32_le(out + MRSH_DISCOVER_BYTES - 4,
                 mesh_crc32(out, MRSH_DISCOVER_BYTES - 4));
    return MRSH_DISCOVER_BYTES;
}

// ────────────────────────────────────────────────────────────────────────
//  Neighbor-table ops
// ────────────────────────────────────────────────────────────────────────
static NeighborEntry* find_or_add_neighbor(const uint8_t hash6[6]) {
    char id_buf[17];
    snprintf(id_buf, sizeof(id_buf), "%02x%02x%02x%02x%02x%02x",
             hash6[0], hash6[1], hash6[2], hash6[3], hash6[4], hash6[5]);
    // find
    for (auto& n : s_neighbors) {
        if (n.active && strncmp(n.node_id, id_buf, 16) == 0) return &n;
    }
    // add — replace oldest / inactive
    NeighborEntry* victim = nullptr;
    uint32_t oldest = 0xFFFFFFFFu;
    for (auto& n : s_neighbors) {
        if (!n.active) { victim = &n; break; }
        if (n.last_seen_ms < oldest) { oldest = n.last_seen_ms; victim = &n; }
    }
    if (!victim) return nullptr;
    memset(victim, 0, sizeof(*victim));
    memcpy(victim->node_id, id_buf, 16);
    victim->node_id[16] = '\0';
    victim->active = 1;
    return victim;
}

static inline uint8_t rssi_to_quality(int8_t rssi, uint8_t snr) {
    // Matches rssiToLinkQuality in worker's topology.ts
    int r = rssi; if (r < -120) r = -120; if (r > -50) r = -50;
    float rn = (float)(r + 120) / 70.0f;
    float sb = (float)snr / 200.0f;
    if (sb > 0.1f) sb = 0.1f;
    float lq = (rn + sb) * 100.0f;
    if (lq < 0) lq = 0; if (lq > 100) lq = 100;
    return (uint8_t)lq;
}

static void neighbor_update(const uint8_t hash6[6], int8_t rssi, uint8_t snr, uint8_t radios, uint8_t hop_gw) {
    NeighborEntry* n = find_or_add_neighbor(hash6);
    if (!n) return;
    // EWMA
    if (n->last_seen_ms == 0) {
        n->rssi_avg = rssi;
        n->snr_avg = snr;
    } else {
        n->rssi_avg = (int8_t)(((int32_t)n->rssi_avg * 3 + rssi) / 4);
        n->snr_avg  = (uint8_t)(((uint32_t)n->snr_avg * 3 + snr) / 4);
    }
    n->last_seen_ms    = mrsh_now_ms();
    n->radios_bitmap   = radios;
    n->hop_to_gateway  = hop_gw;
    n->link_quality    = rssi_to_quality(n->rssi_avg, n->snr_avg);
}

// ────────────────────────────────────────────────────────────────────────
//  Store-and-forward
// ────────────────────────────────────────────────────────────────────────
static bool store_enqueue(const uint8_t* pkt, size_t len, const char* dest_id) {
    // BUG FIX: this previously memcpy'd `len` bytes into s.pkt (a fixed
    // MRSH_ROUTE_MAX_BYTES buffer) with no bounds check. Its only caller
    // (mesh_handle_route_packet's store-and-forward branch) passes through
    // the raw RX length with no upstream size cap on that particular path
    // -- unlike the sibling "forward" branch a few lines below, which DOES
    // check `len > sizeof(buf)` before its own memcpy. An oversized/
    // malformed packet reaching store-and-forward (no known next hop) would
    // overflow this static buffer. Bound-check here, at the point that owns
    // the fixed-size buffer, so it's correct regardless of caller.
    if (!pkt || len == 0 || len > MRSH_ROUTE_MAX_BYTES) return false;
    for (auto& s : s_queue) {
        if (!s.active) {
            s.active = 1;
            memcpy(s.pkt, pkt, len);
            s.pkt_len = len;
            s.enqueued_ms = mrsh_now_ms();
            s.last_retry_ms = s.enqueued_ms;
            strncpy(s.dest_id, dest_id ? dest_id : "", 16);
            s.dest_id[16] = '\0';
            return true;
        }
    }
    return false;
}

// ────────────────────────────────────────────────────────────────────────
//  Forwarding decision
// ────────────────────────────────────────────────────────────────────────
static bool hashes_equal(const uint8_t a[6], const uint8_t b[6]) {
    return memcmp(a, b, 6) == 0;
}

static NeighborEntry* pick_best_next_hop() {
    NeighborEntry* best = nullptr;
    int best_score = -1000;
    const uint32_t now = mrsh_now_ms();
    for (auto& n : s_neighbors) {
        if (!n.active) continue;
        if (now - n.last_seen_ms > 60000) continue;    // stale
        int score = (int)n.link_quality - (int)n.hop_to_gateway * 10;
        if (score > best_score) { best_score = score; best = &n; }
    }
    return best;
}

static NeighborEntry* find_neighbor_by_hash(const uint8_t hash6[6]) {
    char id_buf[17];
    snprintf(id_buf, sizeof(id_buf), "%02x%02x%02x%02x%02x%02x",
             hash6[0], hash6[1], hash6[2], hash6[3], hash6[4], hash6[5]);
    for (auto& n : s_neighbors) {
        if (n.active && strncmp(n.node_id, id_buf, 16) == 0) return &n;
    }
    return nullptr;
}

static int lora_tx(const uint8_t* buf, size_t len) {
    if (mesh_lora_send) return mesh_lora_send(buf, len);
    return -1;
}

// ────────────────────────────────────────────────────────────────────────
//  RX handlers
// ────────────────────────────────────────────────────────────────────────
void mesh_handle_route_packet(const uint8_t* data, size_t len) {
    MrshRouteHeader hdr;
    if (!mesh_parse_route_packet(data, len, &hdr)) { s_pkts_dropped++; return; }

    // Own-node-hash compare
    uint8_t own_hash[6];
    mesh_node_hash(s_own_id, own_hash);

    // Is it for us?
    if (hashes_equal(hdr.dest_hash, own_hash)) {
        s_pkts_delivered++;
        if (s_deliver_cb) s_deliver_cb(hdr.payload, hdr.payload_len);
        else default_deliver(hdr.payload, hdr.payload_len);

        if (hdr.flags & MRSH_FLAG_REQUIRE_ACK) {
            uint8_t ack[MRSH_ACK_BYTES];
            mesh_build_ack_packet(ack, sizeof(ack), hdr.packet_id,
                                  s_own_id, MRSH_ACK_DELIVERED, 1, -70);
            lora_tx(ack, sizeof(ack));
        }
        return;
    }

    // TTL check
    if (hdr.ttl == 0) {
        s_pkts_dropped++;
        uint8_t ack[MRSH_ACK_BYTES];
        mesh_build_ack_packet(ack, sizeof(ack), hdr.packet_id,
                              s_own_id, MRSH_ACK_DROPPED_TTL, 1, -120);
        lora_tx(ack, sizeof(ack));
        return;
    }

    // Route
    NeighborEntry* target = find_neighbor_by_hash(hdr.dest_hash);
    if (!target) target = pick_best_next_hop();

    if (!target) {
        // Store-and-forward?
        if (hdr.flags & MRSH_FLAG_STORE_AND_FORWARD) {
            if (store_enqueue(data, len, nullptr)) {
                uint8_t ack[MRSH_ACK_BYTES];
                mesh_build_ack_packet(ack, sizeof(ack), hdr.packet_id,
                                      s_own_id, MRSH_ACK_STORE_FORWARD_QUEUED, 1, -70);
                lora_tx(ack, sizeof(ack));
                return;
            }
        }
        // Broadcast if flagged
        if (hdr.flags & MRSH_FLAG_BROADCAST_ORPHAN) {
            uint8_t buf[MRSH_ROUTE_MAX_BYTES];
            memcpy(buf, data, len);
            buf[15] = hdr.ttl - 1;    // decrement TTL in copy
            // recompute CRC on new buffer
            const size_t crc_off = MRSH_ROUTE_HEADER_BYTES + hdr.payload_len;
            const uint32_t crc = mesh_crc32(buf, crc_off);
            write_u32_le(buf + crc_off, crc);
            lora_tx(buf, len);
            s_pkts_forwarded++;
        } else {
            s_pkts_dropped++;
            uint8_t ack[MRSH_ACK_BYTES];
            mesh_build_ack_packet(ack, sizeof(ack), hdr.packet_id,
                                  s_own_id, MRSH_ACK_NO_ROUTE, 1, -120);
            lora_tx(ack, sizeof(ack));
        }
        return;
    }

    // Forward: copy + decrement TTL + rewrite CRC
    uint8_t buf[MRSH_ROUTE_MAX_BYTES];
    if (len > sizeof(buf)) { s_pkts_dropped++; return; }
    memcpy(buf, data, len);
    buf[15] = hdr.ttl - 1;
    const size_t crc_off = MRSH_ROUTE_HEADER_BYTES + hdr.payload_len;
    const uint32_t crc = mesh_crc32(buf, crc_off);
    write_u32_le(buf + crc_off, crc);
    if (lora_tx(buf, len) >= 0) {
        s_pkts_forwarded++;
        if (hdr.flags & MRSH_FLAG_REQUIRE_ACK) {
            uint8_t ack[MRSH_ACK_BYTES];
            mesh_build_ack_packet(ack, sizeof(ack), hdr.packet_id,
                                  s_own_id, MRSH_ACK_FORWARDED, 1, -70);
            lora_tx(ack, sizeof(ack));
        }
    } else {
        s_pkts_dropped++;
    }
}

void mesh_handle_ack(const uint8_t* data, size_t len) {
    if (!data || len < MRSH_ACK_BYTES) return;
    if (data[0] != MSG_TYPE_MESH_ACK) return;
    const uint32_t crc_claim = read_u32_le(data + 12);
    const uint32_t crc_comp  = mesh_crc32(data, 12);
    if (crc_claim != crc_comp) return;
    // Upstream: use the packet_id to correlate with a pending TX waiting
    // for an ACK. For the firmware the record is purely bookkeeping — the
    // interesting side is the audit trail, which flows via the gateway to
    // the Worker's /api/mesh/route-ack endpoint. Nothing else to do here.
}

void mesh_handle_discover(const uint8_t* data, size_t len) {
    if (!data || len < MRSH_DISCOVER_BYTES) return;
    if (data[0] != MSG_TYPE_MESH_DISCOVER) return;
    const uint32_t crc_claim = read_u32_le(data + MRSH_DISCOVER_BYTES - 4);
    const uint32_t crc_comp  = mesh_crc32(data, MRSH_DISCOVER_BYTES - 4);
    if (crc_claim != crc_comp) return;

    uint8_t src_hash[6];
    memcpy(src_hash, data + 1, 6);
    const uint8_t radios   = data[7];
    const uint8_t hop_gw   = data[8];
    // We don't have per-frame RSSI here (main dispatcher consumed it).
    // Use last-known link RSSI from prior discover frames; default -75 dBm.
    neighbor_update(src_hash, -75, 8, radios, hop_gw);
}

// ────────────────────────────────────────────────────────────────────────
//  Init + periodic tick
// ────────────────────────────────────────────────────────────────────────
void mesh_routing_init(const char* own_node_id, uint8_t own_radios_bitmap) {
    memset(s_neighbors, 0, sizeof(s_neighbors));
    memset(s_queue, 0, sizeof(s_queue));
    strncpy(s_own_id, own_node_id ? own_node_id : "unknown", 16);
    s_own_id[16] = '\0';
    s_own_radios_bitmap = own_radios_bitmap;
    s_last_discover_ms = 0;
    s_last_stats_ms = 0;
    s_last_store_retry_ms = 0;
    s_next_packet_id = 1;
    s_pkts_forwarded = 0;
    s_pkts_dropped = 0;
    s_pkts_delivered = 0;
}

static void do_discover_broadcast() {
    uint8_t buf[MRSH_DISCOVER_BYTES];
    // Compute own_link_quality as avg of active neighbors
    uint32_t sum = 0; uint32_t cnt = 0;
    for (auto& n : s_neighbors) if (n.active) { sum += n.link_quality; cnt++; }
    uint8_t avg = cnt ? (uint8_t)(sum / cnt) : 50;
    mesh_build_discover_packet(buf, sizeof(buf),
                               s_own_id, s_own_radios_bitmap,
                               s_own_hop_to_gateway, avg);
    lora_tx(buf, sizeof(buf));
}

static void retry_store_forward() {
    const uint32_t now = mrsh_now_ms();
    for (auto& s : s_queue) {
        if (!s.active) continue;
        if (now - s.enqueued_ms > MRSH_STORE_FWD_MAX_MS) {
            s.active = 0; s_pkts_dropped++; continue;
        }
        if (now - s.last_retry_ms < 30000) continue;
        s.last_retry_ms = now;
        // Re-examine — if a neighbor now exists, forward.
        MrshRouteHeader hdr;
        if (!mesh_parse_route_packet(s.pkt, s.pkt_len, &hdr)) { s.active = 0; continue; }
        NeighborEntry* t = find_neighbor_by_hash(hdr.dest_hash);
        if (!t) t = pick_best_next_hop();
        if (t) {
            lora_tx(s.pkt, s.pkt_len);
            s.active = 0;
            s_pkts_forwarded++;
        }
    }
}

static void roll_stats() {
    // Reset 60 s counters after a stats window elapses. Snapshot is
    // readable via mesh_get_stats() between rolls.
    s_pkts_forwarded = 0;
    s_pkts_dropped = 0;
    s_pkts_delivered = 0;
}

void mesh_routing_tick() {
    const uint32_t now = mrsh_now_ms();
    if (now - s_last_discover_ms >= 15000) { s_last_discover_ms = now; do_discover_broadcast(); }
    if (now - s_last_store_retry_ms >= 30000) { s_last_store_retry_ms = now; retry_store_forward(); }
    if (now - s_last_stats_ms >= 60000) { s_last_stats_ms = now; roll_stats(); }
}

// ────────────────────────────────────────────────────────────────────────
//  Public getters + sender
// ────────────────────────────────────────────────────────────────────────
void mesh_get_stats(MrshStats* out) {
    if (!out) return;
    memset(out, 0, sizeof(*out));
    uint32_t sum = 0, cnt = 0;
    uint32_t depth = 0;
    for (auto& n : s_neighbors) if (n.active) { sum += n.link_quality; cnt++; }
    for (auto& s : s_queue) if (s.active) depth++;
    out->pkts_forwarded_last_60s = s_pkts_forwarded;
    out->pkts_dropped_last_60s   = s_pkts_dropped;
    out->pkts_delivered_last_60s = s_pkts_delivered;
    out->active_neighbor_count   = (uint8_t)cnt;
    out->own_link_quality_avg    = cnt ? (uint8_t)(sum / cnt) : 0;
    out->store_forward_depth     = depth;
    out->updated_at_ms           = mrsh_now_ms();
}

size_t mesh_get_neighbors(NeighborEntry* out, size_t max) {
    size_t written = 0;
    for (auto& n : s_neighbors) {
        if (!n.active) continue;
        if (written >= max) break;
        out[written++] = n;
    }
    return written;
}

uint16_t mesh_send(const char* dest_node_id,
                   const uint8_t* payload, size_t payload_len,
                   uint8_t flags)
{
    if (payload_len > MRSH_MAX_PAYLOAD) return 0;
    const uint16_t pkt_id = s_next_packet_id++;
    if (s_next_packet_id == 0) s_next_packet_id = 1;

    uint8_t buf[MRSH_ROUTE_MAX_BYTES];
    const size_t wire_len = mesh_build_route_packet(
        buf, sizeof(buf),
        pkt_id, s_own_id, dest_node_id,
        MRSH_TTL_DEFAULT, flags, /*next_hop_hint=*/0,
        payload, payload_len);
    if (wire_len == 0) return 0;

    // Try immediate TX
    if (lora_tx(buf, wire_len) < 0) {
        // Queue for store-and-forward on failure (if caller enabled)
        if (flags & MRSH_FLAG_STORE_AND_FORWARD) {
            store_enqueue(buf, wire_len, dest_node_id);
        } else {
            s_pkts_dropped++;
            return 0;
        }
    }
    return pkt_id;
}
