Jungfraujoch/receiver/JFJochReceiverService.h

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

#pragma once

#include "RugnuxLiveStatus.h"

#include <mutex>
#include <memory>

#include "JFJochReceiver.h"
#include "../acquisition_device/AcquisitionDeviceGroup.h"

#include "../preview/ZMQMetadataSocket.h"
#include "../preview/ZMQPreviewSocket.h"

struct RugnuxModel;

class JFJochReceiverService {
    std::unique_ptr<JFJochReceiver> receiver;
    AcquisitionDeviceGroup &aq_devices;
    Logger &logger;
    ImageBuffer image_buffer;
    ImagePusher &image_pusher;
    std::shared_ptr<ImagePuller> image_puller;
    int64_t nthreads = 8;

    // Rotation collections from a DECTRIS detector analysed by rugnux while they arrive
    // (JFJochReceiverRugnux), instead of image by image. Off unless asked for.
    bool rugnux_enable = false;
    std::shared_ptr<const RugnuxModel> rugnux_model;
    int rugnux_reserved_cpus = 4;   // left to the forward threads and the pusher
    double rugnux_deadline_s = 60;   // the analysis is cut this long after the last image (0 = never)
    RugnuxLiveStatus rugnux_status;   // of the last collection, kept after it ends

    enum class ReceiverState {Idle, Running};
    ReceiverState state = ReceiverState::Idle;
    mutable std::mutex state_mutex;
    std::condition_variable measurement_done;
    std::future<void> measurement;
    void FinalizeMeasurement();
    void FinalizeMeasurementChangeState();
    SpotFindingSettings spot_finding_settings;

    PreviewImage preview_image;
    std::unique_ptr<ZMQPreviewSocket> zmq_preview_socket;
    std::unique_ptr<ZMQMetadataSocket> zmq_metadata_socket;
    std::unique_ptr<IndexerThreadPool> indexer_thread_pool;

    JFJochReceiverCurrentStatus receiver_status;
    JFJochReceiverPlots plots;
public:
    JFJochReceiverService(AcquisitionDeviceGroup &aq_devices,
                          Logger &logger,
                          ImagePusher &pusher,
                          size_t send_buffer_size_MiB = 512);
    JFJochReceiverService& NumThreads(int64_t input);
    JFJochReceiverService& RugnuxLive(bool enable, std::shared_ptr<const RugnuxModel> model, int reserved_cpus,
                                      double deadline_s);
    RugnuxLiveSnapshot GetRugnuxStatus() const;
    JFJochReceiverService& PreviewSocket(const std::string &addr, const std::optional<int32_t> &watermark = {});
    JFJochReceiverService& MetadataSocket(const std::string &addr);

    JFJochReceiverService& PreviewSocketSettings(const ZMQPreviewSettings &input);
    JFJochReceiverService& MetadataSocketSettings(const ZMQMetadataSettings &input);
    JFJochReceiverService& Indexing(const IndexingSettings& input);

    ZMQPreviewSettings GetPreviewSocketSettings() const;
    ZMQMetadataSettings GetMetadataSocketSettings() const;

    std::string GetPreviewSocketAddress() const;
    std::string GetMetadataSocketAddress() const;

    void LoadInternalGeneratorImage(const DiffractionExperiment& experiment,
                                    const std::vector<uint16_t> &raw_expected_image,
                                    uint64_t image_number);
    // Ask the writer whether this run's output could be written, before anything is started.
    // Throws if it could not.
    void Preflight(const DiffractionExperiment &experiment);
    void Start(const DiffractionExperiment &experiment,
               const PixelMask &pixel_mask,
               const JFCalibration *calibration,
               std::shared_ptr<ImagePuller> puller = {});
    void Cancel(bool silent);
    JFJochReceiverOutput Stop();
    void SetSpotFindingSettings(const SpotFindingSettings& settings);
    MultiLinePlot GetDataProcessingPlot(const PlotRequest& request);
    void GetPlotRaw(std::vector<float> &v, PlotType type, const std::string &roi);

    std::optional<JFJochReceiverStatus> GetStatus();
    std::vector<AcquisitionDeviceNetConfig> GetNetworkConfig();
    std::vector<DeviceStatus> GetDeviceStatus() const;
    std::optional<float> GetProgress() const;

    void GetXFELEventCode(std::vector<uint64_t> &v) const;
    void GetXFELPulseID(std::vector<uint64_t> &v) const;

    void GetStartMessageFromBuffer(std::vector<uint8_t> &v);
    bool GetImageFromBuffer(std::vector<uint8_t> &v, int64_t image_number = ImageBuffer::MaxImage);
    std::string GetTIFFFromBuffer(int64_t image_number = ImageBuffer::MaxImage);
    std::string GetJPEGFromBuffer(const PreviewImageSettings &settings, int64_t image_number = ImageBuffer::MaxImage);
    ImageBufferStatus GetImageBufferStatus() const;
    void ClearImageBuffer();
    // Drop everything the previous run left readable - its status, plots, live rugnux status and
    // buffered images - so nothing of it is served under the next run's tokens. Idle only.
    void ClearResults();

    ImagePusherStatus GetImagePusherStatus() const;
    LiveCollection *GetLiveCollection() { return image_pusher.GetLiveCollection(); }
};