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

#include <catch2/catch_all.hpp>
#include <filesystem>
#include <future>
#include <thread>
#include <sys/stat.h>

#include "../writer/StreamWriter.h"
#include "../image_pusher/ZMQStream2Pusher.h"
#include "../receiver/JFJochReceiverService.h"
#include "../image_pusher/ZMQWriterNotificationPuller.h"
#include "../image_puller/TestImagePuller.h"

TEST_CASE("StreamWriterTest_ZMQ", "[StreamWriter]") {
    RegisterHDF5Filter();

    Logger logger("StreamWriterTest_ZMQ");

    DiffractionExperiment x(DetJF(2));
    x.FilePrefix("subdir/StreamWriterTest").NumTriggers(1).ImagesPerTrigger(5)
            .UseInternalPacketGenerator(true).Raw().PedestalG0Frames(0).OverwriteExistingFiles(true);

    PixelMask pixel_mask(x);

    JFModuleGainCalibration gain;
    AcquisitionDeviceGroup aq_devices;
    for (int i = 0; i < x.GetDataStreamsNum(); i++)
        aq_devices.AddHLSDevice(64);

    ZMQStream2Pusher pusher({"ipc://*"});
    JFJochReceiverService fpga_receiver_service(aq_devices, logger, pusher);

    std::unique_ptr<StreamWriter> writer;
    REQUIRE(x.GetImageNum() == 5);
    auto pusher_addr = pusher.GetAddress();
    REQUIRE(pusher_addr.size() == 1);
    ZMQImagePuller puller(pusher_addr[0]);
    REQUIRE_NOTHROW(writer = std::make_unique<StreamWriter>(logger, puller));
    CHECK(writer->GetStatistics().state == StreamWriterState::Idle);
    REQUIRE_NOTHROW(fpga_receiver_service.Start(x, pixel_mask, nullptr));

    REQUIRE_NOTHROW(writer->Run());

    REQUIRE_NOTHROW(fpga_receiver_service.Stop());
    REQUIRE(fpga_receiver_service.GetStatus()->images_collected == 5);
    REQUIRE(fpga_receiver_service.GetStatus()->images_sent == 5);
    CHECK(writer->GetStatistics().state == StreamWriterState::Finalized);
    CHECK(writer->GetStatistics().processed_images == 5);
    CHECK(writer->GetStatistics().file_prefix == x.GetFilePrefix());

    // HDF5 file can be opened
    std::unique_ptr<HDF5ReadOnlyFile> file;
    REQUIRE_NOTHROW(file = std::make_unique<HDF5ReadOnlyFile>("subdir/StreamWriterTest_data_000001.h5"));
    std::unique_ptr<HDF5DataSet> dataset;
    REQUIRE_NOTHROW(dataset = std::make_unique<HDF5DataSet>(*file, "/entry/data/data"));
    std::unique_ptr<HDF5DataSpace> dataspace;
    REQUIRE_NOTHROW(dataspace = std::make_unique<HDF5DataSpace>(*dataset));

    REQUIRE(dataspace->GetNumOfDimensions() == 3);
    REQUIRE(dataspace->GetDimensions()[0] == 5);
    REQUIRE(dataspace->GetDimensions()[1] == RAW_MODULE_COLS);
    REQUIRE(dataspace->GetDimensions()[2] == 2*RAW_MODULE_LINES);

    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest_master.h5"));
    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest_data_000001.h5"));
    REQUIRE(std::filesystem::remove("subdir"));
}

TEST_CASE("StreamWriterTest_ZMQ_Update", "[StreamWriter]") {
    RegisterHDF5Filter();

    Logger logger("StreamWriterTest_ZMQ_Update");

    DatasetSettings d;
    d.FilePrefix("subdir/StreamWriterTest2").NumTriggers(1).ImagesPerTrigger(5).RunName("run1").RunNumber(256);
    DiffractionExperiment x(DetJF(2));
    x.UseInternalPacketGenerator(true).Raw().PedestalG0Frames(0)
            .ImportDatasetSettings(d).OverwriteExistingFiles(true);
    PixelMask pixel_mask(x);

    JFModuleGainCalibration gain;
    AcquisitionDeviceGroup aq_devices;
    for (int i = 0; i < x.GetDataStreamsNum(); i++)
        aq_devices.AddHLSDevice(64);

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");
    JFJochReceiverService fpga_receiver_service(aq_devices, logger, pusher);

    std::unique_ptr<StreamWriter> writer;
    REQUIRE(x.GetImageNum() == 5);
    auto pusher_addr = pusher.GetAddress();
    REQUIRE(pusher_addr.size() == 1);

    ZMQImagePuller puller(pusher_addr[0]);
    REQUIRE_NOTHROW(writer = std::make_unique<StreamWriter>(logger, puller));
    CHECK(writer->GetStatistics().state == StreamWriterState::Idle);
    REQUIRE_NOTHROW(fpga_receiver_service.Start(x, pixel_mask, nullptr));

    REQUIRE_NOTHROW(writer->Run());

    REQUIRE_NOTHROW(fpga_receiver_service.Stop());
    REQUIRE(fpga_receiver_service.GetStatus()->images_collected == 5);
    REQUIRE(fpga_receiver_service.GetStatus()->images_sent == 5);
    CHECK(writer->GetStatistics().state == StreamWriterState::Finalized);
    CHECK(writer->GetStatistics().processed_images == 5);
    CHECK(writer->GetStatistics().file_prefix == x.GetFilePrefix());

    // HDF5 file can be opened
    std::unique_ptr<HDF5ReadOnlyFile> file;
    REQUIRE_NOTHROW(file = std::make_unique<HDF5ReadOnlyFile>("subdir/StreamWriterTest2_data_000001.h5"));
    std::unique_ptr<HDF5DataSet> dataset;
    REQUIRE_NOTHROW(dataset = std::make_unique<HDF5DataSet>(*file, "/entry/data/data"));
    std::unique_ptr<HDF5DataSpace> dataspace;
    REQUIRE_NOTHROW(dataspace = std::make_unique<HDF5DataSpace>(*dataset));

    REQUIRE(dataspace->GetNumOfDimensions() == 3);
    REQUIRE(dataspace->GetDimensions()[0] == 5);
    REQUIRE(dataspace->GetDimensions()[1] == RAW_MODULE_COLS);
    REQUIRE(dataspace->GetDimensions()[2] == 2*RAW_MODULE_LINES);

    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest2_master.h5"));
    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest2_data_000001.h5"));
    REQUIRE(std::filesystem::remove("subdir"));
}

TEST_CASE("StreamWriterTest_ZMQ_Update_NoNotification", "[StreamWriter]") {
    // This tests simulates what happens if writer notification about writing end is missing
    // Expected end result: receiver ends with an exception

    RegisterHDF5Filter();

    Logger logger("StreamWriterTest_ZMQ_Update_NoNotification");

    DatasetSettings d;
    d.FilePrefix("subdir/StreamWriterTest3").NumTriggers(1).ImagesPerTrigger(5).RunName("run1").RunNumber(256);
    DiffractionExperiment x(DetJF(2));
    x.UseInternalPacketGenerator(true).Raw().PedestalG0Frames(0)
            .ImportDatasetSettings(d).OverwriteExistingFiles(true);
    PixelMask pixel_mask(x);

    JFModuleGainCalibration gain;
    AcquisitionDeviceGroup aq_devices;
    for (int i = 0; i < x.GetDataStreamsNum(); i++)
        aq_devices.AddHLSDevice(64);

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");
    JFJochReceiverService fpga_receiver_service(aq_devices, logger, pusher);

    std::unique_ptr<StreamWriter> writer;
    REQUIRE(x.GetImageNum() == 5);
    auto pusher_addr = pusher.GetAddress();
    REQUIRE(pusher_addr.size() == 1);

    ZMQImagePuller puller(pusher_addr[0]);
    REQUIRE_NOTHROW(writer = std::make_unique<StreamWriter>(logger, puller));
    writer->DebugSkipWriteNotification(true);

    CHECK(writer->GetStatistics().state == StreamWriterState::Idle);
    REQUIRE_NOTHROW(fpga_receiver_service.Start(x, pixel_mask, nullptr));

    REQUIRE_NOTHROW(writer->Run());

    JFJochReceiverOutput r;
    REQUIRE_NOTHROW(r = fpga_receiver_service.Stop());
    REQUIRE(!r.writer_err.empty());
    REQUIRE(fpga_receiver_service.GetStatus()->images_collected == 5);
    REQUIRE(fpga_receiver_service.GetStatus()->images_sent == 5);
    CHECK(writer->GetStatistics().state == StreamWriterState::Finalized);
    CHECK(writer->GetStatistics().processed_images == 5);
    CHECK(writer->GetStatistics().file_prefix == x.GetFilePrefix());

    // HDF5 file can be opened
    std::unique_ptr<HDF5ReadOnlyFile> file;
    REQUIRE_NOTHROW(file = std::make_unique<HDF5ReadOnlyFile>("subdir/StreamWriterTest3_data_000001.h5"));
    std::unique_ptr<HDF5DataSet> dataset;
    REQUIRE_NOTHROW(dataset = std::make_unique<HDF5DataSet>(*file, "/entry/data/data"));
    std::unique_ptr<HDF5DataSpace> dataspace;
    REQUIRE_NOTHROW(dataspace = std::make_unique<HDF5DataSpace>(*dataset));

    REQUIRE(dataspace->GetNumOfDimensions() == 3);
    REQUIRE(dataspace->GetDimensions()[0] == 5);
    REQUIRE(dataspace->GetDimensions()[1] == RAW_MODULE_COLS);
    REQUIRE(dataspace->GetDimensions()[2] == 2*RAW_MODULE_LINES);

    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest3_master.h5"));
    REQUIRE(std::filesystem::remove("subdir/StreamWriterTest3_data_000001.h5"));
    REQUIRE(std::filesystem::remove("subdir"));
}

// Two writers, the second one writing every other data file: the master, written by the first, must
// appear only after the second has closed its data files - here the second one starts only after the
// whole collection has been sent, so it is still writing when END goes out.
TEST_CASE("StreamWriterTest_ZMQ_MasterAfterAllDataFiles", "[StreamWriter]") {
    RegisterHDF5Filter();
    Logger logger("StreamWriterTest_ZMQ_MasterAfterAllDataFiles");

    DatasetSettings d;
    d.FilePrefix("subdir4/StreamWriterTest4").NumTriggers(1).ImagesPerTrigger(4).ImagesPerFile(2)
            .RunName("run1").RunNumber(257);
    DiffractionExperiment x(DetJF(2));
    x.UseInternalPacketGenerator(true).Raw().PedestalG0Frames(0)
            .ImportDatasetSettings(d).OverwriteExistingFiles(true);
    PixelMask pixel_mask(x);

    AcquisitionDeviceGroup aq_devices;
    for (int i = 0; i < x.GetDataStreamsNum(); i++)
        aq_devices.AddHLSDevice(64);

    ZMQStream2Pusher pusher({"ipc://*", "ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");
    JFJochReceiverService fpga_receiver_service(aq_devices, logger, pusher);

    auto pusher_addr = pusher.GetAddress();
    REQUIRE(pusher_addr.size() == 2);
    ZMQImagePuller puller0(pusher_addr[0]), puller1(pusher_addr[1]);
    StreamWriter writer0(logger, puller0), writer1(logger, puller1);

    REQUIRE_NOTHROW(fpga_receiver_service.Start(x, pixel_mask, nullptr));
    auto w0 = std::async(std::launch::async, [&] { return writer0.Run(); });
    for (int i = 0; i < 1200; i++) {
        const auto status = fpga_receiver_service.GetStatus();
        if (status && status->images_sent == 4)
            break;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    auto stop = std::async(std::launch::async, [&] { return fpga_receiver_service.Stop(); });
    std::this_thread::sleep_for(std::chrono::seconds(2));
    auto w1 = std::async(std::launch::async, [&] { return writer1.Run(); });
    JFJochReceiverOutput r;
    REQUIRE_NOTHROW(r = stop.get());
    w0.get();
    w1.get();
    CHECK(r.writer_err.empty());

    auto changed = [](const std::string &path) {
        struct stat st{};
        REQUIRE(stat(path.c_str(), &st) == 0);
        return std::make_pair(st.st_ctim.tv_sec, st.st_ctim.tv_nsec);
    };
    const auto master = changed("subdir4/StreamWriterTest4_master.h5");
    for (int f = 1; f <= 2; f++)
        CHECK(changed("subdir4/StreamWriterTest4_data_00000" + std::to_string(f) + ".h5") <= master);

    std::filesystem::remove_all("subdir4");
}

// A failure that is not a JFJochException must end the run the same way as one that is: error state,
// END processed, state Finalized. Before this test the writer caught only its own exception type at the
// per-image and END sites, and anything else - here the JSON library refusing a sample name that is not
// UTF-8 while publishing a file-done notification - left Run() with the state stuck in Receiving, no END
// acknowledgement and no notification, and on TCP a broker waiting for that acknowledgement for ever.
TEST_CASE("StreamWriterTest_NonJFJochException_EndsRunInErrorState", "[StreamWriter]") {
    RegisterHDF5Filter();
    Logger logger("StreamWriterTest_NonJFJochException");

    DiffractionExperiment x(DetJF(2));
    x.FilePrefix("subdir5/StreamWriterTest5").NumTriggers(1).ImagesPerTrigger(2).ImagesPerFile(1)
            .Compression(CompressionAlgorithm::NO_COMPRESSION).OverwriteExistingFiles(true)
            .SampleName("\xff\xfe not utf-8");

    StartMessage start_message;
    x.FillMessage(start_message);
    start_message.write_master_file = true;

    TestImagePuller puller;
    // The file-done PUB socket is what makes the writer serialise the sample name to JSON.
    StreamWriter writer(logger, puller, "ipc://*");

    puller.Put({.cbor = std::make_shared<CBORStream2DeserializerOutput>(start_message)});
    std::vector<int16_t> image(x.GetPixelsNum());
    for (int i = 0; i < 2; i++) {
        DataMessage message{};
        message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
        message.number = i;
        puller.Put({.cbor = std::make_shared<CBORStream2DeserializerOutput>(message)});
    }
    EndMessage end_message{};
    end_message.max_image_number = 2;
    puller.Put({.cbor = std::make_shared<CBORStream2DeserializerOutput>(end_message)});

    StreamWriterOutput out;
    REQUIRE_NOTHROW(out = writer.Run());

    // Image 0 filled its one-image data file; closing it published the notification, which threw.
    CHECK(out.stats.state == StreamWriterState::Finalized);
    CHECK(out.stats.processed_images == 0);
    CHECK(out.data_file_stats.empty());
    // The master is never written for a run that ended in error; the data file was renamed into
    // place before the notification was attempted.
    CHECK(!std::filesystem::exists("subdir5/StreamWriterTest5_master.h5"));
    CHECK(std::filesystem::exists("subdir5/StreamWriterTest5_data_000001.h5"));

    std::filesystem::remove_all("subdir5");
}