// Copyright (2019-2024) Paul Scherrer Institute

#include "ZMQStream2Pusher.h"

ZMQStream2Pusher::ZMQStream2Pusher(ZMQContext &context, const std::string &addr, int32_t send_buffer_high_watermark,
                                   int32_t send_buffer_size)
        : socket(context, ZMQSocketType::Push) {
    Bind(addr, send_buffer_high_watermark, send_buffer_size);
}

ZMQStream2Pusher::ZMQStream2Pusher(const std::string &addr, int32_t send_buffer_high_watermark,
                                   int32_t send_buffer_size)
        : context(std::make_unique<ZMQContext>()),
          socket(*context, ZMQSocketType::Push) {
    Bind(addr, send_buffer_high_watermark, send_buffer_size);
}

void ZMQStream2Pusher::Bind(const std::string &addr, int32_t send_buffer_high_watermark, int32_t send_buffer_size) {

    if (send_buffer_size > 0)
        socket.SendBufferSize(send_buffer_size);
    if (send_buffer_high_watermark > 0)
        socket.SendWaterMark(send_buffer_high_watermark);
    socket.SendTimeout(std::chrono::seconds(5)); // 5 seconds should be more than enough to flush buffers and to still give fast response
    socket.Bind(addr);
}

void ZMQStream2Pusher::StartDataCollection(const StartMessage &message) {
    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);
    if (!socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
        throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr " + GetAddress());
}

bool ZMQStream2Pusher::SendImage(const DataMessage &message) {
    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);

    return socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), false); // Non-blocking
}

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

    serializer.SerializeSequenceEnd(message);
    return socket.Send(serialization_buffer.data(), serializer.GetBufferSize(), true); // Blocking
}

std::string ZMQStream2Pusher::GetAddress() {
    return socket.GetEndpointName();
}