// 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 <atomic>

#include "../common/ImageBuffer.h"
#include "../receiver/StillsGPUAnalysis.h"

namespace {
    DiffractionExperiment StillsExperiment() {
        DiffractionExperiment x(DetJF(1));
        x.ImagesPerTrigger(StillsGPUAnalysis::PREPASS_IMAGES + 200).NumTriggers(1).DetectorDistance_mm(100).BeamX_pxl(500).BeamY_pxl(250)
                .IncidentEnergy_keV(12.4);
        return x;
    }

    // A frame after the pre-pass's sample (numbers from PREPASS_IMAGES on are never in it), in a slot of `buffer`.
    std::shared_ptr<HeldImage> Frame(ImageBuffer &buffer, int64_t number) {
        auto held = std::make_shared<HeldImage>();
        held->slot = buffer.GetImageSlot();
        REQUIRE(held->slot != nullptr);
        held->message.number = StillsGPUAnalysis::PREPASS_IMAGES + number;
        held->message.image_collection_efficiency = 1.0f;
        return held;
    }

    StartMessage Start(const DiffractionExperiment &x) {
        StartMessage start;
        x.FillMessage(start);
        return start;
    }
}

TEST_CASE("StillsGPUAnalysis_StartFailureReleasesThreads", "[StillsGPUAnalysis]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());
    const DiffractionExperiment x = StillsExperiment();
    const PixelMask mask(x);
    ImageBuffer buffer(64 * 1024 * 1024);
    buffer.StartMeasurement(x);
    RugnuxLiveStatus status;

    std::atomic<int> failed = 0, sunk = 0;
    {
        StillsGPUAnalysis gpu(x, mask, Start(x), nullptr, 4, 100, logger,
                              [](const StillsGPUAnalysis::Setup &) {
                                  throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Timeout on pushing start message");
                              },
                              [&](DataMessage &, const AzimuthalIntegrationProfile &) { ++sunk; },
                              [] { return DiffractionExperiment::DefaultDataProcessingSettings(); },
                              [&](const JFJochException &) { ++failed; }, &status);
        for (int i = 0; i < 3; i++)
            REQUIRE(gpu.Hold(Frame(buffer, i)));
        REQUIRE_THROWS_AS(gpu.Finish(), JFJochException);
    }
    CHECK(failed == 1);
    CHECK(sunk == 0);
    CHECK(status.Get().state == RugnuxLiveState::Failed);
    // Every slot the frames were held in is back.
    CHECK(buffer.GetAvailSlots() == buffer.GetStatus().total_slots);
}

TEST_CASE("StillsGPUAnalysis_HoldIsBounded", "[StillsGPUAnalysis]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());
    const DiffractionExperiment x = StillsExperiment();
    const PixelMask mask(x);
    ImageBuffer buffer(64 * 1024 * 1024);
    buffer.StartMeasurement(x);
    REQUIRE(buffer.GetStatus().total_slots >= 4);

    StillsGPUAnalysis gpu(x, mask, Start(x), nullptr, 2, 3, logger,
                          [](const StillsGPUAnalysis::Setup &) {
                              throw JFJochException(JFJochExceptionCategory::ZeroMQ, "no writer");
                          },
                          [](DataMessage &, const AzimuthalIntegrationProfile &) {},
                          [] { return DiffractionExperiment::DefaultDataProcessingSettings(); });
    // The analysis waits for the pre-pass, which waits for its sample: nothing is taken off the queue.
    for (int i = 0; i < 3; i++)
        REQUIRE(gpu.Hold(Frame(buffer, i)));
    {
        const auto refused = Frame(buffer, 3);
        CHECK(!gpu.Hold(refused));
    }
    CHECK(buffer.GetStatus().preparation_slots == 3);
    REQUIRE_THROWS(gpu.Finish());
    CHECK(buffer.GetStatus().preparation_slots == 0);
}

TEST_CASE("StillsGPUAnalysis_FailingFrameDoesNotStopAnalysis", "[StillsGPUAnalysis]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());
    const DiffractionExperiment x = StillsExperiment();
    const PixelMask mask(x);
    ImageBuffer buffer(64 * 1024 * 1024);
    buffer.StartMeasurement(x);
    RugnuxLiveStatus status;

    std::atomic<int> sunk = 0;
    StillsGPUAnalysis gpu(x, mask, Start(x), nullptr, 1, 100, logger,
                          [](const StillsGPUAnalysis::Setup &) {},
                          [&](DataMessage &m, const AzimuthalIntegrationProfile &) {
                              ++sunk;
                              if (m.number == StillsGPUAnalysis::PREPASS_IMAGES)
                                  throw std::runtime_error("not a JFJochException");
                          },
                          [] { return DiffractionExperiment::DefaultDataProcessingSettings(); },
                          {}, &status);
    // The frames carry no image, so their analysis fails too - they are passed on all the same.
    for (int i = 0; i < 3; i++)
        REQUIRE(gpu.Hold(Frame(buffer, i)));
    REQUIRE_NOTHROW(gpu.Finish());
    CHECK(sunk == 3);
    // The outcome waits for everything to be written, and comes after the latency.
    CHECK(status.Get().state == RugnuxLiveState::Running);
    CHECK(!status.Get().last_image_to_written_s);
    gpu.Written();
    const auto s = status.Get();
    CHECK(s.state == RugnuxLiveState::Failed);
    CHECK(!s.message.empty());
    REQUIRE(s.last_image_to_written_s);
    REQUIRE(s.last_image_s);
    CHECK(s.elapsed_s >= s.last_image_s.value() + s.last_image_to_written_s.value() - 1e-6);
    REQUIRE(!s.steps.empty());
    CHECK(s.steps.back().title == "Writing");
    CHECK(s.steps.back().duration_s);
    CHECK(buffer.GetStatus().preparation_slots == 0);
}