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

#ifndef JUNGFRAUJOCH_IMAGEBUFFER_H
#define JUNGFRAUJOCH_IMAGEBUFFER_H

#include "ThreadSafeFIFO.h"
#include "ZeroCopyReturnValue.h"
#include "DiffractionExperiment.h"

struct ImageBufferStatus {
    std::vector<int64_t> images_in_the_buffer;
    int64_t min_image_number;
    int64_t max_image_number;

    int64_t total_slots;
    int64_t available_slots;
    std::optional<int64_t> current_counter;
};

struct ImageBufferEntry {
    int64_t image_number = -1;
    size_t image_size = 0;
    int64_t readers = 0;
    bool zeromq_processing = false;
    bool indexed = false;
};

class ImageBuffer {
    mutable std::mutex m;
    std::condition_variable cv_zeromq_done;
    std::condition_variable cv_preview_done;

    uint8_t *buffer;
    const size_t buffer_size;

    size_t slot_size = 0;
    int64_t slot_count = 0;

    std::queue<uint32_t> free_slots;
    std::vector<ImageBufferEntry> v;
    std::vector<ZeroCopyReturnValue> send_buffer_zero_copy_ret_val;
    int64_t zeromq_in_process = 0;
    int64_t preview_in_process = 0;

    bool disable_requests = true; // All requests fail in this state

    int64_t counter = 0;

    std::mutex start_message_mutex;
    std::vector<uint8_t> start_message;

    // Assumes existing lock
    bool FinalizeInternal(std::unique_lock<std::mutex> &ul,
                          std::chrono::microseconds timeout = std::chrono::milliseconds(2500));

    std::optional<uint32_t> getHandle(int64_t image_number) const;
    std::optional<uint32_t> getHandleMaxImage() const;
    std::optional<uint32_t> getHandleMaxIndexedImage() const;

    uint8_t *GetBufferLocation(size_t id); // image metadata + image
public:
    constexpr static int64_t MaxImage = -1;
    constexpr static int64_t MaxIndexedImage = -2;

    explicit ImageBuffer(size_t buffer_size_bytes);
    ~ImageBuffer();
    ImageBuffer(const ImageBuffer &) = delete;
    ImageBuffer &operator=(const ImageBuffer &) = delete;

    void StartMeasurement(size_t buffer_location_size);
    void StartMeasurement(const DiffractionExperiment& experiment);

    ZeroCopyReturnValue *GetImageSlot();
    void ReleaseSlot(uint32_t location, int64_t image_number, size_t image_size, bool indexed);

    int64_t GetAvailSlots() const;

    bool CheckIfBufferReturned(std::chrono::microseconds timeout);
    bool Finalize(std::chrono::microseconds timeout);

    bool GetImage(std::vector<uint8_t> &out_v, int64_t image_number = MaxImage);

    void SaveStartMessage(const std::vector<uint8_t> &msg);
    void GetStartMessage(std::vector<uint8_t> &out_v);

    ImageBufferStatus GetStatus() const;
};

#endif //JUNGFRAUJOCH_IMAGEBUFFER_H