// Copyright (2019-2023) Paul Scherrer Institute

#include "ZMQImagePusher.h"
#include "JFJochException.h"

ZMQImagePusher::ZMQImagePusher(ZMQContext &zmq_context, const std::vector<std::string> &addr,
                               int32_t send_buffer_high_watermark, int32_t send_buffer_size) {
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No writer ZMQ address provided");

    for (const auto &a : addr) {
        auto s = std::make_unique<ZMQSocket>(zmq_context, ZMQSocketType::Push);
        if (send_buffer_size > 0)
            s->SendBufferSize(send_buffer_size);
        if (send_buffer_high_watermark > 0)
            s->SendWaterMark(send_buffer_high_watermark);
        s->Bind(a);
        sockets.push_back(std::move(s));
    }
}

ZMQImagePusher::ZMQImagePusher(const std::vector<std::string> &addr,
                               int32_t send_buffer_high_watermark, int32_t send_buffer_size) {
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No writer ZMQ address provided");

    for (const auto &a : addr) {
        auto c = std::make_unique<ZMQContext>();
        auto s = std::make_unique<ZMQSocket>(*c, ZMQSocketType::Push);
        if (send_buffer_size > 0)
            s->SendBufferSize(send_buffer_size);
        if (send_buffer_high_watermark > 0)
            s->SendWaterMark(send_buffer_high_watermark);
        s->Bind(a);
        contexts.push_back(std::move(c));
        sockets.push_back(std::move(s));
    }
}

void ZMQImagePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number)  {
    if (sockets.empty())
        return;
    auto socket_number = (image_number % file_count) % sockets.size();
    sockets[socket_number]->Send(image_data, image_size);
}

void ZMQImagePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number,
                               ZeroCopyReturnValue *z) {
    if (sockets.empty())
        return;
    auto socket_number = (image_number % file_count) % sockets.size();
    sockets[socket_number]->SendZeroCopy(image_data, image_size, z);
}

void ZMQImagePusher::StartDataCollection(const uint8_t *image_data, size_t image_size, int64_t data_file_count) {
    if (data_file_count < 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "File count cannot be zero or negative");
    file_count = data_file_count;

    for (const auto &s: sockets)
        s->Send(image_data, image_size, true);
}


void ZMQImagePusher::EndDataCollection(const EndMessage& message) {
    std::vector<uint8_t> serialization_buffer(80*1024*1024);
    JFJochFrameSerializer serializer(serialization_buffer.data(), serialization_buffer.size());

    EndMessage end_message = message;

    for (const auto &s: sockets) {
        serializer.SerializeSequenceEnd(end_message);
        s->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
        end_message.write_master_file = false;
    }
}