#include "network/switch/ecn_shared_red.h"

#include <algorithm>
#include <numeric>
#include <cmath>

namespace dofs {

SharedREDEngine::Thresholds SharedREDEngine::dynamic_thresholds(
        const Packet &pkt,
        const SwitchBuffer *buf,
        double total_bytes_avg) const {
        const Bytes cap_b = buf->buffer_size();
        const double cap  = static_cast<double>(cap_b);

        if (cap <= 0.0) {
                return {0.0, 0.0, 0.0, false};
        }

        // Active ports heuristic
        const auto &per = buf->ports_buffered();
        const std::size_t active = std::max<std::size_t>(1,
        std::count_if(per.begin(), per.end(), [](Bytes b) {
                return b > 0;
        }));

        const double fair_share = cap / static_cast<double>(active);

        // Class factors: be generous to CTRL, moderate for MICE, strict for ELE.
        double alpha = 0.20, beta = 0.85, pmax = 0.10; // mice defaults
        bool allow_trim = true;

        if (is_ctrl(pkt)) {
                alpha = 0.80;
                beta = 0.98;
                pmax = 0.02;
                allow_trim = false; // almost no marking, never trim
        } else if (is_ele(pkt)) {
                alpha = 0.10;
                beta = 0.60;
                pmax = 0.25;
                allow_trim = true;  // aggressive on elephants
        }

        // Dynamic bias: if total avg usage is high, tighten thresholds a bit
        // shrink factor in [0.8, 1.0] as pool fills
        const double util = std::clamp(total_bytes_avg / cap, 0.0, 1.0);
        const double tighten = 1.0 - 0.2 * util;

        const double min_th = alpha * fair_share * tighten;
        const double max_th = std::max(min_th + 1.0, beta * fair_share * tighten);

        return {min_th, max_th, pmax, allow_trim};
}

Packet &SharedREDEngine::process_packet(Packet &pkt,
                                        SwitchBuffer *buf) noexcept {
        if (!buf)
                return pkt;

        const uint16_t pc = buf->port_cnt();
        ensure_size(pc);

        // Update EWMA of total bytes in the shared pool.
        const auto &per = buf->ports_buffered();
        const double total_now = std::accumulate(per.begin(), per.end(), 0.0,
        [](double acc, Bytes b) {
                return acc + static_cast<double>(b);
        });
        _avg_total_bytes = ewma(total_now, _avg_total_bytes);

        // Also keep per-port EWMA to stabilize active-port detection (optional).
        for (uint16_t p = 0; p < pc; ++p) {
                _avg_port_bytes[p] = ewma(static_cast<double>(per[p]),
                                          _avg_port_bytes[p]);
        }

        // Use per-port instantaneous occupancy for the target port.
        const PortId port = pkt.dst_port();

        if (port >= pc)
                return pkt;

        // Derive class-aware dynamic thresholds
        const auto thr = dynamic_thresholds(pkt, buf, _avg_total_bytes);

        // Choose which queue measure to feed RED: per-port EWMA
        const double q_avg = _avg_port_bytes[port];

        if (q_avg <= thr.min_th) {
                // Well below min_th: no mark
                return pkt;
        }

        if (q_avg >= thr.max_th) {
                // Above max_th: trim (unless forbidden for class)
                if (thr.allow_trim) {
                        header_trim(pkt, false);
                } else {
                        // CTRL: at worst, mark ECN if allowed
                        maybe_mark_ecn(pkt);
                }

                return pkt;
        }

        // Between thresholds: probabilistic ECN
        const double p = thr.p_max * (q_avg - thr.min_th) / std::max(1.0,
                         (thr.max_th - thr.min_th));

        if (_rng) {
                if (_rng->uniform01() < p) {
                        maybe_mark_ecn(pkt);
                }
        } else {
                // Deterministic fallback: mark if over mid-point
                const double mid = 0.5 * (thr.min_th + thr.max_th);

                if (q_avg > mid)
                        maybe_mark_ecn(pkt);
        }

        return pkt;
}

} // namespace dofs
