// Copyright (2019-2023) Paul Scherrer Institute

#include "ZMQStream2Pusher.h"
#include "CBORStream2Serializer.h"

ZMQStream2Pusher::ZMQStream2Pusher(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));
    }
}

ZMQStream2Pusher::ZMQStream2Pusher(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 ZMQStream2Pusher::SendImage(const DataMessage& message)  {
    if (sockets.empty())
        return;
    auto socket_number = (message.number % file_count) % sockets.size();

    size_t approx_size = message.image.size + 2*1024*1024;
    std::vector<uint8_t> serialization_buffer(approx_size);

    CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
    serializer.SerializeImage(message);

    sockets[socket_number]->Send(serialization_buffer.data(), serializer.GetBufferSize());
}

void ZMQStream2Pusher::StartDataCollection(const StartMessage& message) {
    if (message.data_file_count < 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "File count cannot be zero or negative");
    file_count = message.data_file_count;

    size_t approx_size = 1024*1024;
    for (const auto &x : message.pixel_mask)
        approx_size += x.size;

    for (const auto &x : message.calibration)
        approx_size += x.size;

    std::vector<uint8_t> serialization_buffer(approx_size);
    CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
    serializer.SerializeSequenceStart(message);
    for (const auto &s: sockets)
        s->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
}


void ZMQStream2Pusher::EndDataCollection(const EndMessage& message) {
    std::vector<uint8_t> serialization_buffer(80*1024*1024);
    CBORStream2Serializer 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;
    }
}