// Copyright (2019-2024) Paul Scherrer Institute

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

ZMQStream2Pusher::ZMQStream2Pusher(const std::vector<std::string> &addr,
                                             int32_t send_buffer_high_watermark, int32_t send_buffer_size)
        : serialization_buffer(256*1024*1024),
          serializer(serialization_buffer.data(), serialization_buffer.size()),
          preview_counter(std::chrono::seconds(1)) {
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writer ZMQ address provided");

    for (const auto &a : addr) {
        auto s = std::make_unique<ZMQSocket>(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->SendTimeout(std::chrono::seconds(5)); // 5 seconds should be more than enough to flush buffers and to still give fast response
        s->Bind(a);
        socket.emplace_back(std::move(s));
    }
}

bool ZMQStream2Pusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
    if (preview_socket) {
        if (preview_counter.GeneratePreview())
            preview_socket->Send(image_data, image_size, false);
    }

    if (!socket.empty()) {
        auto socket_number = (image_number / images_per_file) % socket.size();
        return socket[socket_number]->Send(image_data, image_size, false);
    } else
        return false;
}

void ZMQStream2Pusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number,
                                      ZeroCopyReturnValue *z) {
    if (preview_socket) {
        if (preview_counter.GeneratePreview())
            preview_socket->Send(image_data, image_size, false);
    }

    if (!socket.empty()) {
        auto socket_number = (image_number / images_per_file) % socket.size();
        socket[socket_number]->SendZeroCopy(image_data, image_size, z);
    } else
        z->release();
}

void ZMQStream2Pusher::StartDataCollection(StartMessage& message) {
      if (message.images_per_file < 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Images per file cannot be zero or negative");
    images_per_file = message.images_per_file;
    run_number = message.run_number;
    run_name = message.run_name;

    for (int i = 0; i < socket.size(); i++) {
        message.socket_number = i;
        if (i > 0)
            message.write_master_file = false;
        serializer.SerializeSequenceStart(message);
        if (!socket[i]->Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
            throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message on addr "
                                                                   + socket[i]->GetEndpointName());
    }

    if (preview_socket)
        preview_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
}

bool ZMQStream2Pusher::SendCalibration(const CompressedImage &message) {
    if (socket.empty())
        return false;

    serializer.SerializeCalibration(message);

    return socket[0]->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
}

bool ZMQStream2Pusher::EndDataCollection(const EndMessage& message) {
    serializer.SerializeSequenceEnd(message);

    bool ret = true;
    for (auto &s: socket) {
        if (!s->Send(serialization_buffer.data(), serializer.GetBufferSize(), true))
            ret = false;
    }

    if (preview_socket)
        preview_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);

    return ret;
}

std::vector<std::string> ZMQStream2Pusher::GetAddress() {
    std::vector<std::string> ret;
    for (auto &p: socket)
        ret.push_back(p->GetEndpointName());
    return ret;
}

ZMQStream2Pusher &ZMQStream2Pusher::PreviewSocket(const std::string &addr) {
    preview_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Pub);
    preview_socket->Bind(addr);
    return *this;
}

std::string ZMQStream2Pusher::GetPreviewAddress() {
    if (preview_socket)
        return preview_socket->GetEndpointName();
    else
        return "";
}

void ZMQStream2Pusher::Finalize() {
    if (writer_notification_socket) {
        for (int i = 0; i < socket.size(); i++) {
            auto n = writer_notification_socket->Receive(run_number, run_name);
            if (!n)
                throw JFJochException(JFJochExceptionCategory::FileWriteError, "No notification received from writer");
            else if (n->socket_number >= socket.size())
                throw JFJochException(JFJochExceptionCategory::FileWriteError, "Wrong socket number provided in the message");
            else if (!n->ok)
                throw JFJochException(JFJochExceptionCategory::FileWriteError, "Writer (socket "
                                                                               + socket[n->socket_number]->GetEndpointName()
                                                                               + ") finished with error");
        }
    }
}

std::string ZMQStream2Pusher::GetWriterNotificationSocketAddress() const {
    if (writer_notification_socket)
        return writer_notification_socket->GetEndpointName();
    else
        return "";
}

ZMQStream2Pusher &ZMQStream2Pusher::WriterNotificationSocket(const std::string &addr) {
    writer_notification_socket = std::make_unique<ZMQWriterNotificationPuller>(addr);
    return *this;
}

ZMQStream2Pusher &ZMQStream2Pusher::PreviewCounterPeriod(std::chrono::microseconds input) {
    preview_counter.Period(input);
    return *this;
}