#include "network/switch/switch_buffer.h"

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

#include "network/network_switch.h"

namespace dofs {

SwitchBuffer::SwitchBuffer(Simulator* const sim,
                           NetworkSwitch* const owner,
                           SwitchBufferType t,
                           Bytes total_bytes,
                           uint16_t ports)
        : _sim(sim),
          _owner(owner),
          _type(t),
          _buffer_bytes(total_bytes),
          _port_cnt(ports),
          _egress_links(ports, nullptr),
          _drain_scheduled(ports, false),
          _per_port_bytes(ports, Bytes(0)),
          _sched(ports),
          _share_ctrl(30), _share_mice(50), _share_ele(20) // sane defaults
{}

Bytes SwitchBuffer::buffer_size() const noexcept {
        return _buffer_bytes;
}
SwitchBufferType SwitchBuffer::type() const noexcept {
        return _type;
}
uint16_t SwitchBuffer::port_cnt() const noexcept {
        return _port_cnt;
}
Bytes SwitchBuffer::port_buffered(PortId p) const noexcept {
        if (p >= _per_port_bytes.size())
                return Bytes(0);

        return _per_port_bytes[p];
}
const std::vector<Bytes> &SwitchBuffer::ports_buffered() const noexcept {
        return _per_port_bytes;
}
uint8_t SwitchBuffer::share_ctrl() const noexcept {
        return _share_ctrl;
}
uint8_t SwitchBuffer::share_mice() const noexcept {
        return _share_mice;
}
uint8_t SwitchBuffer::share_elephant() const noexcept {
        return _share_ele;
}

const Simulator *SwitchBuffer::simulator() const noexcept {
        return _sim;
}
const NetworkSwitch *SwitchBuffer::owner() const noexcept {
        return _owner;
}

void SwitchBuffer::set_egress_links(const std::vector<Link*> &links) {
        if (links.size() != _port_cnt) {
                log_error("ERROR", "SwitchBuffer::set_egress_links size mismatch");
                return;
        }

        _egress_links = links;
}

static void normalize_triple(uint8_t &a, uint8_t &b, uint8_t &c) {
        int sum = int(a) + int(b) + int(c);

        if (sum == 100 || sum == 0) {
                if (sum == 0) {
                        a = 34;
                        b = 33;
                        c = 33;
                }

                return;
        }

        // proportional normalization
        double k = 100.0 / double(sum);
        int na = int(std::round(a * k));
        int nb = int(std::round(b * k));
        int nc = 100 - na - nb;

        if (nc < 0)
                nc = 0;

        a = uint8_t(std::clamp(na, 0, 100));
        b = uint8_t(std::clamp(nb, 0, 100));
        c = uint8_t(std::clamp(nc, 0, 100));
}

void SwitchBuffer::set_share_ctrl(uint8_t pct) noexcept {
        _share_ctrl = pct;
        normalize_triple(_share_ctrl, _share_mice, _share_ele);
}
void SwitchBuffer::set_share_mice(uint8_t pct) noexcept {
        _share_mice = pct;
        normalize_triple(_share_ctrl, _share_mice, _share_ele);
}
void SwitchBuffer::set_share_elephant(uint8_t pct) noexcept {
        _share_ele = pct;
        normalize_triple(_share_ctrl, _share_mice, _share_ele);
}

Bytes SwitchBuffer::queued_bytes_total() const noexcept {
        Bytes sum(0);

        for (auto b : _per_port_bytes)
                sum += b;

        return sum;
}
Bytes SwitchBuffer::queued_bytes_port(PortId p) const noexcept {
        return port_buffered(p);
}

std::optional<Time> SwitchBuffer::try_reserve_and_send(PortId port, Queued& q) {
        if (port >= _egress_links.size() || _egress_links[port] == nullptr) {
                log_error("ERROR", "SwitchBuffer: egress link missing on port");
                return std::nullopt;
        }

        Link* link = _egress_links[port];

        auto r = link->reserve(q.size_bytes, _owner->id());

        if (!r.has_value()) {
                return std::nullopt; // link down or invalid sender
        }

        const Time now = _sim->now();
        const Time to_finish = (r->finish > now) ? r->finish.unsafe_sub(now) : Time(0);
        const Time deliver_after = to_finish + link->propagation_latency();

        link->schedule_delivery_after(q.pkt, _owner->id(), deliver_after);

        return r->finish;
}

void SwitchBuffer::schedule_drain_if_needed(PortId port) {
        if (port >= _port_cnt)
                return;

        if (_drain_scheduled[port])
                return;

        if (queued_bytes_port(port) == Bytes(0))
                return;

        Link* link = _egress_links[port];

        if (!link)
                return;

        const Time now = _sim->now();
        const Time na = link->next_available(_owner->id());
        const Time delay = (na > now) ? na.unsafe_sub(now) : Time(0);

        _drain_scheduled[port] = true;
        _sim->schedule_after(delay,
                             std::mem_fn(&SwitchBuffer::drain_once), this, port);
}

void SwitchBuffer::drain_once(PortId port) {
        _drain_scheduled[port] = false;

        // Try to send one packet. If we made progress, drain_one_common()
        // already scheduled the next drain at the right time.
        if (drain_one_common(port)) {
                return;
        }

        // No progress: either queues empty or link not reservable now.
        // If queues are still non-empty, retry later at next-available.
        auto& qs = queues_for(port);
        bool any = !qs[PRI_CTRL].empty() || !qs[PRI_MICE].empty() ||
                   !qs[PRI_ELE].empty();

        if (any) {
                schedule_drain_if_needed(port);
        }
}

bool SwitchBuffer::drain_one_common(PortId port) {
        if (port >= _port_cnt)
                return false;

        auto& queues = queues_for(port);
        auto& sched  = _sched[port];

        auto head_size = [&](int pri) -> Bytes {
                if (queues[pri].empty())
                        return Bytes(0);

                return queues[pri].front().size_bytes;
        };

        sched.deficit_bytes[PRI_CTRL] += int64_t(_share_ctrl) * QUANTUM_UNIT;
        sched.deficit_bytes[PRI_MICE] += int64_t(_share_mice) * QUANTUM_UNIT;
        sched.deficit_bytes[PRI_ELE]  += int64_t(_share_ele)  * QUANTUM_UNIT;

        for (int k = 0; k < PRI_COUNT; ++k) {
                int pri = (sched.next_pick + k) % PRI_COUNT;
                Bytes sz = head_size(pri);

                if (sz > Bytes(0) && sched.deficit_bytes[pri] >= int64_t(sz)) {
                        auto& q = queues[pri].front();
                        auto finish_opt = try_reserve_and_send(port, q);

                        if (finish_opt.has_value()) {
                                queues[pri].pop_front();
                                sched.deficit_bytes[pri] -= int64_t(sz);
                                on_dequeue_commit(port, sz);
                                sched.next_pick = (pri + 1) % PRI_COUNT;

                                Link* link = _egress_links[port];
                                const Time na = link->next_available(_owner->id());
                                const Time now = _sim->now();
                                const Time delay = (na > now) ? na.unsafe_sub(now) : Time(0);
                                _drain_scheduled[port] = true;
                                _sim->schedule_after(delay,
                                                     std::mem_fn(&SwitchBuffer::drain_once), this, port);

                                return true;
                        } else {
                                // Could not reserve now (link down). Stop early
                                return false;
                        }
                }
        }

        sched.next_pick = (sched.next_pick + 1) % PRI_COUNT;
        return false;
}

} // namespace dofs
