#include "network/network_nic.h"

#include <algorithm>
#include <utility>
#include "core/host.h"
#include "core/error.h"

namespace dofs {

NetworkNic::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
        : NetworkNode(sim, id, NodeType::NIC),
          _buf(buf),
          _port_cnt(total_ports),
          _nic_latency_init(nic_latency),
          _nic_latency_cur(nic_latency),
          _mice_thresh(mice_elephant_threshold),
          _ooo_thresh(ooo_threshold),
          _cc_type(cc_type),
          _lb_type(lb_type),
          _schedw(schedw),
          _ports(total_ports) {

        Node::set_status(NodeStatus::OK);

        _port_blacklisted.assign(_port_cnt, false);

        if (!_buf) {
                log_error("ERROR", "NetworkNic constructed with null SwitchBuffer");
        }

        switch (_cc_type) {
        case CCType::DCQCN:
                _cc = std::make_unique<DCQCN>(cc_init_cwnd, cc_max_cwnd);
                break;

        case CCType::NSCC:
                _cc = std::make_unique<NSCC>(cc_init_cwnd, cc_max_cwnd);
                break;

        default:
                _cc = std::make_unique<DCQCN>(cc_init_cwnd, cc_max_cwnd);
                break;
        }

        switch (_lb_type) {
        case LBType::RANDOM_PACKET_SPRAYING:
                _lb = std::make_unique<LBRandomPacketSpraying>(&_rng);
                break;

        default:
                _lb = std::make_unique<LBRandomPacketSpraying>(&_rng);
                break;
        }

        for (auto& pq : _ports) {
                pq.b_control  = _schedw.control;
                pq.b_retrans  = _schedw.retrans;
                pq.b_mice     = _schedw.mice;
                pq.b_elephant = _schedw.elephant;
        }
}

void NetworkNic::attach_host(Host* host) noexcept {
        _host = host;

        if (_host)
                _host->attach_nic(this);
}

void NetworkNic::detach_host(Host* host) noexcept {
        if (_host && _host != host) {
                log_error("ERROR", "NetworkNic::detach_host pointer mismatch");
                return;
        }

        if (_host)
                _host->detach_nic(this);

        _host = nullptr;
}

void NetworkNic::set_status(NodeStatus s, Time new_latency) noexcept {
        switch (s) {
        case NodeStatus::OK:
                Node::set_status(NodeStatus::OK);
                _nic_latency_cur = _nic_latency_init;
                return;

        case NodeStatus::THROTTLING:
                if (new_latency < _nic_latency_init) {
                        log_error("ERROR", "NetworkNic::set_status THROTTLING new latency < initial");
                        return;
                }

                Node::set_status(NodeStatus::THROTTLING);
                _nic_latency_cur = new_latency;
                return;

        case NodeStatus::DOWN:
        case NodeStatus::BOOTING:
                log_error("ERROR", "NetworkNic::set_status ignored for DOWN/BOOTING");
                return;
        }
}

bool NetworkNic::is_elephant(Bytes sz) const noexcept {
        return sz >= _mice_thresh;
}

std::uint64_t NetworkNic::txkey(FlowId f, PacketSeq s) noexcept {
        return (static_cast<std::uint64_t>(static_cast<std::uint32_t>(f)) << 16) |
               static_cast<std::uint64_t>(static_cast<std::uint16_t>(s));
}

void NetworkNic::record_tx_timestamp(const Packet& pkt) noexcept {
        if (!_sim)
                return;

        const auto key = txkey(pkt.flow_id(), pkt.seq());
        _tx_sent_at[key] = _sim->now();
}

bool NetworkNic::lookup_rtt_and_erase(const Packet& ack_like, Time now,
                                      Time& out_rtt) noexcept {
        const auto key = txkey(ack_like.flow_id(), ack_like.seq());
        auto it = _tx_sent_at.find(key);

        if (it == _tx_sent_at.end())
                return false;

        out_rtt = now.unsafe_sub(it->second);
        _tx_sent_at.erase(it);
        return true;
}

void NetworkNic::recv_pkt(Packet &pkt_in, PortId ingress) {
        if (status() == NodeStatus::DOWN || status() == NodeStatus::BOOTING) {
                return;
        }

        Packet pkt = pkt_in;
        const Time now = _sim ? _sim->now() : Time(0);

        switch (pkt.protocol()) {
        case PacketProtocol::DATA: {
                _telemetry.rx_data++;

                // Auto-ACK after NIC latency
                schedule_ack(pkt, ingress);

                // For OOO detection & NACK: basic placeholder based on seq gap.
                // (A real impl would maintain RxFlowState with expected_seq.)
                // Here we keep it minimal and do not NACK unless explicitly enabled later.

                // Flow completion: if EOF, notify host with total size from payload (accumulation would live in RxFlowState).
                if (pkt.is_eof() && _host) {
                        _host->recv_flow(pkt.src_node(),
                                         pkt.flow_id(), pkt.priority(), pkt.payload_size());
                }

                break;
        }

        case PacketProtocol::ACK: {
                _telemetry.rx_acks++;

                // RTT measurement and CC update
                Time rtt;

                if (lookup_rtt_and_erase(pkt, now, rtt)) {
                        if (_cc) {
                                _cc->update(pkt, rtt);
                                _telemetry.cc_updates++;
                        }
                } else {
                        if (_cc) {
                                _cc->update(pkt, Time(0));
                                _telemetry.cc_updates++;
                        }
                }

                if (_lb) {
                        _lb->update(pkt);
                        _telemetry.lb_updates++;
                }

                if (_host)
                        _host->recv_frame(pkt);

                break;
        }

        case PacketProtocol::NACK: {
                _telemetry.rx_nacks++;

                if (_cc) {
                        _cc->update(pkt, Time(0));
                        _telemetry.cc_updates++;
                }

                if (_lb) {
                        _lb->update(pkt);
                        _telemetry.lb_updates++;
                }

                if (_host)
                        _host->recv_frame(pkt);

                // Retransmission path would re-enqueue missing seq as CONTROL/RETRANS.
                break;
        }

        case PacketProtocol::HEADER_TRIM: {
                _telemetry.rx_trims++;
                schedule_trim_back_response(pkt, ingress);

                if (_cc) {
                        _cc->update(pkt, Time(0));
                        _telemetry.cc_updates++;
                }

                if (_lb) {
                        _lb->update(pkt);
                        _telemetry.lb_updates++;
                }

                if (_host)
                        _host->recv_frame(pkt);

                break;
        }

        case PacketProtocol::HEADER_TRIM_BACK: {
                _telemetry.rx_trims++;

                if (_cc) {
                        _cc->update(pkt, Time(0));
                        _telemetry.cc_updates++;
                }

                if (_lb) {
                        _lb->update(pkt);
                        _telemetry.lb_updates++;
                }

                if (_host)
                        _host->recv_frame(pkt);

                break;
        }

        default:
                // Unknown proto: ignore
                break;
        }
}

void NetworkNic::send_flow(NodeId dst, Bytes size, FlowPriority desired) {
        const bool ele = is_elephant(size);
        const FlowPriority prio = (desired == FlowPriority::AUTO)
                                  ? (ele ? FlowPriority::ELEPHANT : FlowPriority::MICE)
                                  : desired;
        start_tx_unicast(dst, size, prio);
}

void NetworkNic::send_flow(PacketGroups gmask, Bytes size,
                           FlowPriority desired) {
        const bool ele = is_elephant(size);
        const FlowPriority prio = (desired == FlowPriority::AUTO)
                                  ? (ele ? FlowPriority::ELEPHANT : FlowPriority::MICE)
                                  : desired;
        start_tx_multicast(gmask, size, prio);
}

void NetworkNic::set_port_blacklisted(PortId port, bool blacklisted) noexcept {
        if (port >= _port_cnt) {
                log_error("ERROR", "NetworkNic::set_port_blacklisted invalid port");
                return;
        }

        if (port >= _port_blacklisted.size())
                _port_blacklisted.resize(_port_cnt, false);

        _port_blacklisted[port] = blacklisted;
}

bool NetworkNic::is_port_blacklisted(PortId port) const noexcept {
        if (port >= _port_cnt || port >= _port_blacklisted.size())
                return false;

        return _port_blacklisted[port];
}

PortId NetworkNic::pick_src_port_for_flow(NodeId dst) noexcept {
        (void)dst;

        if (_port_cnt == 0)
                return static_cast<PortId>(0);

        // Simple RR scheme
        const uint16_t start = _rr_cursor % _port_cnt;

        for (uint16_t i = 0; i < _port_cnt; ++i) {
                const PortId cand = static_cast<PortId>((start + i) % _port_cnt);

                if (cand < _port_blacklisted.size() && !_port_blacklisted[cand]) {
                        _rr_cursor = static_cast<PortId>((cand + 1) % _port_cnt);
                        return cand;
                }
        }

        log_error("WARN",
                  "NetworkNic: all ports blacklisted; choosing RR port despite blacklist");
        const PortId fallback = static_cast<PortId>(start);
        _rr_cursor = static_cast<PortId>((fallback + 1) % _port_cnt);
        return fallback;
}

Packet NetworkNic::make_data_packet(NodeId dst, PortId out_port,
                                    FlowPriority prio,
                                    FlowId fid, PacketSeq seq, Bytes payload) noexcept {
        Packet p(this->id(), out_port,
                 dst, static_cast<PortId>(0),
                 PacketProtocol::DATA,
                 prio, seq, fid, 0, 0, payload);

        if (_lb) {
                const uint16_t e = _lb->get_entropy(p);
                p.set_entropy(e);
        }

        return p;
}

Packet NetworkNic::make_ack_packet(const Packet& rx_data,
                                   PortId ingress) noexcept {
        Packet ack(rx_data.dst_node(), rx_data.dst_port(),
                   rx_data.src_node(), rx_data.src_port(),
                   PacketProtocol::ACK,
                   FlowPriority::CTRL,
                   rx_data.seq(), rx_data.flow_id(),
                   rx_data.entropy(), 0, Bytes(0));
        (void)ingress;
        return ack;
}

Packet NetworkNic::make_nack_packet(NodeId peer, FlowId flow, PacketSeq miss,
                                    FlowPriority prio,
                                    PortId ingress) noexcept {
        Packet nack(this->id(), ingress,
                    peer, static_cast<PortId>(0),
                    PacketProtocol::NACK,
                    prio, miss, flow, 0, 0, Bytes(0));
        return nack;
}

Packet NetworkNic::make_trim_back_response(const Packet& trim,
                PortId ingress) noexcept {
        Packet tb(trim.dst_node(), trim.dst_port(),
                  trim.src_node(), trim.src_port(),
                  PacketProtocol::HEADER_TRIM_BACK,
                  FlowPriority::CTRL,
                  trim.seq(), trim.flow_id(),
                  trim.entropy(), 0, Bytes(0));
        (void)ingress;
        return tb;
}

void NetworkNic::schedule_ack(const Packet& rx_data, PortId ingress) noexcept {
        if (!_sim)
                return;

        const Packet ack = make_ack_packet(rx_data, ingress);
        const PortId egress = static_cast<PortId> (ingress);

        _sim->schedule_after(_nic_latency_cur,
        [this, ack, egress]() mutable {
                enqueue_packet(egress, QClass::CONTROL, ack);
        });
        _telemetry.tx_acks++;
}

void NetworkNic::schedule_nack(NodeId peer, FlowId flow, PacketSeq missing_seq,
                               FlowPriority prio, PortId ingress) noexcept {
        if (!_sim)
                return;

        const Packet nack = make_nack_packet(peer, flow, missing_seq, prio, ingress);
        const PortId egress = ingress;

        _sim->schedule_after(_nic_latency_cur,
        [this, nack, egress]() mutable {
                enqueue_packet(egress, QClass::CONTROL, nack);
        });
        _telemetry.tx_nacks++;
}

void NetworkNic::schedule_trim_back_response(const Packet& trim,
                PortId ingress) noexcept {
        if (!_sim)
                return;

        const Packet tb = make_trim_back_response(trim, ingress);
        const PortId egress = ingress;

        _sim->schedule_after(_nic_latency_cur,
        [this, tb, egress]() mutable {
                enqueue_packet(egress, QClass::CONTROL, tb);
        });
        _telemetry.tx_trims++;
}

void NetworkNic::start_tx_unicast(NodeId dst, Bytes size, FlowPriority prio) {
        // Minimal segmentation loop: 1 flow_id per call; simple seq starting at 0
        static FlowId next_flow = 1;
        const FlowId fid = next_flow++;

        PortId port = pick_src_port_for_flow(dst);

        Bytes remaining = size;
        PacketSeq seq = 0;

        while (remaining > Bytes(0)) {
                const Bytes payload = std::min(remaining, DEFAULT_MSS_BYTES);
                Packet pkt = make_data_packet(dst, port, prio, fid, seq, payload);

                if (remaining <= DEFAULT_MSS_BYTES) {
                        pkt.set_eof(true);
                }

                enqueue_packet(port, (prio == FlowPriority::ELEPHANT) ? QClass::ELEPHANT :
                               QClass::MICE, pkt);

                _telemetry.tx_data++;
                remaining -= payload;
                seq = static_cast<PacketSeq>(static_cast<std::uint16_t>(seq + 1));
        }
}

void NetworkNic::start_tx_multicast(PacketGroups gmask, Bytes size,
                                    FlowPriority prio) {
        static FlowId next_flow = 1;
        const FlowId fid = next_flow++;

        PortId port = static_cast<PortId>(0);

        Bytes remaining = size;
        PacketSeq seq = 0;

        while (remaining > Bytes(0)) {
                const Bytes payload = std::min(remaining, DEFAULT_MSS_BYTES);

                Packet pkt(this->id(), port, NodeId(0), PortId(0),
                           PacketProtocol::DATA, prio, seq, fid, 0, 0, payload);
                pkt.set_groups(gmask);

                if (_lb) {
                        const uint16_t e = _lb->get_entropy(pkt);
                        pkt.set_entropy(e);
                }

                if (remaining <= DEFAULT_MSS_BYTES) {
                        pkt.set_eof(true);
                }

                enqueue_packet(port, (prio == FlowPriority::ELEPHANT) ? QClass::ELEPHANT :
                               QClass::MICE, pkt);

                _telemetry.tx_data++;
                remaining -= payload;
                seq = static_cast<PacketSeq>(static_cast<std::uint16_t>(seq + 1));
        }
}

void NetworkNic::enqueue_packet(PortId port, QClass cls, Packet pkt) noexcept {
        if (port >= _port_cnt) {
                log_error("ERROR", "NetworkNic::enqueue_packet invalid port");
                return;
        }

        auto& pq = _ports[port];

        switch (cls) {
        case QClass::CONTROL:
                pq.control.push_back(std::move(pkt));
                break;

        case QClass::RETRANS:
                pq.retrans.push_back(std::move(pkt));
                break;

        case QClass::MICE:
                pq.mice.push_back(std::move(pkt));
                break;

        case QClass::ELEPHANT:
                pq.elephant.push_back(std::move(pkt));
                break;
        }

        schedule_port_if_needed(port);
}

void NetworkNic::schedule_port_if_needed(PortId port) noexcept {
        auto& pq = _ports[port];

        if (pq.scheduled || !_sim)
                return;

        pq.scheduled = true;
        _sim->schedule_after(Time(0),
                             std::mem_fn(&NetworkNic::port_drain_task),
                             this, port);
}

bool NetworkNic::pick_next_qclass(const PortQueues& pq,
                                  QClass& out_cls) const noexcept {
        if (!pq.control.empty()  && pq.b_control  > Bytes(0)) {
                out_cls = QClass::CONTROL;
                return true;
        }

        if (!pq.retrans.empty()  && pq.b_retrans  > Bytes(0)) {
                out_cls = QClass::RETRANS;
                return true;
        }

        if (!pq.mice.empty()     && pq.b_mice     > Bytes(0)) {
                out_cls = QClass::MICE;
                return true;
        }

        if (!pq.elephant.empty() && pq.b_elephant > Bytes(0)) {
                out_cls = QClass::ELEPHANT;
                return true;
        }

        return false;
}

bool NetworkNic::try_send_one(PortId port, QClass cls) noexcept {
        auto& pq = _ports[port];
        std::deque<Packet> *q = nullptr;
        Bytes* budget = nullptr;

        switch (cls) {
        case QClass::CONTROL:
                q = &pq.control;
                budget = &pq.b_control;
                break;

        case QClass::RETRANS:
                q = &pq.retrans;
                budget = &pq.b_retrans;
                break;

        case QClass::MICE:
                q = &pq.mice;
                budget = &pq.b_mice;
                break;

        case QClass::ELEPHANT:
                q = &pq.elephant;
                budget = &pq.b_elephant;
                break;
        }

        if (!q || q->empty() || *budget <= Bytes(0))
                return false;

        Packet pkt = std::move(q->front());
        q->pop_front();

        const bool is_ctrl = (cls == QClass::CONTROL || cls == QClass::RETRANS);

        if (!is_ctrl && _cc) {
                const Bytes next_bytes = pkt.total_size();
                Bytes& out = _tx_outstanding[pkt.flow_id()];

                if (!_cc->is_allowed_to_send(out, next_bytes)) {
                        q->push_front(std::move(pkt));
                        return false;
                }

                out += next_bytes;
        }

        if (_sim) {
                _sim->schedule_after(_nic_latency_cur,
                [this, pkt, port]() mutable {
                        if (pkt.protocol() == PacketProtocol::DATA) {
                                record_tx_timestamp(pkt);
                        }
                        (void)_buf->enqueue_packet(pkt, port, pkt.priority());
                });
        }

        const Bytes cost = pkt.total_size();

        if (cost >= *budget)
                *budget = Bytes(0);
        else *budget = *budget - cost;

        return true;
}

void NetworkNic::port_drain_task(PortId port) noexcept {
        auto& pq = _ports[port];

        // Attempt to send while there is any budget and any queue has packets.
        // Refill budgets when all depleted or all queues empty.
        for (;;) {
                QClass cls;

                if (pick_next_qclass(pq, cls)) {
                        if (!try_send_one(port, cls)) {
                                // Could not send, stop and retry later.
                                break;
                        }

                        continue;
                }

                const bool any_nonempty = !(pq.control.empty() && pq.retrans.empty() &&
                                            pq.mice.empty() && pq.elephant.empty());
                const bool any_budget = (pq.b_control  > Bytes(0) ||
                                         pq.b_retrans  > Bytes(0) ||
                                         pq.b_mice     > Bytes(0) ||
                                         pq.b_elephant > Bytes(0));

                if (!any_budget && any_nonempty) {
                        pq.b_control  = _schedw.control;
                        pq.b_retrans  = _schedw.retrans;
                        pq.b_mice     = _schedw.mice;
                        pq.b_elephant = _schedw.elephant;
                        continue;
                }

                break;
        }

        pq.scheduled = false;

        if (!(pq.control.empty() && pq.retrans.empty() &&
              pq.mice.empty() && pq.elephant.empty())) {
                schedule_port_if_needed(port);
        }
}

} // namespace dofs
