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

#pragma once

#include <vector>

#include "../common/DiffractionExperiment.h"
#include "../common/JFJochMessages.h"
#include "../common/ZeroCopyReturnValue.h"
#include "../common/Logger.h"
#include "../frame_serialize/CBORStream2Serializer.h"

enum class ImagePusherType {HDF5, CBOR, TCP, ZMQ, Test, None};

struct ImagePusherStatus {
    ImagePusherType pusher_type = ImagePusherType::None;
    std::vector<std::string> address = {};
    size_t connected_writers = 0;
    std::optional<size_t> images_written = 0;
    std::optional<size_t> images_write_error = 0;
    std::vector<int64_t> writer_fifo_utilization;
};

struct ImagePusherQueueElement {
    const uint8_t *image_data;
    ZeroCopyReturnValue *z;
    bool end;
};

void PrepareCBORImage(DataMessage& message,
                      const DiffractionExperiment &experiment,
                      void *image, size_t image_size);

class ImagePusher {
protected:
    std::vector<uint8_t> serialization_buffer;
    CBORStream2Serializer serializer;
    ImagePusher();
public:
    virtual void StartDataCollection(StartMessage& message) = 0;
    virtual bool EndDataCollection(const EndMessage& message) = 0; // Non-blocking
    virtual bool SendImage(const uint8_t *image_data, size_t image_size,  int64_t image_number) = 0;
    // Returns true if the image was accepted (handed to the writer / enqueued for
    // sending), false if it was dropped (e.g. no TCP connection, or the enqueue
    // deadline expired). Ownership is settled either way: on a false return the
    // slot has already been release()d. The authoritative count of images actually
    // delivered is GetImagesWritten() (ACK-based for TCP).
    virtual bool SendImage(ZeroCopyReturnValue &z);
    virtual bool SendCalibration(const CompressedImage& message) = 0;
    virtual std::string Finalize(); // Ensure that all streams are closed, can throw exception
    virtual std::string GetWriterNotificationSocketAddress() const;
    virtual ~ImagePusher() = default;
    virtual std::string PrintSetup() const = 0;

    virtual std::vector<std::string> GetAddress() const { return {}; }
    virtual std::optional<uint64_t> GetImagesWritten() const { return std::nullopt; }
    virtual std::optional<uint64_t> GetImagesWriteError() const { return std::nullopt; }
    virtual size_t GetConnectedWriters() const { return 0; };
    virtual ImagePusherType GetType() const = 0;
    virtual std::vector<int64_t> GetWriterFifoUtilization() const { return {}; }

    ImagePusherStatus GetStatus() const {
        return ImagePusherStatus{
            GetType(),
            GetAddress(),
            GetConnectedWriters(),
            GetImagesWritten(),
            GetImagesWriteError(),
            GetWriterFifoUtilization()
        };
    }
};