#include "network/network_switch.h"

#include <algorithm>
#include <iterator>
#include <set>

#include "network/switch/multicast_table.h"
#include "network/switch/unicast_table.h"

namespace dofs {

NetworkSwitch::NetworkSwitch(Simulator *const sim,
                             NodeId id,
                             uint16_t total_ports,
                             ECNEngine *const ecn,
                             SwitchBuffer *const buf,
                             const RoutingTables *const rt,
                             Time forwarding_latency,
                             Time multicast_dup_delay) noexcept
        : NetworkNode(sim, id, NodeType::SWITCH),
          _ecn(ecn),
          _buf(buf),
          _rt(rt),
          _port_cnt(total_ports),
          _fwd_latency_init(forwarding_latency),
          _fwd_latency_cur(forwarding_latency),
          _dup_delay(multicast_dup_delay) {

        Node::set_status(NodeStatus::OK);

        if (!_ecn) {
                log_error("ERROR", "NetworkSwitch constructed with null ECN engine");
        }

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

        if (!_rt) {
                log_error("ERROR", "NetworkSwitch constructed with null RoutingTables");
        }
}

void NetworkSwitch::set_status(NodeStatus s,
                               Time new_forward_latency) noexcept {
        switch (s) {
        case NodeStatus::OK:
                Node::set_status(NodeStatus::OK);
                _fwd_latency_cur = _fwd_latency_init;
                return;

        case NodeStatus::THROTTLING:
                if (new_forward_latency < _fwd_latency_init) {
                        log_error("ERROR",
                                  "NetworkSwitch::set_status(THROTTLING) new latency < initial");
                        return;
                }

                Node::set_status(NodeStatus::THROTTLING);
                _fwd_latency_cur = new_forward_latency;
                return;

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

void NetworkSwitch::recv_pkt(Packet &pkt_in, PortId ingress) {
        // If switch is DOWN or BOOTING, drop silently
        if (status() == NodeStatus::DOWN || status() == NodeStatus::BOOTING) {
                return;
        }

        Packet pkt = pkt_in;

        if (_ecn && _buf) {
                pkt = _ecn->process_packet(pkt, _buf);
        }

        if (_sim) {
                _sim->schedule_after(_fwd_latency_cur,
                                     std::mem_fn(&NetworkSwitch::forward_after_delay),
                                     this, pkt, ingress);
        }
}

void NetworkSwitch::forward_after_delay(Packet pkt, PortId ingress) noexcept {
        if (!_buf || !_rt)
                return;

        // Decide multicast vs unicast based on groups mask
        std::vector<PortId> fanout;
        fanout.reserve(8); // small typical fanout

        const bool is_mc = (pkt.groups() != static_cast<PacketGroups>(0));

        if (is_mc) {
                build_mc_fanout(pkt, ingress, fanout);
        } else {
                build_uc_egress(pkt, fanout);
        }

        // Filter invalid/self ports and keep determinism
        fanout.erase(std::remove_if(fanout.begin(), fanout.end(),
        [&](PortId p) {
                return (p >= _port_cnt) || (p == ingress);
        }),
        fanout.end());

        if (fanout.empty())
                return;

        const FlowPriority prio = pkt.priority();

        for (std::size_t i = 0; i < fanout.size(); ++i) {
                const PortId egress = fanout[i];
                const Time offset = _dup_delay * Time(i);

                if (_sim && offset > Time(0)) {
                        _sim->schedule_after(offset,
                                             std::mem_fn(&NetworkSwitch::enqueue_one),
                                             this, pkt, egress, prio);
                } else {
                        enqueue_one(pkt, egress, prio);
                }
        }
}

void NetworkSwitch::enqueue_one(Packet pkt, PortId egress,
                                FlowPriority prio) noexcept {
        (void)_buf->enqueue_packet(pkt, egress, prio);
}

void NetworkSwitch::build_uc_egress(const Packet& pkt,
                                    std::vector<PortId> &out) const {
        // Unicast: consult table; if multiple candidates, ECMP-pick one by entropy
        std::vector<PortId> candidates =
                _rt->unicast.get_port_list(pkt.dst_node(), pkt.dst_port());

        if (candidates.empty()) {
                // Fallback to dst_port if sane
                if (pkt.dst_port() < _port_cnt) {
                        out.clear();
                        out.push_back(pkt.dst_port());
                }

                return;
        }

        if (candidates.size() == 1) {
                out = candidates;
                return;
        }

        // Multiple candidates: pick one deterministically via hash_ecmp
        const uint32_t differentiator =
                (static_cast<uint32_t>(pkt.entropy()) << 16) ^
                static_cast<uint32_t>(pkt.dst_port());

        const uint16_t idx = hash_ecmp(pkt.src_node(), pkt.src_port(),
                                       pkt.dst_node(), pkt.dst_port(),
                                       differentiator,
                                       static_cast<uint16_t>(candidates.size()));

        out.clear();
        out.push_back(candidates[idx % candidates.size()]);
}

void NetworkSwitch::build_mc_fanout(const Packet& pkt, PortId ingress,
                                    std::vector<PortId> &fanout) const {
        if (!_rt)
                return;

        const PacketGroups mask = pkt.groups();
        const std::size_t gcount = _rt->multicast.group_count();

        for (std::size_t gid = 0; gid < gcount; ++gid) {
                if (!group_bit_set(mask, gid))
                        continue;

                const auto* trees = _rt->multicast.trees_of(gid);

                if (!trees || trees->empty())
                        continue;

                const uint16_t k = static_cast<uint16_t>(trees->size());
                uint16_t chosen_idx = 0;

                if (k > 1) {
                        // Entropy-based multipath spraying:
                        // differentiator mixes 16b entropy with 16b group index
                        const uint32_t differentiator =
                                (static_cast<uint32_t>(pkt.entropy()) << 16) ^
                                static_cast<uint32_t>(gid & 0xFFFF);

                        chosen_idx = hash_ecmp(pkt.src_node(), pkt.src_port(),
                                               pkt.dst_node(), pkt.dst_port(),
                                               differentiator, k);
                        chosen_idx %= k;
                }

                const McTree& t = (*trees)[chosen_idx];

                if (t.parent_port.has_value() && ingress != t.parent_port.value()) {
                        // Skip this group's replication from the wrong direction
                        continue;
                }

                merge_sorted_unique(fanout, t.child_ports);
        }
}

void NetworkSwitch::merge_sorted_unique(std::vector<PortId> &base,
                                        const std::vector<PortId> &add) {
        if (add.empty())
                return;

        if (base.empty()) {
                base = add;
                return;
        }

        std::vector<PortId> merged;
        merged.reserve(base.size() + add.size());
        std::set_union(base.begin(), base.end(),
                       add.begin(), add.end(),
                       std::back_inserter(merged));
        base.swap(merged);
}

inline bool NetworkSwitch::group_bit_set(PacketGroups mask,
                std::size_t gid) noexcept {
        return (mask >> gid) & static_cast<PacketGroups>(1);
}

} // namespace dofs
