#include "hosts/publisher.h"
#include "core/error.h"
#include "network/network_nic.h"

namespace dofs {

Publisher::Publisher(Simulator *const sim,
                     NodeId id,
                     Time update_latency_base,
                     std::unique_ptr<PubBasePolicy> policy,
                     Time mgmt_latency) noexcept
        : Host(sim, id)
        , _sim(sim)
        , _update_latency_base(update_latency_base)
        , _update_latency_cur(update_latency_base)
        , _mgmt_latency(mgmt_latency)
        , _policy(std::move(policy)) {}

void Publisher::set_status(NodeStatus s, Time new_latency) noexcept {
        // Enforce clamp: new_latency == 0 means "no-op"
        if (new_latency.count() != 0 && new_latency < _update_latency_base) {
                DOFS_ERROR("publisher", "set_status: new latency below base; ignored");
        } else if (new_latency.count() != 0) {
                _update_latency_cur = new_latency;
        }

        // Forward to Node base if available (Host derives from Node).
        // If Node::set_status is public, this works; otherwise adjust when you expose it.
        // (We keep this call for correctness; remove if not yet available.)
        // Node::set_status(s);

        if (s == NodeStatus::DOWN) {
                if (_staging_evt != NULL_EVENT) {
                        _sim->cancel(_staging_evt);
                        _staging_evt = NULL_EVENT;
                }

                _state = State::Quiesced;
        }
}

void Publisher::recv_update(Bytes size,
                            PacketGroups update_groups_mask) noexcept {
        if (size == 0 || update_groups_mask == 0) {
                DOFS_ERROR("publisher", "recv_update: invalid size or groups");
                return;
        }

        _queue.push_back(Pending{size, update_groups_mask, _sim->now()});
        ++_updates_in;

        if (_state == State::Idle) {
                _state = State::Staging;
                arm_staging_if_needed();
        }
}

void Publisher::arm_staging_if_needed() noexcept {
        if (_staging_evt != NULL_EVENT)
                return;

        if (_queue.empty()) {
                _state = State::Idle;
                return;
        }

        _staging_evt = _sim->schedule_after(_update_latency_cur, [this]() {
                this->on_staging_timer();
        });
}

void Publisher::on_staging_timer() noexcept {
        _staging_evt = NULL_EVENT;

        if (_state == State::Quiesced)
                return;

        if (_queue.empty()) {
                _state = State::Idle;
                return;
        }

        Pending p = _queue.front();
        _queue.pop_front();

        const PacketGroups mcast_mask = _policy ? _policy->select_multicast_groups(
                                                p.mask) : PacketGroups{0};

        if (mcast_mask == 0) {
                DOFS_ERROR("publisher", "policy produced empty multicast mask");
        } else {
                auto* n = this->nic();

                if (!n) {
                        DOFS_ERROR("publisher", "no NIC attached; dropping update");
                } else {
                        n->send_flow(mcast_mask, p.size, FlowPriority::MICE);
                        _bytes_out += p.size;
                }
        }

        // Next
        if (_queue.empty()) {
                _state = State::Idle;
        } else {
                arm_staging_if_needed();
        }
}

void Publisher::recv_mgmt_msg(MgmtMsgPtr msg) noexcept {
        if (!msg)
                return;

        ++_mgmt_in;

        switch (msg->kind()) {
        case MgmtKind::END_SIMULATION: {
                if (_staging_evt != NULL_EVENT) {
                        _sim->cancel(_staging_evt);
                        _staging_evt = NULL_EVENT;
                }

                _state = State::Quiesced;
        }
        break;

        case MgmtKind::HEARTBEAT:
        case MgmtKind::JOB_FINISHED:
        default:
                // informational only for now
                break;
        }
}

void Publisher::recv_flow(NodeId, FlowId, FlowPriority, Bytes) {
        // Publisher is mainly a sender in this workload; ignore data-plane flow arrivals.
}

void Publisher::recv_frame(const Packet &) {
        // Control frames not used by Publisher in this stage.
}

} // namespace dofs
