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

#pragma once

#include <mutex>
#include <atomic>
#include <chrono>

#include "../common/DiffractionExperiment.h"
#include "../image_analysis/spot_finding/SpotFindingSettings.h"
#include "JFJochReceiverOutput.h"
#include "../image_puller/ZMQImagePuller.h"
#include "../common/ImageBuffer.h"
#include "../image_pusher/ImagePusher.h"
#include "../preview/ZMQPreviewSocket.h"
#include "../preview/ZMQMetadataSocket.h"
#include "../preview/PreviewImage.h"
#include "../common/JFJochReceiverPlots.h"
#include "LossyFilter.h"
#include "../common/MovingAverage.h"
#include "../common/ScanResultGenerator.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"
#include "../image_analysis/IndexAndRefine.h"

class JFJochReceiver {
protected:
    Logger &logger;

    DiffractionExperiment experiment;

    SpotFindingSettings spot_finding_settings;
    std::mutex spot_finding_settings_mutex;
    SpotFindingSettings GetSpotFindingSettings();

    std::atomic<bool> writer_queue_full{false};
    std::atomic<bool> cancelled{false};
    bool push_images_to_writer;

    std::atomic<size_t> compressed_size{0};
    std::atomic<uint64_t> uncompressed_size{0};
    std::atomic<size_t> images_sent{0};
    std::atomic<size_t> images_collected{0};
    std::atomic<size_t> images_skipped{0};

    std::chrono::time_point<std::chrono::system_clock> start_time;
    std::chrono::time_point<std::chrono::system_clock> end_time;

    // After END, how long the image pusher may take to return every image buffer it still holds
    constexpr static auto SendFinalizeTimeout = std::chrono::seconds(30);

    int64_t max_image_number_sent = 0;
    std::mutex max_image_number_sent_mutex;
    void UpdateMaxImageSent(int64_t image_number);

    int64_t max_image_number_received = 0;
    std::mutex max_image_number_received_mutex;
    void UpdateMaxImageReceived(int64_t image_number);

    ImageBuffer &image_buffer;
    ImagePusher &image_pusher;

    ZMQPreviewSocket *zmq_preview_socket;
    ZMQMetadataSocket *zmq_metadata_socket;

    JFJochReceiverCurrentStatus &current_status;
    JFJochReceiverPlots &plots;

    ScanResultGenerator scan_result;

    LossyFilter serialmx_filter;

    MovingAverage saturated_pixels{1000};
    MovingAverage error_pixels{1000};
    MovingAverage roi_beam_npixel{1000};
    MovingAverage roi_beam_sum{1000};

    std::optional<int64_t> images_written;

    std::vector<std::unique_ptr<ADUHistogram>> adu_histogram_module;
    PixelMask pixel_mask;

    std::unique_ptr<AzimuthalIntegrationMapping> az_int_mapping;

    std::optional<uint64_t> max_delay;
    std::mutex max_delay_mutex;

    IndexAndRefine indexer;
    std::optional<CrystalLattice> rotation_indexing_lattice;
    std::optional<LatticeMessage> rotation_indexing_lattice_type;

    std::string writer_error;

    void UpdateMaxDelay(uint64_t delay);
    JFJochReceiverStatus GetStatus() const;
    void SendStartMessage();
    void SendEndMessage();
    void SaveStartMessageToImageBuffer(const StartMessage &msg);
    // The last word on the end message before it is sent: a receiver whose analysis is not the
    // per-image one fills in what it found.
    virtual void AddToEndMessage(EndMessage &message) {}
    virtual void Cancel(const JFJochException &e);
public:
    JFJochReceiver(const DiffractionExperiment& experiment,
                   ImageBuffer &image_buffer,
                   ImagePusher &image_pusher,
                   JFJochReceiverCurrentStatus &current_status,
                   JFJochReceiverPlots &plots,
                   const SpotFindingSettings &spot_finding_settings,
                   Logger &logger,
                   const PixelMask &pixel_mask,
                   ZMQPreviewSocket *zmq_preview_socket = nullptr,
                   ZMQMetadataSocket *zmq_metadata_socket = nullptr,
                   IndexerThreadPool *indexing_thread_pool = nullptr);
    virtual ~JFJochReceiver();

    virtual void StopReceiver() = 0;

    virtual float GetEfficiency() const = 0;
    virtual void Cancel(bool silent);
    virtual float GetProgress() const = 0;

    virtual void SetSpotFindingSettings(const SpotFindingSettings &spot_finding_settings) = 0;

    virtual JFJochReceiverOutput GetFinalStatistics() const;
};