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

#include <catch2/catch_all.hpp>

#include "../common/DiffractionExperiment.h"
#include "../writer/FileWriter.h"
#include "../xds-plugin/plugin.h"

namespace {
    void CleanupFiles(const std::string& prefix, int file_count) {
        remove((prefix + "_master.h5").c_str());
        for (int i = 1; i <= file_count; ++i) {
            char data_file[256];
            snprintf(data_file, sizeof(data_file), "%s_data_%06d.h5", prefix.c_str(), i);
            remove(data_file);
        }
    }

    void ClosePlugin() {
        int error_flag = 0;
        plugin_close(&error_flag);
        REQUIRE(error_flag == 0);
    }
}

TEST_CASE("XDSPlugin_GetData_VDS", "[HDF5][XDS][Plugin]") {
    DiffractionExperiment x(DetJF(1));

    x.FilePrefix("test_xds_vds").ImagesPerTrigger(3).ImagesPerFile(1).OverwriteExistingFiles(true);
    x.SetFileWriterFormat(FileWriterFormat::NXmxVDS);
    x.BitDepthImage(16).PixelSigned(false);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
        .IncidentEnergy_keV(WVL_1A_IN_KEV)
        .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    RegisterHDF5Filter();

    std::vector<uint32_t> pixel_mask(x.GetPixelsNum(), 0);
    std::vector<uint16_t> image(x.GetPixelsNum(), 7);
    image[0] = UINT16_MAX;
    image[1] = 123;

    {
        StartMessage start_message;
        x.FillMessage(start_message);
        start_message.pixel_mask["default"] = pixel_mask;

        FileWriter file_set(start_message);

        for (int i = 0; i < 3; ++i) {
            DataMessage message{};
            image[5678] = static_cast<uint16_t>(200 + i);
            message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
            message.number = i;
            REQUIRE_NOTHROW(file_set.WriteHDF5(message));
        }

        EndMessage end_message;
        end_message.max_image_number = 3;
        file_set.WriteHDF5(end_message);
        file_set.Finalize();
    }

    {
        int info[1024] = {};
        int error_flag = 0;

        REQUIRE_NOTHROW(plugin_open("test_xds_vds_master.h5", info, &error_flag));
        REQUIRE(error_flag == 0);

        int nx = 0, ny = 0, nbytes = 0, number_of_frames = 0;
        float qx = 0.0f, qy = 0.0f;

        REQUIRE_NOTHROW(plugin_get_header(&nx, &ny, &nbytes, &qx, &qy, &number_of_frames, info, &error_flag));
        REQUIRE(error_flag == 0);

        CHECK(nx == static_cast<int>(x.GetXPixelsNum()));
        CHECK(ny == static_cast<int>(x.GetYPixelsNum()));
        CHECK(nbytes == 2);
        CHECK(number_of_frames == 3);

        std::vector<int> data(static_cast<size_t>(nx) * static_cast<size_t>(ny), 0);

        int frame_number = 2;
        REQUIRE_NOTHROW(plugin_get_data(&frame_number, &nx, &ny, data.data(), info, &error_flag));
        REQUIRE(error_flag == 0);

        CHECK(data[0] == -1);
        CHECK(data[1] == 123);
        CHECK(data[5678] == 201);

        ClosePlugin();
    }

    CleanupFiles("test_xds_vds", 3);
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("XDSPlugin_GetData_Legacy", "[HDF5][XDS][Plugin]") {
    DiffractionExperiment x(DetJF(1));

    x.FilePrefix("test_xds_legacy").ImagesPerTrigger(2).ImagesPerFile(1).OverwriteExistingFiles(true);
    x.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);
    x.BitDepthImage(16).PixelSigned(true);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
        .IncidentEnergy_keV(WVL_1A_IN_KEV)
        .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    RegisterHDF5Filter();

    std::vector<uint32_t> pixel_mask(x.GetPixelsNum(), 0);
    std::vector<int16_t> image(x.GetPixelsNum(), 0);
    image[0] = INT16_MAX;
    image[1] = INT16_MIN;
    image[2] = 456;

    {
        StartMessage start_message;
        x.FillMessage(start_message);
        start_message.pixel_mask["default"] = pixel_mask;

        FileWriter file_set(start_message);

        for (int i = 0; i < 2; ++i) {
            DataMessage message{};
            image[5678] = static_cast<int16_t>(10 + i);
            message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
            message.number = i;
            REQUIRE_NOTHROW(file_set.WriteHDF5(message));
        }

        EndMessage end_message;
        end_message.max_image_number = 2;
        file_set.WriteHDF5(end_message);
        file_set.Finalize();
    }

    {
        int info[1024] = {};
        int error_flag = 0;

        REQUIRE_NOTHROW(plugin_open("test_xds_legacy_master.h5", info, &error_flag));
        REQUIRE(error_flag == 0);

        int nx = 0, ny = 0, nbytes = 0, number_of_frames = 0;
        float qx = 0.0f, qy = 0.0f;

        REQUIRE_NOTHROW(plugin_get_header(&nx, &ny, &nbytes, &qx, &qy, &number_of_frames, info, &error_flag));
        REQUIRE(error_flag == 0);
        REQUIRE(number_of_frames == 2);

        std::vector<int> data(static_cast<size_t>(nx) * static_cast<size_t>(ny), 0);

        int frame_number = 1;
        REQUIRE_NOTHROW(plugin_get_data(&frame_number, &nx, &ny, data.data(), info, &error_flag));
        REQUIRE(error_flag == 0);

        CHECK(data[0] == -1);
        CHECK(data[1] == -1);
        CHECK(data[2] == 456);
        CHECK(data[5678] == 10);

        frame_number = 2;
        REQUIRE_NOTHROW(plugin_get_data(&frame_number, &nx, &ny, data.data(), info, &error_flag));
        REQUIRE(error_flag == 0);
        CHECK(data[5678] == 11);

        ClosePlugin();
    }

    CleanupFiles("test_xds_legacy", 2);
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("XDSPlugin_GetData_Integrated", "[HDF5][XDS][Plugin]") {
    DiffractionExperiment x(DetJF(1));

    x.FilePrefix("test_xds_integrated").ImagesPerTrigger(3).OverwriteExistingFiles(true);
    x.SetFileWriterFormat(FileWriterFormat::NXmxIntegrated);
    x.BitDepthImage(16).PixelSigned(false);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
        .IncidentEnergy_keV(WVL_1A_IN_KEV)
        .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    x.Compression(CompressionAlgorithm::NO_COMPRESSION);

    RegisterHDF5Filter();

    AzimuthalIntegrationSettings azint_settings;
    azint_settings.AzimuthalBinCount(4);
    x.ImportAzimuthalIntegrationSettings(azint_settings);

    std::vector<uint16_t> image(x.GetPixelsNum(), 0);
    image[0] = UINT16_MAX;
    image[1] = 321;
    image[5678] = 777;

    {
        StartMessage start_message;
        x.FillMessage(start_message);
        start_message.pixel_mask["default"] = std::vector<uint32_t>(x.GetPixelsNum(), 0);

        FileWriter file_set(start_message);

        for (int i = 0; i < 3; ++i) {
            DataMessage message{};
            image[5678] = static_cast<uint16_t>(777 + i);
            message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
            message.number = i;
            message.az_int_profile = std::vector<float>(azint_settings.GetBinCount(), static_cast<float>(10 + i));
            REQUIRE_NOTHROW(file_set.WriteHDF5(message));
        }

        EndMessage end_message;
        end_message.max_image_number = 3;
        file_set.WriteHDF5(end_message);
        file_set.Finalize();
    }

    {
        int info[1024] = {};
        int error_flag = 0;

        REQUIRE_NOTHROW(plugin_open("test_xds_integrated_master.h5", info, &error_flag));
        REQUIRE(error_flag == 0);

        int nx = 0, ny = 0, nbytes = 0, number_of_frames = 0;
        float qx = 0.0f, qy = 0.0f;

        REQUIRE_NOTHROW(plugin_get_header(&nx, &ny, &nbytes, &qx, &qy, &number_of_frames, info, &error_flag));
        REQUIRE(error_flag == 0);

        CHECK(nx == static_cast<int>(x.GetXPixelsNum()));
        CHECK(ny == static_cast<int>(x.GetYPixelsNum()));
        CHECK(qx == Catch::Approx(0.075));
        CHECK(qy == Catch::Approx(0.075));
        CHECK(nbytes == 2);
        CHECK(number_of_frames == 3);

        std::vector<int> data(static_cast<size_t>(nx) * static_cast<size_t>(ny), 0);

        int frame_number = 3;
        REQUIRE_NOTHROW(plugin_get_data(&frame_number, &nx, &ny, data.data(), info, &error_flag));
        REQUIRE(error_flag == 0);

        CHECK(data[0] == -1);
        CHECK(data[1] == 321);
        CHECK(data[5678] == 779);

        ClosePlugin();
    }

    remove("test_xds_integrated_master.h5");
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}