#include "network/nic/congestion_control.h"

namespace dofs {

CongestionControl::CongestionControl(Bytes init_cwnd, Bytes max_cwnd) noexcept
        : _cwnd(init_cwnd), _cwnd_max(max_cwnd) {
        if (_cwnd > _cwnd_max)
                _cwnd = _cwnd_max;

        if (_cwnd == Bytes(0))
                _cwnd = Bytes(1);
}

bool CongestionControl::is_allowed_to_send(Bytes bytes_outstanding,
                Bytes next_bytes) const noexcept {
        const Bytes needed = bytes_outstanding + next_bytes;
        return needed <= _cwnd;
}

// IMPORTANT: placeholders
DCQCN::DCQCN(Bytes init_cwnd, Bytes max_cwnd) noexcept
        : CongestionControl(init_cwnd, max_cwnd) {}

void DCQCN::update(const Packet& pkt, Time rtt) noexcept {
        // Extremely simple placeholder logic:
        // - If ECN marked or TRIM_BACK/NACK, cut cwnd multiplicatively.
        // - If ACK of data (no ECN), increase cwnd additively (slowly).
        // This gives you a cwnd-based shape; replace with real DCQCN later.

        const bool is_ctrl = (pkt.protocol() == PacketProtocol::ACK ||
                              pkt.protocol() == PacketProtocol::NACK ||
                              pkt.protocol() == PacketProtocol::HEADER_TRIM_BACK);

        if (pkt.protocol() == PacketProtocol::HEADER_TRIM_BACK ||
            pkt.protocol() == PacketProtocol::NACK ||
            pkt.is_ecn()) {
                // MD: halve cwnd, floor at 1 MSS
                Bytes new_cwnd = _cwnd / 2;

                if (new_cwnd < DEFAULT_MSS_BYTES)
                        new_cwnd = DEFAULT_MSS_BYTES;

                _cwnd = new_cwnd;

                if (_cwnd > _cwnd_max)
                        _cwnd = _cwnd_max;

                (void)rtt;
                (void)is_ctrl;
                return;
        }

        // Additive increase on clean ACKs; bound by cwnd_max.
        if (pkt.protocol() == PacketProtocol::ACK) {
                Bytes inc = DEFAULT_MSS_BYTES; // 1 MSS per RTT-ish sample
                Bytes new_cwnd = _cwnd + inc;

                if (new_cwnd > _cwnd_max)
                        new_cwnd = _cwnd_max;

                _cwnd = new_cwnd;
                (void)rtt;
                (void)is_ctrl;
        }
}

NSCC::NSCC(Bytes init_cwnd, Bytes max_cwnd) noexcept
        : CongestionControl(init_cwnd, max_cwnd) {}

void NSCC::update(const Packet& pkt, Time rtt) noexcept {
        if (pkt.protocol() == PacketProtocol::HEADER_TRIM_BACK ||
            pkt.protocol() == PacketProtocol::NACK ||
            pkt.is_ecn()) {
                Bytes new_cwnd = (_cwnd * 3) / 4; // 25% cut

                if (new_cwnd < DEFAULT_MSS_BYTES)
                        new_cwnd = DEFAULT_MSS_BYTES;

                _cwnd = new_cwnd;

                if (_cwnd > _cwnd_max)
                        _cwnd = _cwnd_max;

                (void)rtt;
                return;
        }

        if (pkt.protocol() == PacketProtocol::ACK) {
                Bytes inc = DEFAULT_MSS_BYTES;
                Bytes new_cwnd = _cwnd + inc;

                if (new_cwnd > _cwnd_max)
                        new_cwnd = _cwnd_max;

                _cwnd = new_cwnd;
                (void)rtt;
        }
}

} // namespace dofs
