// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "../common/JFJochException.h"
#include "StreamWriter.h"
#include "HDF5Writer.h"
#include "HDF5NXmx.h"
#include "MakeDirectory.h"

StreamWriter::StreamWriter(ZMQContext &context, Logger &in_logger, const std::string &zmq_addr)
: image_puller(context), logger(in_logger) {
    image_puller.Connect(zmq_addr);
    logger.Info("Connected via ZMQ to {}", zmq_addr);
}

void StreamWriter::StartDataCollection() {
    image_puller.WaitForImage();
    while (image_puller.GetFrameType() != JFJochFrameDeserializer::Type::START) {
        logger.Error("Expected START image");
        image_puller.WaitForImage();
    }
    start_message = image_puller.GetStartMessage();
    logger.Info("Starting writing for dataset {} of {} images", start_message.file_prefix,
                start_message.number_of_images);

    MakeDirectory(start_message.file_prefix);
}

void StreamWriter::CollectImages(std::vector<HDF5DataFileStatistics> &v) {
    HDF5Writer writer(start_message);
    image_puller.WaitForImage();
    while (image_puller.GetFrameType() == JFJochFrameDeserializer::Type::IMAGE) {
        auto image_array = image_puller.GetDataMessage();
        writer.Write(image_array);
        image_puller.WaitForImage();
    }
    writer.GetStatistics(v);
}

void StreamWriter::EndDataCollection() {
    while (image_puller.GetFrameType() != JFJochFrameDeserializer::Type::END) {
        logger.Error("Expected END image");
        image_puller.WaitForImage();
    }
    EndMessage end_message = image_puller.GetEndMessage();
    if (end_message.write_master_file)
        HDF5Metadata::NXmx(start_message, end_message);
}


void StreamWriter::Abort() {
    image_puller.Abort();
}

StreamWriterStatistics  StreamWriter::Run() {
    StreamWriterStatistics ret;
    StartDataCollection();
    try {
        CollectImages(ret.data_file_stats);
    } catch (JFJochException &e) {
        // Error during collecting images will skip to end data collection
        // End data collection will consume all images till the end
        logger.ErrorException(e);
    }
    EndDataCollection();

    ret.image_puller_stats = image_puller.GetStatistics();
    logger.Info("Write task done. Images = {} Throughput = {:.0f} MB/s Frame rate = {:.0f} Hz",
                ret.image_puller_stats.processed_images, ret.image_puller_stats.performance_MBs, ret.image_puller_stats.performance_Hz);
    return ret;
}