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

// Based on Durin plugin code from Diamond Light Source Ltd. and Global Phasing (BSD-3 license)

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <vector>

#include "../writer/HDF5Objects.h"

#include "plugin.h"

#include "JFJochDecompress.h"
#include "../common/GitInfo.h"
#include "../common/JFJochMessages.h"

namespace {
    class PluginError : public std::runtime_error {
        int err_code;
    public:
        PluginError(int code, const std::string &message) : std::runtime_error(message), err_code(code) {}
    };

    std::shared_mutex plugin_mutex;
    std::unique_ptr<HDF5ReadOnlyFile> hdf5_file;
    FileWriterFormat format = FileWriterFormat::NoFile;
    uint64_t images_per_file = 0;
    size_t image_size_x, image_size_y;
    float pixel_size_x, pixel_size_y;
    size_t total_image_number;
    uint32_t pixel_byte_depth = 0;
    bool pixel_signed = false;

    std::vector<uint8_t> one_byte_mask;

    template<class T>
    void ConvertAndMaskTyped(const std::vector<uint8_t> &in_8bit,
                             int marker_value,
                             int *out) {
        auto in = reinterpret_cast<const T *>(in_8bit.data());
        size_t size = in_8bit.size() / sizeof(T);

        for (size_t i = 0; i < size; ++i) {
            if (!one_byte_mask.empty() && (one_byte_mask.at(i) == 1))
                out[i] = -1;
            else if (!one_byte_mask.empty() && (one_byte_mask.at(i) == 2))
                out[i] = -2;
            else if (marker_value != 0 && in[i] == static_cast<T>(marker_value))
                out[i] = -1;
            else if (std::is_signed_v<T> && in[i] < 0)
                out[i] = -1;
            else if (in[i] > INT32_MAX)
                out[i] = INT32_MAX;
            else
                out[i] = static_cast<int>(in[i]);
        }
    }

    void LoadDataset() {
        std::string dataset_name;
        if (hdf5_file->Exists("/entry/data/data")) {
            format = FileWriterFormat::NXmxVDS;
            dataset_name = "/entry/data/data";
        } else if (hdf5_file->Exists("/entry/data/data_000001")) {
            format = FileWriterFormat::NXmxLegacy;
            dataset_name = "/entry/data/data_000001";
        } else {
            throw PluginError(-1, "Could not locate detector dataset");
        }

        HDF5DataSet data(*hdf5_file, dataset_name);
        HDF5DataSpace dataspace(data);
        HDF5DataType datatype(data);

        auto dim = dataspace.GetDimensions();
        if (dim.size() != 3)
            throw PluginError(-1, "Wrong dimension of /entry/data/data");
        image_size_x = dim[2];
        image_size_y = dim[1];
        images_per_file = dim[0];
        total_image_number = dim[0];

        if (format == FileWriterFormat::NXmxLegacy) {
            int dataset = 2;
            while (dataset < 100000) {
                char name[255];
                snprintf(name, sizeof(name), "/entry/data/data_%06d", dataset);

                if (!hdf5_file->Exists(name))
                    break;
                auto leg_dims = hdf5_file->GetDimension(name);
                if (leg_dims.size() != 3)
                    throw PluginError(-1, "Wrong dimension of " + std::string(name));
                if (leg_dims[2] != image_size_x)
                    throw PluginError(-1, "Image size of " + std::string(name) + " does not match");
                if (leg_dims[1] != image_size_y)
                    throw PluginError(-1, "Image size of " + std::string(name) + " does not match");
                total_image_number += leg_dims[0];
                dataset++;
            }
        }

        if (!datatype.IsInteger())
            throw PluginError(-1, "Data type of /entry/data/data is not integer");
        pixel_signed = datatype.IsSigned();
        pixel_byte_depth = datatype.GetElemSize();

        pixel_size_x = hdf5_file->GetFloat("/entry/instrument/detector/x_pixel_size");
        pixel_size_y = hdf5_file->GetFloat("/entry/instrument/detector/y_pixel_size");

        one_byte_mask = std::vector<uint8_t>(image_size_x * image_size_y, 0);

        auto mask_tmp = hdf5_file->ReadVector<uint32_t>(
                "/entry/instrument/detector/pixel_mask",
                {0, 0},
                {image_size_y, image_size_x}
        );

        for (int i = 0; i < mask_tmp.size(); i++) {
            if (mask_tmp[i] & (1 << 30))
                one_byte_mask[i] = 2;
            else if (mask_tmp[i] & 0xFFFF)
                one_byte_mask[i] = 1;
        }
    }

    void ConvertToIntAndMask(const std::vector<uint8_t> &in_8bit, int *out_buffer) {
        switch (pixel_byte_depth) {
            case 1:
                if (pixel_signed)
                    return ConvertAndMaskTyped<int8_t>(in_8bit, INT8_MAX, out_buffer);
                else
                    return ConvertAndMaskTyped<uint8_t>(in_8bit, UINT8_MAX, out_buffer);
            case 2:
                if (pixel_signed)
                    return ConvertAndMaskTyped<int16_t>(in_8bit, INT16_MAX, out_buffer);
                else
                    return ConvertAndMaskTyped<uint16_t>(in_8bit, UINT16_MAX, out_buffer);
            case 4:
                if (pixel_signed)
                    return ConvertAndMaskTyped<int32_t>(in_8bit, INT32_MAX, out_buffer);
                else
                    return ConvertAndMaskTyped<uint32_t>(in_8bit, UINT32_MAX, out_buffer);
            default:
                throw PluginError(-1, "Unsupported conversion to int");
        }
    }

    void FillInfoArray(int info[1024]) {
        info[0] = 0x01;
        info[1] = VERSION_MAJOR;
        info[2] = VERSION_MINOR;
        info[3] = VERSION_PATCH;
        info[4] = VERSION_TIMESTAMP;
        info[5] = 0;
        info[6] = -1;
    }
} // namespace

extern "C" {
void plugin_open(const char *filename, int info[1024], int *error_flag) {
    std::unique_lock sl(plugin_mutex);

    std::cout << "********** Jungfraujoch XDS plugin **********" << std::endl;
    std::cout << "Jungfraujoch version " << jfjoch_version() << std::endl;
    std::cout << "Plugin version " << VERSION_MAJOR << "." << VERSION_MINOR << "." << VERSION_PATCH << std::endl << std::endl;
    std::cout << "Copyright (C) 2024-2026 Paul Scherrer Institute" << std::endl;
    std::cout << "This program comes with ABSOLUTELY NO WARRANTY" << std::endl;
    std::cout << "This is free software, and you are welcome to redistribute it" << std::endl;
    std::cout << "under certain conditions (GPLv3)" << std::endl << std::endl;
    std::cout << "Based on durin plugin from Diamond Light Source Ltd. with modification from the Global Phasing Ltd." << std::endl;
    std::cout << "(BSD-3 license)" << std::endl << std::endl;
    try {
        FillInfoArray(info);

        RegisterHDF5Filter();
        hdf5_file = std::make_unique<HDF5ReadOnlyFile>(filename);
        LoadDataset();
        *error_flag = 0;
    } catch (std::exception &e) {
        std::cerr << e.what() << std::endl;
        *error_flag = -1;
        hdf5_file.reset();
        total_image_number = 0;
    }
}

void plugin_get_header(int *nx, int *ny, int *nbytes, float *qx, float *qy,
                       int *number_of_frames, int info[1024], int *error_flag) {
    std::shared_lock sl(plugin_mutex);

    try {
        FillInfoArray(info);

        if (!hdf5_file)
            throw PluginError(-1, "HDF5 file not open");

        *nx = image_size_x;
        *ny = image_size_y;
        *nbytes = pixel_byte_depth;
        *number_of_frames = total_image_number;
        *qx = pixel_size_x * 1e3; // mm
        *qy = pixel_size_y * 1e3; // mm
        *error_flag = 0;
    } catch (std::exception &e) {
        std::cerr << e.what() << std::endl;
        *error_flag = -1;
    }
}

void plugin_get_data(int *frame_number, int *nx, int *ny, int *data_array,
                     int info[1024], int *error_flag) {
    std::shared_lock sl(plugin_mutex);

    try {
        if (!hdf5_file)
            throw PluginError(-1, "HDF5 file not open");

        FillInfoArray(info);

        if (*frame_number <= 0 || *frame_number > total_image_number)
            throw PluginError(-1, "Frame number out of range");

        std::string dataset_name;
        hsize_t image_id;

        if (format == FileWriterFormat::NXmxLegacy) {
            char str[256];
            size_t dataset_index = (*frame_number - 1) / images_per_file + 1;
            snprintf(str, sizeof(str), "/entry/data/data_%06ld", dataset_index);
            dataset_name = std::string(str);
            image_id = (*frame_number - 1) % images_per_file;
        } else {
            dataset_name = "/entry/data/data";
            image_id = *frame_number - 1;
        }

        HDF5DataSet dataset(*hdf5_file, dataset_name);
        HDF5Dcpl dcpl(dataset);

        std::vector<uint8_t> tmp;
        std::vector<hsize_t> start = {image_id, 0, 0};

        CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION;
        auto chunk_size = dcpl.GetChunking();

        if ((chunk_size.size() == 3)
            && (chunk_size[0] == 1)
            && (chunk_size[1] == image_size_y)
            && (chunk_size[2] == image_size_x)) {
            dataset.ReadDirectChunk(tmp, start);
            algorithm = dcpl.GetCompression();
            } else {
                dataset.ReadVectorToU8(tmp, start, {1, image_size_y, image_size_x});
                algorithm = CompressionAlgorithm::NO_COMPRESSION;
            }

        if (algorithm != CompressionAlgorithm::NO_COMPRESSION) {
            std::vector<uint8_t> decompressed_image(image_size_x * image_size_y * pixel_byte_depth);
            JFJochDecompressPtr(decompressed_image.data(),
                                algorithm,
                                tmp.data(),
                                tmp.size(),
                                image_size_x * image_size_y,
                                pixel_byte_depth);
            ConvertToIntAndMask(decompressed_image, data_array);
        } else {
            ConvertToIntAndMask(tmp, data_array);
        }
        *error_flag = 0;
    } catch (std::exception &e) {
        std::cerr << e.what() << std::endl;
        *error_flag = -1;
    }
}

void plugin_close(int *error_flag) {
    std::unique_lock sl(plugin_mutex);
    try {
        hdf5_file.reset();
        one_byte_mask.clear();
        total_image_number = 0;
        format = FileWriterFormat::NoFile;
        images_per_file = 0;
        image_size_x = 0;
        image_size_y = 0;
        pixel_size_x = 0;
        pixel_size_y = 0;
        *error_flag = 0;
    } catch (std::exception &e) {
        std::cerr << e.what() << std::endl;
        *error_flag = -1;
    }
}
} /* extern "C" */