// Copyright (2019-2023) Paul Scherrer Institute

#include "HDF5Writer.h"
#include "HDF5NXmx.h"
#include <nlohmann/json.hpp>

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),
          user_data(request.user_data) {}

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)
        AddStats(files[file_number]->Close());
}

std::vector<HDF5DataFileStatistics> HDF5Writer::Finalize() {
    std::lock_guard<std::mutex> lock(hdf5_mutex);
    for (auto &f: files) {
        if (f)
            AddStats(f->Close());
    }
    return stats;
}

void HDF5Writer::AddStats(const std::optional<HDF5DataFileStatistics>& s) {
    if (!s)
        return;

    stats.push_back(*s);
    if (socket) {
        nlohmann::json j;
        j["filename"] = s->filename;
        j["nimages"] = s->total_images;
        if (!user_data.empty()) {
            nlohmann::json j_userdata;

            // if user_data is valid json, interpret it as such, otherwise embed as string
            try {
                j_userdata = nlohmann::json::parse(user_data);
            } catch (...) {
                j_userdata = user_data;
            }
            j["user_data"] = j_userdata;
        }
        socket->Send(j.dump());
    }
}

void HDF5Writer::SetupSocket(const std::string &addr) {
    socket = std::make_unique<ZMQSocket>(ZMQSocketType::Pub);
    socket->Bind(addr);
}

std::optional<std::string> HDF5Writer::GetZMQAddr() {
    if (socket) {
        return socket->GetEndpointName();
    } else
        return {};
}