// Copyright (2019-2023) Paul Scherrer Institute

#include "ZMQImagePuller.h"

ZMQImagePuller::ZMQImagePuller(const std::string &repub_address) :
        socket (ZMQSocketType::Pull) {
    socket.ReceiveWaterMark(ReceiverWaterMark);
    socket.ReceiveTimeout(ReceiveTimeout);

    if (!repub_address.empty()) {
        repub_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
        repub_socket->SendWaterMark(100);
        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);

    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 {
            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;
}