#ifndef NETWORK_NETWORK_NIC_H
#define NETWORK_NETWORK_NIC_H

#include <cstdint>
#include <deque>
#include <unordered_map>
#include <vector>

#include "core/time.h"
#include "core/types.h"
#include "core/rng.h"
#include "network/network_node.h"
#include "network/packet.h"
#include "network/switch/switch_buffer.h"
#include "network/nic/nic_telemetry.h"
#include "network/nic/congestion_control.h"
#include "network/nic/load_balance.h"

namespace dofs {

class Host;

struct NicSchedulingWeights {
        // Bytes of budget per round (refilled when all budgets deplete).
        Bytes control  { 64 * 1024 };
        Bytes retrans  { 256 * 1024 };
        Bytes mice     { 1024 * 1024 };
        Bytes elephant { 1024 * 1024 };
};

class NetworkNic final : public NetworkNode {
public:
        NetworkNic(Simulator *const sim,
                   NodeId id,
                   uint16_t total_ports,
                   SwitchBuffer *const buf,
                   Time nic_latency,
                   Bytes mice_elephant_threshold,
                   PacketSeq ooo_threshold,
                   CCType cc_type,
                   LBType lb_type,
                   Bytes cc_init_cwnd,
                   Bytes cc_max_cwnd,
                   const NicSchedulingWeights& schedw) noexcept;

        virtual void recv_pkt(Packet &pkt, PortId ingress) override;

        void attach_host(Host* host) noexcept;
        void detach_host(Host* host) noexcept;

        // Host API (TX):
        void send_flow(NodeId dst, Bytes size,
                       FlowPriority desired = FlowPriority::AUTO);
        void send_flow(PacketGroups group_mask, Bytes size,
                       FlowPriority desired = FlowPriority::AUTO);

        void set_status(NodeStatus s, Time new_latency = Time(0)) noexcept;

        const NicTelemetry &telemetry() const noexcept {
                return _telemetry;
        }

        void set_port_blacklisted(PortId port, bool blacklisted) noexcept;
        bool is_port_blacklisted(PortId port) const noexcept;

private:
        enum class QClass : uint8_t { CONTROL = 0, RETRANS = 1, MICE = 2, ELEPHANT = 3 };

        struct PortQueues {
                std::deque<Packet> control;
                std::deque<Packet> retrans;
                std::deque<Packet> mice;
                std::deque<Packet> elephant;

                Bytes b_control  {0};
                Bytes b_retrans  {0};
                Bytes b_mice     {0};
                Bytes b_elephant {0};

                bool scheduled {false};
        };

        void schedule_port_if_needed(PortId port) noexcept;
        void port_drain_task(PortId port) noexcept;

        bool pick_next_qclass(const PortQueues& pq, QClass& out_cls) const noexcept;

        bool try_send_one(PortId port, QClass cls) noexcept;

        void enqueue_packet(PortId port, QClass cls, Packet pkt) noexcept;

        void schedule_ack(const Packet& rx_data, PortId ingress) noexcept;
        void schedule_nack(NodeId peer, FlowId flow, PacketSeq missing_seq,
                           FlowPriority prio, PortId ingress) noexcept;
        void schedule_trim_back_response(const Packet& trim, PortId ingress) noexcept;

        bool is_elephant(Bytes sz) const noexcept;

        static inline std::uint64_t txkey(FlowId f, PacketSeq s) noexcept;

        void record_tx_timestamp(const Packet& pkt) noexcept;
        bool lookup_rtt_and_erase(const Packet& ack_like, Time now,
                                  Time& out_rtt) noexcept;

        PortId pick_src_port_for_flow(NodeId dst) noexcept;

        void start_tx_unicast(NodeId dst, Bytes size, FlowPriority prio);
        void start_tx_multicast(PacketGroups gmask, Bytes size, FlowPriority prio);

        Packet make_data_packet(NodeId dst, PortId out_port, FlowPriority prio,
                                FlowId fid, PacketSeq seq, Bytes payload) noexcept;
        Packet make_ack_packet(const Packet& rx_data, PortId ingress) noexcept;
        Packet make_nack_packet(NodeId peer, FlowId flow, PacketSeq miss,
                                FlowPriority prio,
                                PortId ingress) noexcept;
        Packet make_trim_back_response(const Packet& trim, PortId ingress) noexcept;

private:
        Host                  *_host {nullptr};
        SwitchBuffer          *_buf  {nullptr};

        const uint16_t         _port_cnt;
        Time                   _nic_latency_init;
        Time                   _nic_latency_cur;

        const Bytes            _mice_thresh;
        const PacketSeq        _ooo_thresh;

        CCType                 _cc_type;
        LBType                 _lb_type;

        NicSchedulingWeights   _schedw;

        NicTelemetry           _telemetry;
        Rng                    _rng;

        std::vector<PortQueues> _ports;

        std::vector<bool>       _port_blacklisted;
        PortId                  _rr_cursor {0};

        std::unique_ptr<CongestionControl> _cc;
        std::unique_ptr<LoadBalance>       _lb;

        // Outstanding bytes per flow (TX) for cwnd gating
        std::unordered_map<FlowId, Bytes> _tx_outstanding;

        // TX timestamps for RTT: (flow, seq) -> send_time
        std::unordered_map<std::uint64_t, Time> _tx_sent_at;
};

} // namespace dofs

#endif // NETWORK_NETWORK_NIC_H
