// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ZMQStream2Pusher.h"
#include "../frame_serialize/CBORStream2Serializer.h"

ZMQStream2Pusher::ZMQStream2Pusher(const std::vector<std::string> &addr,
                                   std::optional<int32_t> send_buffer_high_watermark,
                                   std::optional<int32_t> send_buffer_size)
        : serialization_buffer(256*1024*1024),
          serializer(serialization_buffer.data(), serialization_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>(ZMQSocketType::Push);
        if (send_buffer_size)
            s->SendBufferSize(send_buffer_size.value());
        if (send_buffer_high_watermark)
            s->SendWaterMark(send_buffer_high_watermark.value());
        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 (!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 zmq_socket_free(void *data, void *hint) {
    auto z = (ZeroCopyReturnValue *) hint;
    z->release();
}

void ZMQStream2Pusher::SendImage(ZeroCopyReturnValue &z) {
    if (!socket.empty()) {
        auto socket_number = (z.GetImageNumber() / images_per_file) % socket.size();
        socket[socket_number]->SendZeroCopy(z.GetImage(), z.GetImageSize(), zmq_socket_free, &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; // Only writer on first socket is asked to write master file
        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());
    }
}

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;
    }
    return ret;
}

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

std::string ZMQStream2Pusher::Finalize() {
    std::string ret;
    if (writer_notification_socket) {
        for (int i = 0; i < socket.size(); i++) {
            auto n = writer_notification_socket->Receive(run_number, run_name);
            if (!n)
                ret += "Writer " + std::to_string(i) + ": no end notification received within 1 minute from collection end";
            else if (n->socket_number >= socket.size())
                ret += "Writer " + std::to_string(i) + ": mismatch in socket number";
            else if (!n->ok)
                ret += "Writer " + std::to_string(i) + ": " + n->error;
        }
    }
    return ret;
}

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, std::chrono::minutes(1));
    return *this;
}

std::string ZMQStream2Pusher::PrintSetup() const {
    std::string output = "ZMQStream2Pusher: Sending images to the following ZeroMQ sockets: ";
    for (const auto &s: socket)
        output += s->GetEndpointName() + " ";
    return output;
}