// Copyright (2019-2023) Paul Scherrer Institute

#include "HDF5Writer.h"
#include "HDF5NXmx.h"

HDF5Writer::HDF5Writer(const StartMessage &request)
: images_per_file(request.images_per_file),
file_prefix(request.file_prefix),
max_spot_count(request.max_spot_count),
az_int_bin_to_q(request.az_int_bin_to_q) {}

void HDF5Writer::Write(const DataMessage& message) {
    std::lock_guard<std::mutex> lock(hdf5_mutex);
    if (message.image.size == 0)
        return;

    if (message.number < 0)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "No support for negative images");

    size_t file_number = 0;
    size_t image_number = message.number;
    if (images_per_file > 0) {
        file_number = message.number / images_per_file;
        image_number = message.number % images_per_file;
    }
    if (files.size() <= file_number)
        files.resize(file_number + 1);

    if (!files[file_number])
        files[file_number] = std::make_unique<HDF5DataFile>(HDF5Metadata::DataFileName(file_prefix, file_number),
                                                  az_int_bin_to_q,
                                                  file_number * images_per_file,
                                                  max_spot_count);
    // Ignore zero size images
    if (message.image.size > 0)
        files[file_number]->Write(message, image_number);

    if (files[file_number]->GetNumImages() == images_per_file) {
        stats.emplace_back(files[file_number]->GetStatistics());
        files[file_number].reset();
    }
}

std::vector<HDF5DataFileStatistics> HDF5Writer::Finalize() {
    std::lock_guard<std::mutex> lock(hdf5_mutex);
    for (auto &f: files) {
        if (f) {
            auto tmp = f->GetStatistics();
            if (tmp.total_images > 0)
                stats.push_back(tmp);
            f.reset();
        }
    }
    return stats;
}