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

#ifndef JUNGFRAUJOCH_STREAMWRITER_H
#define JUNGFRAUJOCH_STREAMWRITER_H

#include <future>

#include "../common/Logger.h"
#include "../image_puller/ImagePuller.h"
#include "HDF5DataFile.h"
#include "FileWriter.h"

enum class StreamWriterState {Idle, Started, Receiving, Error, Finalized};

struct StreamWriterStatistics {
    uint64_t processed_images;
    float performance_MBs;
    float performance_Hz;
    std::string file_prefix;
    std::string run_name;
    uint64_t run_number;
    uint64_t socket_number;
    StreamWriterState state;
    size_t max_puller_fifo_utilization;
};

struct StreamWriterOutput {
    StreamWriterStatistics stats;
    std::vector<HDF5DataFileStatistics> data_file_stats;
};

class StreamWriter {
    std::atomic<bool> abort = false;
    const bool verbose;
    std::string file_done_address;

    StreamWriterState state = StreamWriterState::Idle;

    ImagePullerOutput image_puller_output;
    std::unique_ptr<FileWriter> file_writer;

    std::atomic<uint64_t> processed_images;
    std::atomic<uint64_t> processed_image_size;
    std::chrono::time_point<std::chrono::system_clock> start_time;
    std::chrono::time_point<std::chrono::system_clock> end_time;
    std::string file_prefix;
    std::string run_name;
    uint64_t run_number;
    uint64_t socket_number;
    uint64_t max_image_number;
    std::string writer_notification_zmq_addr;
    std::string err;
    bool mute_data_msg_in_idle = false;
    std::vector<HDF5DataFileStatistics> hdf5_data_file_statistics;

    bool debug_skip_write_notification = false;
    bool tcp_data_fatal_sent = false;

    ImagePuller &image_puller;
    Logger &logger;
    void CollectImages();
    bool WaitForImage();
    void NotifyReceiverOnFinalizedWrite(const std::string &detector_update_zmq_addr);
    void NotifyTcpAck(TCPFrameType ack_for, bool ok, bool fatal, TCPAckCode code, const std::string &error_text = "");
    void ProcessStartMessage();
    void ProcessEndMessage();
    void ProcessDataImage();
    void ProcessCalibrationImage();
    void FinalizeDataCollection();
public:
    StreamWriter(Logger &in_logger,
                 ImagePuller &in_image_puller,
                 std::string in_file_done_address = "",
                 bool verbose = false);
    StreamWriterOutput Run();
    void Cancel();
    StreamWriterStatistics GetStatistics() const;
    void DebugSkipWriteNotification(bool input);
};


#endif //JUNGFRAUJOCH_STREAMWRITER_H