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

#ifndef JUNGFRAUJOCH_JFJOCHSTATEMACHINE_H
#define JUNGFRAUJOCH_JFJOCHSTATEMACHINE_H

#include <string>
#include <mutex>
#include <future>
#include <optional>

#include "../common/DiffractionExperiment.h"
#include "../jungfrau/JFCalibration.h"
#include "../common/Logger.h"

#include "JFJochServices.h"
#include "../common/ROIMap.h"

enum class JFJochState {Inactive, Idle, Measuring, Error, Busy, Pedestal};

struct BrokerStatus {
    JFJochState state = JFJochState::Inactive;
    std::optional<float> progress;
    std::optional<std::string> message;
    enum class MessageSeverity {Error, Info, Warning, Success} message_severity = MessageSeverity::Error;
    int64_t gpu_count;
    bool ml_resolution_estimation;
};

struct DetectorListElement {
    std::string description;
    std::string serial_number;
    std::string base_ipv4_addr;
    int64_t udp_interface_count;
    int64_t nmodules;
    int64_t width;
    int64_t height;
    std::chrono::microseconds readout_time;
    std::chrono::microseconds min_frame_time;
    std::chrono::microseconds min_count_time;
    DetectorType detector_type;
    float pixel_size_mm;
};

struct DetectorList {
    std::vector<DetectorListElement> detector;
    int64_t current_id;
};

struct MeasurementStatistics {
    std::string file_prefix;
    std::string experiment_group;
    int64_t run_number;

    int64_t images_expected;
    int64_t images_collected;
    int64_t images_sent;
    int64_t images_skipped;
    int64_t max_image_number_sent;
    std::optional<float> collection_efficiency;
    std::optional<float> compression_ratio;

    bool cancelled;
    std::optional<int64_t> max_receive_delay;

    std::optional<float> indexing_rate;

    int64_t detector_width;
    int64_t detector_height;
    int64_t detector_pixel_depth;

    std::optional<float> bkg_estimate;
    std::optional<std::pair<float, float>> beam_center_drift_pxl;

    std::string unit_cell;

    std::optional<float> error_pixels;
    std::optional<float> saturated_pixels;
    std::optional<float> roi_beam_npixel;
    std::optional<float> roi_beam_sum;
};

class JFJochStateMachine {
    Logger &logger;
    JFJochServices &services;

    std::future<void> measurement;

    // assuming immutable during normal operation
    std::vector<DetectorSetup> detector_setup;
    std::vector<JFModuleGainCalibration> gain_calibration;

    mutable std::mutex experiment_detector_settings_mutex;
    mutable std::mutex experiment_azimuthal_integration_settings_mutex;
    mutable std::mutex experiment_instrument_metadata_mutex;
    mutable std::mutex experiment_image_format_settings_mutex;
    mutable std::mutex experiment_file_writer_settings_mutex;
    mutable std::mutex experiment_indexing_settings_mutex;
    DiffractionExperiment experiment;

    // mutex m is protecting:
    mutable std::mutex m;
    std::condition_variable c;
    volatile JFJochState state = JFJochState::Inactive; // state should not be set directly, but through SetState function
    volatile bool cancel_sequence = false;
    std::unique_ptr<JFCalibration> calibration;
    PixelMask pixel_mask;
    int64_t current_detector_setup; // Lock only on change
    std::optional<ScanResult> scan_result;

    mutable std::mutex calibration_statistics_mutex;
    std::vector<JFCalibrationModuleStatistics> calibration_statistics;

    mutable std::mutex data_processing_settings_mutex;
    SpotFindingSettings data_processing_settings;

    mutable std::mutex pixel_mask_statistics_mutex;
    PixelMaskStatistics pixel_mask_statistics;

    mutable std::mutex broker_status_mutex;
    BrokerStatus broker_status;

    mutable std::mutex roi_mutex;
    ROIDefinition roi;

    bool indexing_possible;
    bool resolution_estimate_possible;

    const int32_t gpu_count;

    void UpdatePixelMaskStatistics(const PixelMaskStatistics &input);

    // Private functions assume that lock m is acquired
    void SetState(JFJochState curr_state,
        const std::optional<std::string> &message = {},
        BrokerStatus::MessageSeverity message_severity = BrokerStatus::MessageSeverity::Info);
    void MeasurementThread();
    void PedestalThread(std::unique_lock<std::mutex> ul);
    void InitializeThread(std::unique_lock<std::mutex> ul);
    bool ImportPedestalG1G2(const JFJochReceiverOutput &receiver_output, size_t gain_level, size_t storage_cell = 0);
    bool ImportPedestalG0(const JFJochReceiverOutput &receiver_output);
    bool IsRunning() const; // Is state Busy/Pedestal/Measure
    void ResetError() noexcept;
    void TakePedestalInternalAll(std::unique_lock<std::mutex> &ul);
    void TakePedestalInternalG0(std::unique_lock<std::mutex> &ul);
    void TakePedestalInternalG1(std::unique_lock<std::mutex> &ul, int32_t storage_cell = 0);
    void TakePedestalInternalG2(std::unique_lock<std::mutex> &ul, int32_t storage_cell = 0);
    bool ImportDetectorSettings(const DetectorSettings& input);
    void UpdateROIDefinition();
public:
    JFJochStateMachine(const DiffractionExperiment& experiment,
                       JFJochServices &in_services,
                       Logger &logger);
    ~JFJochStateMachine();

    void Initialize();
    void Pedestal();
    void Deactivate();
    void Start(const DatasetSettings& settings);
    BrokerStatus WaitTillMeasurementDone();
    BrokerStatus WaitTillMeasurementDone(std::chrono::milliseconds timeout);
    void Trigger();

    void Cancel();

    void SetCalibrationStatistics(const std::vector<JFCalibrationModuleStatistics> &input);

    DetectorSettings GetDetectorSettings() const;
    void LoadDetectorSettings(const DetectorSettings& settings);

    InstrumentMetadata GetInstrumentMetadata() const;
    void LoadInstrumentMetadata(const InstrumentMetadata& settings);

    FileWriterSettings GetFileWriterSettings() const;
    void LoadFileWriterSettings(const FileWriterSettings& settings);

    ImageFormatSettings GetImageFormatSettings() const;
    void LoadImageFormatSettings(const ImageFormatSettings& settings);
    void RawImageFormatSettings();
    void ConvImageFormatSettings();

    // return by value to ensure thread safety
    std::optional<MeasurementStatistics> GetMeasurementStatistics() const;
    std::vector<JFCalibrationModuleStatistics> GetCalibrationStatistics() const;

    BrokerStatus GetStatus() const;
    MultiLinePlot GetPlots(const PlotRequest &request) const;

    void SetSpotFindingSettings(const SpotFindingSettings& settings);
    SpotFindingSettings GetSpotFindingSettings() const;

    DetectorList GetDetectorsList() const;
    void SelectDetector(int64_t id);
    std::optional<DetectorStatus> GetDetectorStatus() const;

    void SetRadialIntegrationSettings(const AzimuthalIntegrationSettings& settings);
    AzimuthalIntegrationSettings GetRadialIntegrationSettings() const;

    std::string GetPreviewJPEG(const PreviewImageSettings& settings, int64_t image_number) const;
    std::string GetPreviewTIFF(int64_t image_number) const;
    std::string GetPedestalTIFF(size_t gain_level, size_t sc) const;

    void LoadInternalGeneratorImage(const void *data, size_t size, uint64_t image_number);
    void LoadInternalGeneratorImageTIFF(const std::string &s, uint64_t image_number);

    // Not thread safe - only for configuration in serial context
    DiffractionExperiment Experiment();

    // Function for debug only - UNSAFE for real operation
    void DebugOnly_SetState(JFJochState state,
        const std::optional<std::string> &message = {},
        BrokerStatus::MessageSeverity message_severity = BrokerStatus::MessageSeverity::Info);

    void SetROIDefinition(const ROIDefinition& input);
    ROIDefinition GetROIDefintion() const;

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

    std::string GetFullPixelMaskTIFF() const;
    std::string GetUserPixelMaskTIFF() const;
    std::vector<uint32_t> GetFullPixelMask() const;
    std::vector<uint32_t> GetUserPixelMask() const;

    void SetUserPixelMask(const std::vector<uint32_t> &v);

    std::vector<DeviceStatus> GetDeviceStatus() const;

    void SetPreviewSocketSettings(const ZMQPreviewSettings &input);
    ZMQPreviewSettings GetPreviewSocketSettings();

    void SetMetadataSocketSettings(const ZMQMetadataSettings &input);
    ZMQMetadataSettings GetMetadataSocketSettings();

    void SetIndexingSettings(const IndexingSettings &input);
    IndexingSettings GetIndexingSettings() const;
    PixelMaskStatistics GetPixelMaskStatistics() const;

    void GetStartMessageFromBuffer(std::vector<uint8_t> &v);
    void GetImageFromBuffer(std::vector<uint8_t> &v, int64_t image_number = -1);
    ImageBufferStatus GetImageBufferStatus() const;
    void ClearImageBuffer() const;
    void AddDetectorSetup(const DetectorSetup& setup); // Not thread safe, only during setup

    std::optional<ScanResult> GetScanResult() const;
};


#endif //JUNGFRAUJOCH_JFJOCHSTATEMACHINE_H