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

#pragma once

#include "ImagePusher.h"
#include "../common/ZMQWrappers.h"
#include "../preview/PreviewCounter.h"
#include "ZMQWriterNotificationPuller.h"
#include "ZMQStream2PusherSocket.h"

class ZMQStream2Pusher : public ImagePusher {
    std::vector<std::unique_ptr<ZMQStream2PusherSocket>> socket;

    std::unique_ptr<ZMQWriterNotificationPuller> writer_notification_socket;

    int64_t images_per_file = 1;
    uint64_t run_number = 0;
    std::string run_name;
    std::atomic<bool> transmission_error = false;

    mutable std::mutex images_written_mutex;
    std::optional<uint64_t> images_written;
public:
    explicit ZMQStream2Pusher(const std::vector<std::string>& addr,
                              std::optional<int32_t> send_buffer_high_watermark = {},
                              std::optional<int32_t> send_buffer_size = {});

    ZMQStream2Pusher& WriterNotificationSocket(const std::string& addr);
    std::string GetWriterNotificationSocketAddress() const override;

    std::vector<std::string> GetAddress();

    // Strictly serial, as order of these is important
    void StartDataCollection(StartMessage& message) override;
    bool EndDataCollection(const EndMessage& message) override;
    bool SendCalibration(const CompressedImage& message) override;

    // Thread-safe
    bool SendImage(ZeroCopyReturnValue &z) override;
    bool SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) override;

    std::string Finalize() override;

    std::string PrintSetup() const override;
    std::optional<uint64_t> GetImagesWritten() const override;

    size_t GetConnectedWriters() const override;

    ImagePusherType GetType() const override { return ImagePusherType::ZMQ; }
};