// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#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,
                               std::binary_semaphore *sempahore) {
    if (sockets.empty())
        return;
    auto socket_number = (image_number % file_count) % sockets.size();
    sockets[socket_number]->SendZeroCopy(image_data, image_size, sempahore);
}

void ZMQImagePusher::StartDataCollection(const StartMessage& message) {
    JFJochFrameSerializer serializer(80*1024*1024); // 80 MiB should be safe even for 16M

    if (message.data_file_count < 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "File count cannot be zero or negative");
    file_count = message.data_file_count;

    serializer.SerializeSequenceStart(message);

    for (const auto &s: sockets)
        s->Send(serializer.GetBuffer(), serializer.GetBufferSize(), true);
}


void ZMQImagePusher::EndDataCollection(const EndMessage& message) {
    JFJochFrameSerializer serializer(1024*1024);

    EndMessage end_message = message;

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