// Copyright (2019-2023) Paul Scherrer Institute

#include <sys/stat.h>
#include "HDF5DataFile.h"
#include "../compression/JFJochCompressor.h"

HDF5DataFile::HDF5DataFile(const std::string &in_filename,
                           const std::vector<float>& in_rad_int_bin_to_q,
                           size_t in_max_spots) {
    xpixel = 0;
    ypixel = 0;
    max_image_number = 0;
    nimages = 0;
    filename = in_filename;

    plugins.emplace_back(std::make_unique<HDF5DataFilePluginJUNGFRAU>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginAzInt>(in_rad_int_bin_to_q));
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginXFEL>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginMX>(in_max_spots));
}

HDF5DataFile::~HDF5DataFile() {

    if (data_file) {
        HDF5Group group_exp(*data_file, "/entry/detector");
        group_exp.NXClass("NXcollection");

        group_exp.SaveVector("timestamp", timestamp);
        group_exp.SaveVector("exptime", exptime);
    }
}

void HDF5DataFile::CreateFile(const DataMessage& msg) {
    SetupSWMRFile(msg);

    data_file = std::make_unique<HDF5File>(filename, false, true);
    data_set = std::make_unique<HDF5DataSet>(*data_file, "/entry/data/data");
    for (auto &p: plugins)
        p->OpenFile(*data_file, msg);
}

void HDF5DataFile::Write(const DataMessage &msg, uint64_t image_number) {
    std::lock_guard<std::mutex> lock(hdf5_mutex);
    bool new_file = false;

    if (!data_file) {
        CreateFile(msg);
        new_file = true;
    }

    if (new_file || (static_cast<int64_t>(image_number) > max_image_number)) {
        max_image_number = image_number;
        data_set->SetExtent({max_image_number+1, ypixel, xpixel});
        timestamp.resize(max_image_number + 1);
        exptime.resize(max_image_number + 1);
    }

    nimages++;
    data_set->WriteDirectChunk(msg.image.data, msg.image.size, {image_number, 0, 0});

    for (auto &p: plugins)
        p->Write(msg);

    timestamp[image_number] = msg.timestamp;
    exptime[image_number] = msg.exptime;

    auto now_time = std::chrono::system_clock::now();
    auto time_diff = now_time - (last_flush + swmr_flush_period);
    if (time_diff.count() >= 0) {
        data_set->Flush();

        for (auto &p: plugins)
            p->Flush();

        last_flush = now_time;
    }
}

HDF5DataFileStatistics HDF5DataFile::GetStatistics() const {
    HDF5DataFileStatistics ret;
    ret.max_image_number = max_image_number;
    ret.total_images = nimages;
    ret.filename = filename;
    return ret;
}

size_t HDF5DataFile::GetNumImages() const {
    return nimages;
}

void HDF5DataFile::SetupSWMRFile(const DataMessage& msg) {
    HDF5Dcpl dcpl;
    HDF5DataType data_type(msg.image.pixel_depth_bytes, msg.image.pixel_is_signed);

    xpixel = msg.image.xpixel;
    ypixel = msg.image.ypixel;

    dcpl.SetCompression(msg.image.algorithm, msg.image.pixel_depth_bytes, JFJochBitShuffleCompressor::DefaultBlockSize);
    dcpl.SetChunking( {1, ypixel, xpixel});
    if (msg.image.pixel_is_signed) {
        if (msg.image.pixel_depth_bytes == 2)
            dcpl.SetFillValue16(INT16_MIN);
        else
            dcpl.SetFillValue32(INT32_MIN);
    }

    HDF5File data_file_local(filename, true, true);
    chmod(filename.c_str(), S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP); // default permissions

    HDF5Group(data_file_local, "/entry").NXClass("NXentry");
    HDF5Group(data_file_local, "/entry/data").NXClass("NXdata");

    HDF5DataSpace data_space({1, ypixel, xpixel}, {H5S_UNLIMITED, ypixel, xpixel});
    HDF5DataSet(data_file_local, "/entry/data/data", data_type, data_space, dcpl);

    for (auto &p: plugins)
        p->SetupSWMRFile(data_file_local, msg);

    last_flush = std::chrono::system_clock::now();
}