// 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 "../broker/JFJochStateMachine.h"
#include "../acquisition_device/HLSSimulatedDevice.h"
#include "../receiver/JFJochReceiverService.h"

using namespace std::literals::chrono_literals;

TEST_CASE("JFJochStateMachine_States") {
    Logger logger("JFJochStateMachine_States");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF4M());

    DatasetSettings setup;

    REQUIRE(state_machine.GetStatus().state == JFJochState::Inactive);

    REQUIRE_THROWS(state_machine.Start(setup));
    REQUIRE_THROWS(state_machine.Pedestal());

    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);

    REQUIRE_NOTHROW(state_machine.Deactivate());
    REQUIRE(state_machine.GetStatus().state == JFJochState::Inactive);
}

TEST_CASE("JFJochStateMachine_State_Pedestal") {
    Logger logger("JFJochStateMachine_State_Pedestal");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF4M());

    DatasetSettings setup;

    state_machine.DebugOnly_SetState(JFJochState::Calibration);

    REQUIRE(state_machine.GetStatus().state == JFJochState::Calibration);

    REQUIRE_THROWS(state_machine.Start(setup));
    REQUIRE_THROWS(state_machine.Pedestal());
    REQUIRE_THROWS(state_machine.Initialize());
    REQUIRE(state_machine.WaitTillMeasurementDone(std::chrono::milliseconds(1)).state == JFJochState::Calibration);
}

TEST_CASE("JFJochStateMachine_State_Measure") {
    Logger logger("JFJochStateMachine_State_Measure");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF4M());

    DatasetSettings setup;

    state_machine.DebugOnly_SetState(JFJochState::Measuring);

    REQUIRE(state_machine.GetStatus().state == JFJochState::Measuring);

    REQUIRE_THROWS(state_machine.Start(setup));
    REQUIRE_THROWS(state_machine.Pedestal());
    REQUIRE_THROWS(state_machine.Initialize());
    REQUIRE(state_machine.WaitTillMeasurementDone(std::chrono::milliseconds(1)).state == JFJochState::Measuring);

    DetectorSettings settings{};
    REQUIRE_THROWS(state_machine.LoadDetectorSettings(settings));
}

TEST_CASE("JFJochStateMachine_State_Error") {
    Logger logger("JFJochStateMachine_State_Error");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF4M());

    DatasetSettings setup;

    state_machine.DebugOnly_SetState(JFJochState::Error,
        "msg1234",
        BrokerStatus::MessageSeverity::Error);

    REQUIRE(state_machine.GetStatus().state == JFJochState::Error);
    REQUIRE(state_machine.GetStatus().message.has_value());
    REQUIRE(state_machine.GetStatus().message == "msg1234");
    REQUIRE(state_machine.GetStatus().message_severity == BrokerStatus::MessageSeverity::Error);

    state_machine.DebugOnly_SetState(JFJochState::Inactive, "msg3456", BrokerStatus::MessageSeverity::Info);
    REQUIRE(state_machine.GetStatus().state == JFJochState::Inactive);
    REQUIRE(state_machine.GetStatus().message.has_value());
    REQUIRE(state_machine.GetStatus().message == "msg3456");
    REQUIRE(state_machine.GetStatus().message_severity == BrokerStatus::MessageSeverity::Info);

    state_machine.DebugOnly_SetState(JFJochState::Error);
    REQUIRE(state_machine.GetStatus().state == JFJochState::Error);
    REQUIRE(!state_machine.GetStatus().message.has_value());

    REQUIRE_THROWS(state_machine.Start(setup));
    REQUIRE_THROWS(state_machine.Pedestal());

    REQUIRE(state_machine.WaitTillMeasurementDone(std::chrono::milliseconds(1)).state == JFJochState::Error);

    DetectorSettings settings;
    REQUIRE_NOTHROW(state_machine.LoadDetectorSettings(settings));

    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);
}

TEST_CASE("JFJochStateMachine_NoDetectorSetup") {
    Logger logger("JFJochStateMachine_NoDetectorSetup");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    REQUIRE_THROWS(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
}

TEST_CASE("JFJochStateMachine_AddDetectorSetup") {
    Logger logger("JFJochStateMachine_AddDetectorSetup");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    DetectorSetup setup = DetJF4M();
    state_machine.AddDetectorSetup(setup);
    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
}

TEST_CASE("JFJochStateMachine_AddDetectorSetup_Gain") {
    Logger logger("JFJochStateMachine_AddDetectorSetup_Gain");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    DetectorSetup setup = DetJF(4);
    setup.LoadGain({"../../tests/test_data/gainMaps_M049.bin",
                    "../../tests/test_data/gainMaps_M049.bin",
                    "../../tests/test_data/gainMaps_M049.bin",
                    "../../tests/test_data/gainMaps_M049.bin"});
    state_machine.AddDetectorSetup(setup);
    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
}

TEST_CASE("JFJochStateMachine_AddDetectorSetup_Multiple") {
    Logger logger("JFJochStateMachine_AddDetectorSetup_Multiple");
    JFJochServices services(logger);
    DiffractionExperiment experiment;

    JFJochStateMachine state_machine(experiment, services, logger);
    REQUIRE_NOTHROW(state_machine.AddDetectorSetup(DetJF(4, 1, 0,0,false,"Det1", {"mx1", "mx2", "mx3", "mx4"})));
    REQUIRE_NOTHROW(state_machine.AddDetectorSetup(DetEIGER(2, 1, 0, 0, false, "Det2", {"mx1", "mx2", "mx3", "mx4"})));
    REQUIRE_NOTHROW(state_machine.AddDetectorSetup(DetDECTRIS(23,45, "Det3", {"mx5"})));


    auto dl = state_machine.GetDetectorsList();
    REQUIRE(dl.detector.size() == 3);
    REQUIRE(dl.detector[0].description == "Det1");
    REQUIRE(dl.detector[0].nmodules == 4);
    REQUIRE(dl.detector[0].detector_type == DetectorType::JUNGFRAU);
    REQUIRE(dl.detector[1].description == "Det2");
    REQUIRE(dl.detector[1].nmodules == 2);
    REQUIRE(dl.detector[1].detector_type == DetectorType::EIGER);
    REQUIRE(dl.detector[2].description == "Det3");
    REQUIRE(dl.detector[2].nmodules == 1);
    REQUIRE(dl.detector[2].detector_type == DetectorType::DECTRIS);

    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
    REQUIRE(state_machine.Experiment().GetModulesNum() == 4);
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);

    REQUIRE_THROWS(state_machine.SelectDetector(7));
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);

    REQUIRE_NOTHROW(state_machine.SelectDetector(2));
    REQUIRE(state_machine.Experiment().GetModulesNum() == 1);
    REQUIRE(state_machine.Experiment().GetDetectorDescription() == "Det3");
    REQUIRE(state_machine.GetStatus().state == JFJochState::Inactive);
}

TEST_CASE("JFJochStateMachine_LoadDetectorSettings_Error") {
    Logger logger("JFJochStateMachine_LoadDetectorSettings_Error");
    JFJochServices services(logger);
    DiffractionExperiment experiment;
    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF4M());

    DatasetSettings setup;

    state_machine.DebugOnly_SetState(JFJochState::Idle);

    DetectorSettings settings;
    settings.FrameTime(std::chrono::microseconds(1));

    REQUIRE_THROWS(state_machine.LoadDetectorSettings(settings));
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);
}

namespace {
    // Stands in for a writer that refuses to start the run - e.g. the output file exists and cannot be
    // overwritten. The refusal reaches the broker as an ordinary exception thrown while the start
    // message is being sent, which is what separates it from a critical detector fault.
    class RefusingImagePusher : public ImagePusher {
    public:
        bool refuse = true;
        void StartDataCollection(StartMessage &) override {
            if (refuse)
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "writer_refused_1234");
        }
        bool EndDataCollection(const EndMessage &) override { return true; }
        bool SendImage(const uint8_t *, size_t, int64_t) override { return true; }
        bool SendCalibration(const CompressedImage &) override { return true; }
        std::string PrintSetup() const override { return "RefusingImagePusher"; }
        ImagePusherType GetType() const override { return ImagePusherType::Test; }
    };
}

// An asynchronous start returns before the measurement thread runs, so a failure there has to be
// reported to whoever asks next. Without this the state is a plain Idle, indistinguishable from a
// timeout, and the writer's message is lost.
TEST_CASE("JFJochStateMachine_AsyncStartFailure") {
    Logger logger("JFJochStateMachine_AsyncStartFailure");

    DiffractionExperiment experiment(DetJF(2));
    experiment.Conversion().PedestalG0Frames(0).NumTriggers(1).UseInternalPacketGenerator(true)
            .ImagesPerTrigger(4).IncidentEnergy_keV(12.4);

    AcquisitionDeviceGroup aq_devices;
    for (int i = 0; i < experiment.GetDataStreamsNum(); i++)
        aq_devices.Add(std::make_unique<HLSSimulatedDevice>(i, 64));

    RefusingImagePusher pusher;
    JFJochReceiverService receiver_service(aq_devices, logger, pusher);

    JFJochServices services(logger);
    services.Receiver(&receiver_service);

    JFJochStateMachine state_machine(experiment, services, logger);
    state_machine.AddDetectorSetup(DetJF(2));
    state_machine.DebugOnly_SetState(JFJochState::Idle);

    DatasetSettings setup;
    setup.ImagesPerTrigger(4).NumTriggers(1);

    // The asynchronous start itself succeeds - it only launches the measurement thread.
    REQUIRE_NOTHROW(state_machine.Start(setup, true));

    // Both wait functions report the failure, and repeatedly: the exception_ptr is not consumed.
    REQUIRE_THROWS_WITH(state_machine.WaitTillNotBusy(std::chrono::seconds(10)),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    REQUIRE_THROWS_WITH(state_machine.WaitTillNotBusy(std::chrono::seconds(10)),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    REQUIRE_THROWS_WITH(state_machine.WaitTillMeasurementDone(std::chrono::seconds(10)),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));

    // An ordinary failure leaves the detector usable, so the run can be retried from Idle. The
    // status keeps the writer's message, and nothing is left of the run that never happened - a
    // progress figure in particular would show up in the frontend as a stalled acquisition.
    auto status = state_machine.GetStatus();
    REQUIRE(status.state == JFJochState::Idle);
    REQUIRE(status.message_severity == BrokerStatus::MessageSeverity::Error);
    REQUIRE_THAT(status.message.value_or(""), Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    REQUIRE_FALSE(status.progress.has_value());

    auto statistics = state_machine.GetMeasurementStatistics();
    REQUIRE(statistics.has_value());
    REQUIRE(statistics->images_collected == 0);
    REQUIRE(statistics->images_sent == 0);
    REQUIRE_FALSE(statistics->collection_efficiency.has_value());

    // A synchronous start reports the same failure directly, and leaves it visible to the wait
    // functions afterwards.
    REQUIRE_THROWS_WITH(state_machine.Start(setup),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    REQUIRE_THROWS_WITH(state_machine.WaitTillNotBusy(std::chrono::seconds(10)),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));

    // A start rejected before the measurement thread is launched must not leave the previous
    // failure behind for the next wait call to report.
    DatasetSettings rejected_setup;
    rejected_setup.ImagesPerTrigger(4).NumTriggers(1)
            .ImageTime(state_machine.Experiment().GetFrameTime() + std::chrono::nanoseconds(1));
    REQUIRE_THROWS(state_machine.Start(rejected_setup, true));
    REQUIRE_NOTHROW(state_machine.WaitTillNotBusy(std::chrono::seconds(10)));

    // Any other operation started from Idle supersedes the pending failure - it must not be
    // attributed to the pedestal, which is allowed from Idle and has nothing to do with it.
    REQUIRE_THROWS_WITH(state_machine.Start(setup),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    REQUIRE_NOTHROW(state_machine.Pedestal());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());

    // Nothing is left half-started: once the writer accepts the run, the next start goes through
    // and completes without any intervening re-initialisation.
    pusher.refuse = false;
    REQUIRE_NOTHROW(state_machine.Start(setup));
    REQUIRE(state_machine.GetStatus().state == JFJochState::Measuring);
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);
    REQUIRE(state_machine.GetMeasurementStatistics()->images_sent == 4);

    // Re-initialising discards a pending failure too, so a detector brought back up does not report
    // the failure of the run before it.
    pusher.refuse = true;
    REQUIRE_THROWS_WITH(state_machine.Start(setup),
                        Catch::Matchers::ContainsSubstring("writer_refused_1234"));
    services.Receiver(nullptr);
    REQUIRE_NOTHROW(state_machine.Initialize());
    REQUIRE_NOTHROW(state_machine.WaitTillMeasurementDone());
    REQUIRE(state_machine.GetStatus().state == JFJochState::Idle);
}