// 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 <fstream>
#include <future>
#include <thread>

#include "../image_pusher/HDF5FilePusher.h"
#include "../image_pusher/TCPStreamPusher.h"
#include "../image_pusher/ZMQStream2Pusher.h"
#include "../image_puller/TCPImagePuller.h"
#include "../writer/StreamWriter.h"
#include "../common/PixelMask.h"

namespace {
    const std::string preflight_dir = "preflight_test";

    void PlaceFile(const std::string &name) {
        std::filesystem::path path(name);
        if (path.has_parent_path())
            std::filesystem::create_directories(path.parent_path());
        std::ofstream f(name);
        f << "in the way";
    }

    // A run number that is neither zero nor the default: an ACK is only matched to its request if
    // both sides put this number on it, so anything that forgets it shows up as an ACK timeout.
    constexpr uint64_t preflight_run_number = 4711;

    StartMessage PreflightMessage(const std::string &prefix) {
        StartMessage msg{};
        msg.file_prefix = prefix;
        msg.run_number = preflight_run_number;
        msg.run_name = "preflight";
        msg.images_per_file = 2;
        msg.number_of_images = 5;
        msg.file_format = FileWriterFormat::NXmxVDS;
        msg.overwrite = false;
        return msg;
    }

    void WaitForWriters(const TCPStreamPusher &pusher, size_t expected) {
        for (int i = 0; i < 200 && pusher.GetConnectedWriters() < expected; i++)
            std::this_thread::sleep_for(std::chrono::milliseconds(25));
    }
}

// The pre-flight is a dry run of the writer's own start-time checks, sent before anything is
// started. What it has to get right is that the answer comes back at all (an ACK carries the run
// number of the run being checked, and either side forgetting it looks like a timeout), that a
// refusal is reported with the writer's own message, and that a refusal costs nothing: the writer
// stays idle and the connection stays usable for the next attempt.
TEST_CASE("Preflight_TCP", "[Preflight][TCP]") {
    Logger logger("Preflight_TCP");

    std::filesystem::remove_all(preflight_dir);

    TCPStreamPusher pusher("tcp://127.0.0.1:*", 1);
    TCPImagePuller puller(pusher.GetAddress()[0], 8 * 1024 * 1024);
    StreamWriter writer(logger, puller);

    auto writer_future = std::async(std::launch::async, [&] { return writer.Run(); });

    WaitForWriters(pusher, 1);
    REQUIRE(pusher.GetConnectedWriters() == 1);

    // Nothing in the way: the check passes, and it creates the output directory on the way - the
    // prefix names a subdirectory that does not exist yet.
    auto clean = PreflightMessage(preflight_dir + "/sub/clean");
    REQUIRE_NOTHROW(pusher.Preflight(clean));
    CHECK(std::filesystem::is_directory(preflight_dir + "/sub"));

    // A dry run starts nothing: the writer is still waiting for a START.
    CHECK(writer.GetStatistics().state == StreamWriterState::Idle);

    // A master file already there is refused, and the writer's own message comes back with it.
    auto collision = PreflightMessage(preflight_dir + "/collision");
    PlaceFile(preflight_dir + "/collision_master.h5");
    REQUIRE_THROWS_WITH(pusher.Preflight(collision),
                        Catch::Matchers::ContainsSubstring("collision_master.h5"));

    // Refusing a run is not a transport failure - the connection carries the next attempt.
    CHECK(pusher.GetConnectedWriters() == 1);
    CHECK(writer.GetStatistics().state == StreamWriterState::Idle);
    auto second = PreflightMessage(preflight_dir + "/second");
    REQUIRE_NOTHROW(pusher.Preflight(second));

    // The same collision passes once overwriting is allowed.
    auto overwriting = PreflightMessage(preflight_dir + "/collision");
    overwriting.overwrite = true;
    REQUIRE_NOTHROW(pusher.Preflight(overwriting));

    // A data file in the way is refused too, with no master file present. This is the case the
    // start-time check used to let through: the run was written in full and only failed when the
    // first data file was renamed into place at the end of it.
    auto data_collision = PreflightMessage(preflight_dir + "/data_collision");
    PlaceFile(HDF5Metadata::DataFileName(data_collision, 1));
    REQUIRE_FALSE(std::filesystem::exists(preflight_dir + "/data_collision_master.h5"));
    REQUIRE_THROWS_WITH(pusher.Preflight(data_collision),
                        Catch::Matchers::ContainsSubstring("Output file already exists"));

    // A prefix that escapes the output directory is refused by the same path guard the writer
    // applies at start time.
    auto escaping = PreflightMessage("../escape");
    REQUIRE_THROWS(pusher.Preflight(escaping));

    // A run the pre-flight passed still starts: the run number the pre-flight left behind must not
    // stop the START ACK from being matched.
    auto started = PreflightMessage(preflight_dir + "/started");
    REQUIRE_NOTHROW(pusher.Preflight(started));
    REQUIRE_NOTHROW(pusher.StartDataCollection(started));
    CHECK(writer.GetStatistics().run_number == preflight_run_number);
    EndMessage end{};
    end.run_number = preflight_run_number;
    CHECK(pusher.EndDataCollection(end));

    writer.Cancel();
    REQUIRE_NOTHROW(writer_future.get());

    std::filesystem::remove_all(preflight_dir);
}

// With more than one writer every connection is asked, but only the one holding the master file
// checks the output - the same assignment the start makes, so a pre-flight cannot pass a run the
// start would then refuse.
TEST_CASE("Preflight_TCP_TwoWriters", "[Preflight][TCP]") {
    Logger logger("Preflight_TCP_TwoWriters");

    std::filesystem::remove_all(preflight_dir);

    TCPStreamPusher pusher("tcp://127.0.0.1:*", 2);

    std::vector<std::unique_ptr<TCPImagePuller>> pullers;
    std::vector<std::unique_ptr<StreamWriter>> writers;
    std::vector<std::future<StreamWriterOutput>> futures;

    for (int i = 0; i < 2; i++) {
        pullers.push_back(std::make_unique<TCPImagePuller>(pusher.GetAddress()[0], 8 * 1024 * 1024));
        writers.push_back(std::make_unique<StreamWriter>(logger, *pullers.back()));
        futures.push_back(std::async(std::launch::async, [w = writers.back().get()] { return w->Run(); }));
    }

    WaitForWriters(pusher, 2);
    REQUIRE(pusher.GetConnectedWriters() == 2);

    auto clean = PreflightMessage(preflight_dir + "/two_clean");
    REQUIRE_NOTHROW(pusher.Preflight(clean));

    // Both connections were asked, and each was told which socket it is.
    CHECK(writers[0]->GetStatistics().socket_number != writers[1]->GetStatistics().socket_number);

    auto collision = PreflightMessage(preflight_dir + "/two_collision");
    PlaceFile(preflight_dir + "/two_collision_master.h5");
    REQUIRE_THROWS_WITH(pusher.Preflight(collision),
                        Catch::Matchers::ContainsSubstring("two_collision_master.h5"));

    for (auto &w : writers)
        w->Cancel();
    for (auto &f : futures)
        REQUIRE_NOTHROW(f.get());

    std::filesystem::remove_all(preflight_dir);
}

TEST_CASE("Preflight_TCP_NoWriters", "[Preflight][TCP]") {
    TCPStreamPusher pusher("tcp://127.0.0.1:*", 1);

    auto msg = PreflightMessage(preflight_dir + "/no_writers");
    REQUIRE_THROWS_WITH(pusher.Preflight(msg),
                        Catch::Matchers::ContainsSubstring("No writers connected"));
}

// The in-process writer has no transport to ask, so its pre-flight is the check itself.
TEST_CASE("Preflight_HDF5FilePusher", "[Preflight]") {
    std::filesystem::remove_all(preflight_dir);

    HDF5FilePusher pusher;

    auto clean = PreflightMessage(preflight_dir + "/hdf5_clean");
    clean.write_master_file = true;
    REQUIRE_NOTHROW(pusher.Preflight(clean));
    CHECK(std::filesystem::is_directory(preflight_dir));

    auto collision = PreflightMessage(preflight_dir + "/hdf5_collision");
    collision.write_master_file = true;
    PlaceFile(preflight_dir + "/hdf5_collision_master.h5");
    REQUIRE_THROWS_WITH(pusher.Preflight(collision),
                        Catch::Matchers::ContainsSubstring("hdf5_collision_master.h5"));

    // A pusher that is only asked to write data files does not speak for the master file.
    auto not_master = PreflightMessage(preflight_dir + "/hdf5_collision");
    not_master.write_master_file = false;
    REQUIRE_NOTHROW(pusher.Preflight(not_master));

    std::filesystem::remove_all(preflight_dir);
}

// The pre-flight frame goes out once per run per writer, and it describes the run without carrying
// any of the per-pixel arrays a start message does. Keeping it that way is the point of building it
// from FillMessage alone: measured on a JUNGFRAU 9M it is 1.4 kB, against 359 kB once the pixel mask
// is added and 543 kB with the azimuthal and ROI maps on top. The bound below is loose - it is here
// to catch a per-pixel array finding its way into FillMessage, not to pin the exact size.
TEST_CASE("Preflight_MessageStaysSmall", "[Preflight]") {
    DiffractionExperiment x(DetJF(18)); // JUNGFRAU 9M - the largest detector in the standard set
    x.FilePrefix("size_probe").NumTriggers(1).ImagesPerTrigger(3600).ImagesPerFile(1000)
     .UseInternalPacketGenerator(true).IncidentEnergy_keV(12.4);

    // Built exactly as JFJochReceiverService::Preflight builds it.
    StartMessage message{};
    x.FillMessage(message);

    std::vector<uint8_t> buffer(MESSAGE_SIZE_FOR_START_END);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());
    serializer.SerializeSequenceStart(message);
    CHECK(serializer.GetBufferSize() < 64 * 1024);

    // For contrast: one per-pixel array is two orders of magnitude larger than the whole message.
    PixelMask mask(x);
    message.pixel_mask["default"] = mask.GetMask(x);
    CBORStream2Serializer with_mask(buffer.data(), buffer.size());
    with_mask.SerializeSequenceStart(message);
    CHECK(with_mask.GetBufferSize() > 100 * 1024);
}

// The transport with no back-channel cannot ask anything, so its pre-flight has to be silent
// rather than guess - the ZeroMQ writer still fails at the final rename, as it always has.
TEST_CASE("Preflight_NoBackChannelIsSilent", "[Preflight]") {
    std::filesystem::remove_all(preflight_dir);

    ZMQStream2Pusher pusher({"ipc://*"});

    auto collision = PreflightMessage(preflight_dir + "/zmq_collision");
    collision.write_master_file = true;
    PlaceFile(preflight_dir + "/zmq_collision_master.h5");
    REQUIRE_NOTHROW(pusher.Preflight(collision));

    std::filesystem::remove_all(preflight_dir);
}