meshtastic_firmware/src/mesh/NextHopRouter.h

#pragma once

#include "FloodingRouter.h"
#include <optional>
#include <unordered_map>

/**
 * An identifier for a globally unique message - a pair of the sending nodenum and the packet id assigned
 * to that message
 */
struct GlobalPacketId {
    NodeNum node;
    PacketId id;

    bool operator==(const GlobalPacketId &p) const { return node == p.node && id == p.id; }

    explicit GlobalPacketId(const meshtastic_MeshPacket *p)
    {
        node = getFrom(p);
        id = p->id;
    }

    GlobalPacketId(NodeNum _from, PacketId _id)
    {
        node = _from;
        id = _id;
    }
};

/**
 * A packet queued for retransmission
 */
struct PendingPacket {
    meshtastic_MeshPacket *packet;

    /** The next time we should try to retransmit this packet */
    uint32_t nextTxMsec = 0;

    /** Starts at NUM_RETRANSMISSIONS -1 and counts down.  Once zero it will be removed from the list */
    uint8_t numRetransmissions = 0;

    PendingPacket() {}
    explicit PendingPacket(meshtastic_MeshPacket *p, uint8_t numRetransmissions);
};

/**
 * RAM-only per-destination route health. Tracks how fresh a learned next_hop is and how many
 * consecutive directed deliveries to it have failed, so getNextHop() can proactively decay a stale or
 * repeatedly-failing route back to flooding instead of trusting a dead hop on the next (and on dense
 * meshes, slowest) attempt. Not persisted: the learned next_hop itself lives in NodeInfoLite; this is
 * just freshness/failure metadata.
 */
struct RouteHealth {
    NodeNum dest = 0;                             ///< destination this record describes; 0 == empty slot
    uint32_t learnedAtMsec = 0;                   ///< millis() when next_hop was last (re)learned (rollover-aware)
    uint8_t consecutiveFailures = 0;              ///< directed deliveries to `dest` that went un-ACKed
    uint8_t lastNextHop = NO_NEXT_HOP_PREFERENCE; ///< the relay byte this health refers to
};

// M4 (optional, off by default): when a route is not proven healthy, fall back to flooding one retry
// earlier instead of spending a second directed attempt. Trades airtime for recovery latency on dense
// meshes; leaves the sparse-mesh happy path (fresh, verified routes) unchanged. Measure on the
// simulator before enabling broadly.
#ifndef NEXTHOP_EARLY_FLOOD_ON_UNVERIFIED
#define NEXTHOP_EARLY_FLOOD_ON_UNVERIFIED 0
#endif

class GlobalPacketIdHashFunction
{
  public:
    size_t operator()(const GlobalPacketId &p) const { return (std::hash<NodeNum>()(p.node)) ^ (std::hash<PacketId>()(p.id)); }
};

/*
  Router for direct messages, which only relays if it is the next hop for a packet. The next hop is set by the current
  relayer of a packet, which bases this on information from a previous successful delivery to the destination via flooding.
  Namely, in the PacketHistory, we keep track of (up to 3) relayers of a packet. When the ACK is delivered back to us via a node
  that also relayed the original packet, we use that node as next hop for the destination from then on. This makes sure that only
  when there’s a two-way connection, we assign a next hop. Both the ReliableRouter and NextHopRouter will do retransmissions (the
  NextHopRouter only 1 time). For the final retry, if no one actually relayed the packet, it will reset the next hop in order to
  fall back to the FloodingRouter again. Note that thus also intermediate hops will do a single retransmission if the intended
  next-hop didn’t relay, in order to fix changes in the middle of the route.
*/
class NextHopRouter : public FloodingRouter
{
  public:
    /**
     * Constructor
     *
     */
    NextHopRouter();

    /**
     * Send a packet
     * @return an error code
     */
    virtual ErrorCode send(meshtastic_MeshPacket *p) override;

    /** Do our retransmission handling */
    virtual int32_t runOnce() override
    {
        // Note: We must doRetransmissions FIRST, because it might queue up work for the base class runOnce implementation
        doRetransmissions();

        int32_t r = FloodingRouter::runOnce();

        // Also after calling runOnce there might be new packets to retransmit
        auto d = doRetransmissions();
        return min(d, r);
    }

    // The number of retransmissions intermediate nodes will do (actually 1 less than this)
    constexpr static uint8_t NUM_INTERMEDIATE_RETX = 2;
    // The number of retransmissions the original sender will do
    constexpr static uint8_t NUM_RELIABLE_RETX = 3;

    // M3: bounded RAM route-health table (reuse-oldest eviction, like PacketHistory)
    constexpr static uint8_t ROUTE_HEALTH_MAX = 32;              // ~12B/slot -> ~384B
    constexpr static uint32_t ROUTE_TTL_MSEC = 30UL * 60 * 1000; // re-discover a route unconfirmed for 30 min
    constexpr static uint8_t ROUTE_FAILURE_THRESHOLD = 3;        // consecutive un-ACKed directed deliveries -> dead

  protected:
    /**
     * Pending retransmissions
     */
    std::unordered_map<GlobalPacketId, PendingPacket, GlobalPacketIdHashFunction> pending;

    /**
     * Per-destination route health (M3). Bounded array, reuse-oldest eviction. RAM-only.
     */
    RouteHealth routeHealth[ROUTE_HEALTH_MAX] = {};

    /**
     * Should this incoming filter be dropped?
     *
     * Called immediately on reception, before any further processing.
     * @return true to abandon the packet
     */
    virtual bool shouldFilterReceived(const meshtastic_MeshPacket *p) override;
    bool relayOpaquePacket(const meshtastic_MeshPacket *p) override;

    /**
     * Look for packets we need to relay
     */
    virtual void sniffReceived(const meshtastic_MeshPacket *p, const meshtastic_Routing *c) override;

    /**
     * Try to find the pending packet record for this ID (or NULL if not found)
     */
    PendingPacket *findPendingPacket(NodeNum from, PacketId id) { return findPendingPacket(GlobalPacketId(from, id)); }
    PendingPacket *findPendingPacket(GlobalPacketId p);

    /**
     * Add p to the list of packets to retransmit occasionally.  We will free it once we stop retransmitting.
     */
    PendingPacket *startRetransmission(meshtastic_MeshPacket *p, uint8_t numReTx = NUM_INTERMEDIATE_RETX);

    // Return true if we're allowed to cancel a packet in the txQueue (so we may never transmit it even once)
    bool roleAllowsCancelingFromTxQueue(const meshtastic_MeshPacket *p);

    /**
     * Stop any retransmissions we are doing of the specified node/packet ID pair
     *
     * @return true if we found and removed a transmission with this ID
     */
    bool stopRetransmission(NodeNum from, PacketId id);
    bool stopRetransmission(GlobalPacketId p);

    /**
     * Do any retransmissions that are scheduled (FIXME - for the time being called from loop)
     *
     * @return the number of msecs until our next retransmission or MAXINT if none scheduled
     */
    int32_t doRetransmissions();

    void setNextTx(PendingPacket *pending);

    // --- M3 route-health helpers (RAM-only). Protected so ReliableRouter (a subclass) can record
    // delivery success, and so the unit-test shim can reach them via `using`. All take `now` where
    // time matters so the decay logic is pure and testable without a clock mock. ---

    /// @return the health record for `dest`, or nullptr if we hold none.
    RouteHealth *findRouteHealth(NodeNum dest);
    /// @return an existing record for `dest`, else a freshly claimed slot (reuse-oldest on overflow).
    RouteHealth *getOrAllocRouteHealth(NodeNum dest, uint32_t now);
    /// Record that we (re)learned `nextHop` for `dest`. Resets the failure count only when the hop
    /// changed (so a flapping reverse-path re-learn of the same dead hop still ages out).
    void noteRouteLearned(NodeNum dest, uint8_t nextHop, uint32_t now);
    /// Record an end-to-end delivery success to `dest` (clears failures, refreshes freshness).
    void noteRouteSuccess(NodeNum dest, uint32_t now);
    /// Record that a directed delivery to `dest` went un-ACKed (no-op if we hold no record).
    void noteRouteFailure(NodeNum dest);
    /// @return true if the route is too old (TTL) or has failed too many times in a row.
    bool isRouteStale(const RouteHealth &h, uint32_t now) const;
    /// Forget any health record for `dest`.
    void clearRouteHealth(NodeNum dest);

#ifdef PIO_UNIT_TESTING
  public: // expose getNextHop to the test shim without widening production visibility
#else
  private:
#endif
    /**
     * Get the next hop for a destination, given the relay node
     * @return the node number of the next hop, 0 if no preference (fallback to FloodingRouter)
     */
    std::optional<uint8_t> getNextHop(NodeNum to, uint8_t relay_node);

#ifdef PIO_UNIT_TESTING
  public: // expose perhapsRebroadcast to the test shim
#else
  private:
#endif
    /** Check if we should be rebroadcasting this packet if so, do so.
     *  @return true if we did rebroadcast */
    bool perhapsRebroadcast(const meshtastic_MeshPacket *p) override;
};