// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ZMQImagePuller.h"
#include "../frame_serialize/CBORStream2Serializer.h"

#include <cstring>

ZMQImagePuller::ZMQImagePuller(const std::string &in_addr,
                               const std::string &repub_address,
    const std::optional<int32_t> &rcv_watermark,
    const std::optional<int32_t> &repub_watermark,
    const std::optional<size_t> &ring_size) :
        ImagePuller(DefaultQueueSize, ring_size.value_or(0)),
        socket (ZMQSocketType::Pull), addr(in_addr) {
    auto start_time = std::chrono::steady_clock::now();
    puller_thread = std::thread(&ZMQImagePuller::PullerThread, this);
    cbor_thread = std::thread(&ZMQImagePuller::CBORThread, this);

    socket.ReceiveWaterMark(rcv_watermark.value_or(default_receive_watermark));
    socket.ReceiveTimeout(ReceiveTimeout);
    socket.Connect(addr);

    if (!repub_address.empty()) {
        repub_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
        repub_socket->SendWaterMark(repub_watermark.value_or(default_repub_watermark));
        repub_socket->SendTimeout(RepubTimeout);
        repub_socket->Bind(repub_address);
        repub_thread = std::thread(&ZMQImagePuller::RepubThread, this);
    }
    auto end_time = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration<float>(end_time - start_time);
    logger.Info("ZMQImagePuller connected to {} in {:.3f} s", addr, duration.count());
}

ZMQImagePuller::~ZMQImagePuller() {
    ZMQImagePuller::Disconnect();
}

void ZMQImagePuller::Disconnect() {
    disconnect = 1;

    // The threads below hand messages on through bounded blocking queues, and by the time a puller
    // is disconnected nothing drains outside_fifo any more - the receiver that did is gone. A full
    // queue would then park CBORThread in PutBlocking for ever, back up cbor_fifo and park
    // PullerThread too, and the joins below would never return. Stop() releases them.
    cbor_fifo.Stop();
    repub_fifo.Stop();
    outside_fifo.Stop();
    // And PullerThread may be waiting for ring space that the stopped consumers will never free.
    ring.Stop();

    if (puller_thread.joinable())
        puller_thread.join();
    if (cbor_thread.joinable())
        cbor_thread.join();
    if (repub_thread.joinable())
        repub_thread.join();

    if (!addr.empty())
        socket.Disconnect(addr);

    addr = "";
}

void ZMQImagePuller::PullerThread() {
    while (true) {
        auto msg = std::make_shared<ZMQMessage>();
        bool received = false;
        while (!received) {
            if (disconnect) {
                cbor_fifo.PutBlocking(ImagePullerOutput{});
                return;
            }
            try {
                received = socket.Receive(*msg, false);
                if (!received)
                    std::this_thread::sleep_for(std::chrono::milliseconds(1));
            } catch (const JFJochException &e) {
                logger.ErrorException(e);
            }
        }
        ImagePullerOutput ret;
        if (ring.GetCapacity() == 0) {
            // No ring: the ZeroMQ message itself is the buffer, kept alive by the region. No copy.
            ret.raw = std::shared_ptr<ReceiveBufferRegion>(new ReceiveBufferRegion{msg->data(), msg->size()},
                                                           [msg](ReceiveBufferRegion *r) { delete r; });
        } else {
            // Copied into the ring, and the ZeroMQ message is released with the next receive. When the
            // ring is full this waits and the socket is not read: its queue fills to the watermark and
            // the sender's send blocks - for as long as the sender is prepared to wait. The broker's
            // ZeroMQ pusher waits one second, then drops the image and sends without waiting for the
            // rest of the run, so a ring that is too small for the writer's stalls costs images.
            ret.raw = ring.Allocate(msg->size());
            if (!ret.raw) {
                // Stop() - the puller is being disconnected.
                cbor_fifo.PutBlocking(ImagePullerOutput{});
                return;
            }
            if (msg->size() > 0)
                memcpy(ret.raw->data, msg->data(), msg->size());
        }
        cbor_fifo.PutBlocking(ret);
    }
}

void ZMQImagePuller::CBORThread() {
    auto ret = cbor_fifo.GetBlocking();

    while (ret.raw) {
        try {
            ret.cbor = CBORStream2Deserialize(ret.raw->data, ret.raw->size);
            // Even if we suspend consuming of the messages by the receiver,
            // it is still reasonable to have them republished
            // so republish functionality is not affected by Suspend()
            if (!suspend)
                outside_fifo.PutBlocking(ret);
            if (repub_socket) {
                if ((ret.cbor->msg_type == CBORImageType::START)
                    || (ret.cbor->msg_type == CBORImageType::END))
                    repub_fifo.PutBlocking(ret);
                else if (RepubCanHold(ret.raw->size) && repub_fifo.Put(ret))
                    repub_bytes += ret.raw->size;
            }
        } catch (const std::exception &e) {
            // std::exception: the deserialiser allocates a vector per per-image array of an END message,
            // and a bad_alloc leaving this thread would end the process.
            logger.ErrorException(e);
        }
        ret = cbor_fifo.GetBlocking();
    }
    if (repub_socket)
        repub_fifo.PutBlocking(ret);
    outside_fifo.PutBlocking(ret);
}

void ZMQImagePuller::RepubThread() {
    auto ret = repub_fifo.GetBlocking();
    bool repub_active = false;

    while (ret.raw) {
        try {
            if (ret.cbor->msg_type == CBORImageType::START) {
                // Start message needs to be cleaned when running republish
                StartMessage msg = ret.cbor->start_message.value();
                msg.writer_notification_zmq_addr = "";
                std::vector<uint8_t> serialization_buffer(256*1024*1024);
                CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
                serializer.SerializeSequenceStart(msg);
                repub_active = repub_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
                if (repub_active)
                    logger.Info("Republish active");
            } else {
                if (repub_active)
                    repub_socket->Send(ret.raw->data, ret.raw->size, true);
            }
        } catch (const std::exception &e) {
            logger.ErrorException(e);
        }
        if (ret.cbor && (ret.cbor->msg_type != CBORImageType::START) && (ret.cbor->msg_type != CBORImageType::END))
            repub_bytes -= ret.raw->size;
        ret = repub_fifo.GetBlocking();
    }
    if (repub_active)
        logger.Info("Republish finished");
}

void ZMQImagePuller::Suspend() {
    suspend = true;
}

void ZMQImagePuller::ResumeAndClear() {
    outside_fifo.Clear();
    outside_fifo.ClearMaxUtilization();
    suspend = false;
}