#ifndef NETWORK_SWITCH_ECN_ENGINE_H
#define NETWORK_SWITCH_ECN_ENGINE_H

#include <cstdint>
#include "core/types.h"
#include "network/packet.h"

namespace dofs {

class SwitchBuffer;

class ECNEngine {
public:
        virtual ~ECNEngine() = default;

        // Decide ECN mark / header-trim for a packet about to enter 'buf'.
        // Returns the same packet reference, possibly modified in-place.
        virtual Packet &process_packet(Packet &pkt, SwitchBuffer* buf) noexcept = 0;

protected:
        // Helpers for derived engines
        // Now defined directly in the header
        static inline void maybe_mark_ecn(Packet &pkt) noexcept {
                if (pkt.is_ecn_enabled())
                        pkt.set_ecn_marked(true);
        }

        // Trim payload to zero and set protocol accordingly.
        // If back_to_sender, flip src/dst like an RX NIC reflection.
        static inline void header_trim(Packet &pkt, bool back_to_sender) noexcept {
                pkt.set_payload_size(Bytes(0));

                if (back_to_sender) {
                        const NodeId s_n = pkt.src_node();
                        const PortId s_p = pkt.src_port();
                        const NodeId d_n = pkt.dst_node();
                        const PortId d_p = pkt.dst_port();
                        pkt.set_src_node(d_n);
                        pkt.set_src_port(d_p);
                        pkt.set_dst_node(s_n);
                        pkt.set_dst_port(s_p);
                        pkt.set_protocol(PacketProtocol::HEADER_TRIM_BACK);
                } else {
                        pkt.set_protocol(PacketProtocol::HEADER_TRIM);
                }
        }
};

} // namespace dofs

#endif // NETWORK_SWITCH_ECN_ENGINE_
