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

#include <catch2/catch_all.hpp>
#include <filesystem>

#include "../common/DiffractionExperiment.h"
#include "../common/ScanResultGenerator.h"
#include "../writer/FileWriter.h"
#include "../reader/JFJochHDF5Reader.h"
#include "../process/JFJochProcess.h"
#include "../process/JFJochProcessCommandLine.h"

namespace {
    // Write a small VDS dataset of `n` flat images and return nothing (prefix_master.h5 +
    // prefix_data_000001.h5 land in the test working directory).
    void WriteTestDataset(const std::string &prefix, int n) {
        RegisterHDF5Filter();

        DiffractionExperiment x(DetJF(1));
        x.FilePrefix(prefix).ImagesPerTrigger(n).OverwriteExistingFiles(true);
        x.BitDepthImage(16).ImagesPerFile(n).SetFileWriterFormat(FileWriterFormat::NXmxVDS).PixelSigned(true);
        x.Compression(CompressionAlgorithm::NO_COMPRESSION);
        x.BeamX_pxl(512).BeamY_pxl(256).DetectorDistance_mm(150).IncidentEnergy_keV(WVL_1A_IN_KEV)
                .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));

        std::vector<int16_t> image(x.GetPixelsNum(), 5);
        StartMessage start_message;
        x.FillMessage(start_message);
        FileWriter file_set(start_message);
        ScanResultGenerator generator(x);
        for (int i = 0; i < n; i++) {
            DataMessage message{};
            message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
            message.number = i;
            REQUIRE_NOTHROW(file_set.WriteHDF5(message));
            generator.Add(message);
        }
        EndMessage end_message;
        end_message.max_image_number = n;
        generator.FillEndMessage(end_message);
        file_set.WriteHDF5(end_message);
        file_set.Finalize();
    }
}

TEST_CASE("JFJochProcess_AzInt", "[HDF5][Full]") {
    WriteTestDataset("process_azint_in", 8);

    JFJochHDF5Reader reader;
    REQUIRE_NOTHROW(reader.ReadFile("process_azint_in_master.h5"));
    auto dataset = reader.GetDataset();
    REQUIRE(dataset);

    ProcessConfig config;
    config.mode = ProcessMode::AzimuthalIntegration;
    config.nthreads = 2;
    config.output_prefix = "process_azint_out";

    JFJochProcess process(reader, dataset->experiment, dataset->pixel_mask, config);
    ProcessResult result;
    REQUIRE_NOTHROW(result = process.Run());

    CHECK_FALSE(result.cancelled);
    CHECK(result.images_processed == 8);
    REQUIRE(result.written_master_path.has_value());

    {
        // The _process.h5 links back to the source images and carries an azimuthal profile per image.
        JFJochHDF5Reader out;
        REQUIRE_NOTHROW(out.ReadFile("process_azint_out_process.h5"));
        CHECK(out.GetNumberOfImages() == 8);
        std::shared_ptr<JFJochReaderImage> img;
        REQUIRE_NOTHROW(img = out.LoadImage(0));
        REQUIRE(img);
        CHECK_FALSE(img->ImageData().az_int_profile.empty());
    }

    reader.Close();
    remove("process_azint_in_master.h5");
    remove("process_azint_in_data_000001.h5");
    remove("process_azint_out_process.h5");
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochProcess_NoOutput", "[HDF5][Full]") {
    WriteTestDataset("process_noout_in", 6);

    JFJochHDF5Reader reader;
    REQUIRE_NOTHROW(reader.ReadFile("process_noout_in_master.h5"));
    auto dataset = reader.GetDataset();

    // Empty output prefix => process without writing any file.
    ProcessConfig config;
    config.mode = ProcessMode::AzimuthalIntegration;
    config.nthreads = 3;

    JFJochProcess process(reader, dataset->experiment, dataset->pixel_mask, config);
    auto result = process.Run();

    CHECK_FALSE(result.cancelled);
    CHECK(result.images_processed == 6);
    CHECK_FALSE(result.written_master_path.has_value());

    reader.Close();
    remove("process_noout_in_master.h5");
    remove("process_noout_in_data_000001.h5");
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochProcess_Cancel", "[HDF5][Full]") {
    WriteTestDataset("process_cancel_in", 8);

    JFJochHDF5Reader reader;
    REQUIRE_NOTHROW(reader.ReadFile("process_cancel_in_master.h5"));
    auto dataset = reader.GetDataset();

    ProcessConfig config;
    config.mode = ProcessMode::AzimuthalIntegration;
    config.nthreads = 2;

    JFJochProcess process(reader, dataset->experiment, dataset->pixel_mask, config);
    process.Cancel();   // cancel before running: the worker loop stops immediately
    auto result = process.Run();

    CHECK(result.cancelled);
    CHECK(result.images_processed == 0);
    CHECK_FALSE(result.written_master_path.has_value());

    reader.Close();
    remove("process_cancel_in_master.h5");
    remove("process_cancel_in_data_000001.h5");
    REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}

TEST_CASE("JFJochProcessCommandLine_Full", "[process]") {
    DiffractionExperiment x(DetJF(1));
    IndexingSettings idx;
    idx.Algorithm(IndexingAlgorithmEnum::FFT);
    idx.GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum::BeamCenter);
    x.ImportIndexingSettings(idx);
    x.SpaceGroupNumber(96);

    ProcessConfig config;
    config.mode = ProcessMode::FullAnalysis;
    config.nthreads = 8;
    config.output_prefix = "run1";
    config.end_image = 500;
    config.rotation_indexing = true;
    config.two_pass_rotation = true;
    config.rotation_indexing_image_count = 30;
    config.spot_finding = DiffractionExperiment::DefaultDataProcessingSettings();

    const std::string cmd = JFJochProcessCommandLine(config, x, "/data/lyso_master.h5");
    CHECK(cmd.rfind("jfjoch_process", 0) == 0);
    CHECK(cmd.find("-N 8") != std::string::npos);
    CHECK(cmd.find("-e 500") != std::string::npos);
    CHECK(cmd.find("-o run1") != std::string::npos);
    CHECK(cmd.find("-X fft") != std::string::npos);
    CHECK(cmd.find("-S 96") != std::string::npos);
    CHECK(cmd.find("-R 30") != std::string::npos);
    CHECK(cmd.find("/data/lyso_master.h5") != std::string::npos);
}

TEST_CASE("JFJochProcessCommandLine_AzInt", "[process]") {
    DiffractionExperiment x(DetJF(1));
    AzimuthalIntegrationSettings a;
    a.AzimuthalBinCount(4);
    x.ImportAzimuthalIntegrationSettings(a);

    ProcessConfig config;
    config.mode = ProcessMode::AzimuthalIntegration;
    config.nthreads = 2;
    config.output_prefix = "az";

    const std::string cmd = JFJochProcessCommandLine(config, x, "in.h5");
    CHECK(cmd.rfind("jfjoch_azint", 0) == 0);
    CHECK(cmd.find("--azimuthal-bins 4") != std::string::npos);
    CHECK(cmd.find("--min-q") != std::string::npos);
    CHECK(cmd.find("in.h5") != std::string::npos);
}