#ifndef NETWORK_SWITCH_SWITCH_BUFFER_H
#define NETWORK_SWITCH_SWITCH_BUFFER_H

#include <cstdint>
#include <vector>
#include <deque>
#include <optional>

#include "core/time.h"
#include "core/types.h"
#include "core/simulator.h"
#include "core/error.h"
#include "network/packet.h"
#include "network/link.h"

namespace dofs {

class NetworkSwitch;

// Base class for switch egress buffering with weighted priorities.
// Implementations: SharedBuffer, DedicatedBuffer.
class SwitchBuffer {
public:
        struct Queued {
                Packet pkt;
                PortId egress;
                FlowPriority prio;
                Time enq_time;
                Bytes size_bytes;
        };

        virtual ~SwitchBuffer() = default;

        // --- API: enqueue & drain ---
        // NOTE: ECN/drop decisions are expected to be done *before* enqueue;
        // we still guard and return false on capacity violations.
        virtual bool enqueue_packet(const Packet& pkt,
                                    PortId egress,
                                    FlowPriority prio) = 0;

        // Attempt to send exactly one packet on 'port' using weighted priorities.
        // Returns true if a packet was reserved and scheduled.
        virtual bool drain_one(PortId port) = 0;

        // --- Accessors (implemented in .cc) ---
        Bytes buffer_size() const noexcept;
        SwitchBufferType type() const noexcept;
        uint16_t port_cnt() const noexcept;
        Bytes port_buffered(PortId p) const noexcept;
        const std::vector<Bytes> &ports_buffered() const noexcept;

        // Priority shares: percentages [0..100]. Sum must be 100.
        uint8_t share_ctrl() const noexcept;
        uint8_t share_mice() const noexcept;
        uint8_t share_elephant() const noexcept;

        // --- Mutators for priority shares (implemented in .cc) ---
        // If the 3 shares don't sum to 100, we normalize proportionally.
        void set_share_ctrl(uint8_t pct) noexcept;
        void set_share_mice(uint8_t pct) noexcept;
        void set_share_elephant(uint8_t pct) noexcept;

        // Owner & env
        const Simulator *simulator() const noexcept;
        const NetworkSwitch *owner() const noexcept;

        // Link wiring (egress side). Size must equal port_cnt().
        void set_egress_links(const std::vector<Link*> &links);

protected:
        SwitchBuffer(Simulator* const sim,
                     NetworkSwitch* const owner,
                     SwitchBufferType t,
                     Bytes total_bytes,
                     uint16_t ports);

        // Per-priority index and helpers
        static constexpr int PRI_COUNT = 3;
        static constexpr int PRI_CTRL = 0;
        static constexpr int PRI_MICE = 1;
        static constexpr int PRI_ELE  = 2;

        static int to_idx(FlowPriority p) noexcept {
                switch (p) {
                case FlowPriority::CTRL:
                        return PRI_CTRL;

                case FlowPriority::MICE:
                        return PRI_MICE;

                case FlowPriority::ELEPHANT:
                        return PRI_ELE;

                default:
                        return PRI_CTRL;
                }

                return PRI_ELE; // defensive
        }

        // Weighted Deficit Round Robin across priorities (per port):
        // Maintain byte-deficits per prio; each drain attempt adds quantum
        // proportional to the configured share. Pick first prio whose head fits.
        struct PerPortSched {
                std::array<int64_t, PRI_COUNT> deficit_bytes {0, 0, 0};
                int next_pick = 0; // rotating pointer among priorities
        };

        // Quantum baseline (bytes) per "share unit". Calibrated to MSS-ish.
        // Effective quantum = share_pct * QUANTUM_UNIT.
        static constexpr int64_t QUANTUM_UNIT = 128; // bytes per % share

        // For derived classes
        bool drain_one_common(PortId port);

        // Per-port self-scheduling
        void schedule_drain_if_needed(PortId port);
        void drain_once(PortId port); // event-loop callback

        // Derived classes must expose queues and capacity guards:

        // Access the 3 per-priority queues for a port.
        virtual std::array<std::deque<Queued>, PRI_COUNT> &queues_for(PortId p) = 0;
        virtual const std::array<std::deque<Queued>, PRI_COUNT> &queues_for(
                PortId p) const = 0;

        // Capacity checks and accounting; return false if not enough space.
        virtual bool on_enqueue_cap_check(PortId port, Bytes sz) = 0;
        virtual void on_enqueue_commit(PortId port, Bytes sz) = 0;
        virtual void on_dequeue_commit(PortId port, Bytes sz) = 0;

        // Attempt reservation+schedule on the egress link for (port, packet).
        // Returns the reservation finish time if successful (for scheduling the next drain).
        std::optional<Time> try_reserve_and_send(PortId port, Queued &q);

        // Helpers
        Bytes queued_bytes_total() const noexcept;
        Bytes queued_bytes_port(PortId p) const noexcept;

protected:
        Simulator *const _sim;
        NetworkSwitch *const _owner;

        SwitchBufferType _type;
        Bytes _buffer_bytes;
        uint16_t _port_cnt;

        // Egress link per port (owned by switch; not by buffer)
        std::vector<Link *> _egress_links;

        std::vector<bool> _drain_scheduled;

        // Per-port accounting and sched state
        std::vector<Bytes> _per_port_bytes;
        std::vector<PerPortSched> _sched; // per-port deficits

        // Priority shares (percentages). Maintained to sum to 100.
        uint8_t _share_ctrl;
        uint8_t _share_mice;
        uint8_t _share_ele;
};

} // namespace dofs

#endif // NETWORK_SWITCH_SWITCH_BUFFER_H