// 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"

ZMQImagePuller::ZMQImagePuller(const std::string &repub_address,
    const std::optional<int32_t> &rcv_watermark,
    const std::optional<int32_t> &repub_watermark) :
        socket (ZMQSocketType::Pull) {

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

    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(std::chrono::milliseconds(100));
        repub_socket->Bind(repub_address);
    }
}

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

void ZMQImagePuller::Connect(const std::string &in_address) {
    Disconnect();

    disconnect = 0;
    abort = 0;
    addr = in_address;
    socket.Connect(in_address);

    cbor_fifo.ClearMaxUtilization();
    repub_fifo.ClearMaxUtilization();
    outside_fifo.ClearMaxUtilization();

    puller_thread = std::thread(&ZMQImagePuller::PullerThread, this);
    cbor_thread = std::thread(&ZMQImagePuller::CBORThread, this);
    if (repub_socket)
        repub_thread = std::thread(&ZMQImagePuller::RepubThread, this);
}

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

    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::Abort() {
    abort = 1;
}

void ZMQImagePuller::PullerThread() {
    while (true) {
        ZMQImagePullerOutput ret;
        ret.msg = std::make_shared<ZMQMessage>();
        bool received = false;
        while (!received) {
            if (disconnect) {
                cbor_fifo.Put(ZMQImagePullerOutput{});
                return;
            }
            try {
                received = socket.Receive(*ret.msg, true);
            } catch (const JFJochException &e) {
                logger.ErrorException(e);
            }
        }
        cbor_fifo.Put(ret);
    }
}

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

    while (ret.msg) {
        try {
            ret.cbor = CBORStream2Deserialize(ret.msg->data(), ret.msg->size());
            outside_fifo.Put(ret);
            if (repub_socket)
                repub_fifo.Put(ret);
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }
        ret = cbor_fifo.GetBlocking();
    }
    if (repub_socket)
        repub_fifo.Put(ret);
    outside_fifo.Put(ret);
}

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

    while (ret.msg) {
        // Republishing is non-blocking for images
        // and blocking (with 100ms timeout) for START/END
        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_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
            } else
                repub_socket->Send(ret.msg->data(), ret.msg->size(), ret.cbor->msg_type != CBORImageType::IMAGE);
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }
        ret = repub_fifo.GetBlocking();
    }
}

ZMQImagePullerOutput ZMQImagePuller::WaitForImage() {
    ZMQImagePullerOutput ret{};
    while (!outside_fifo.Get(ret) && !abort)
        std::this_thread::sleep_for(std::chrono::microseconds(100));
    return ret;
}

ZMQImagePullerStatistics ZMQImagePuller::GetStatistics() const {
    return {
        .cbor_fifo_max_util = cbor_fifo.GetMaxUtilization(),
        .outside_fifo_max_util = outside_fifo.GetMaxUtilization(),
        .repub_fifo_max_util = repub_fifo.GetMaxUtilization()
    };
}