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

#pragma once

#include <QObject>
#include <QString>
#include <QVector>
#include <QRect>
#include <QMutex>
#include <QTimer>
#include <QElapsedTimer>

#include "../reader/JFJochHDF5Reader.h"
#include "../reader/JFJochCBFReader.h"
#include "../reader/JFJochMarCCDReader.h"
#include "../reader/JFJochSMVReader.h"
#include "../reader/JFJochXFELReader.h"
#include "../reader/JFJochXTCReader.h"
#include "../common/Logger.h"
#include "JFJochHttpReader.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../common/BraggIntegrationSettings.h"
#include "../common/ScalingSettings.h"
#include "../image_analysis/bragg_prediction/BraggPrediction.h"
#include "../image_analysis/LoadFCalcFromMtz.h"   // ReferenceMtzData + LoadReferenceMtz
#include "../gemmi_gph/gemmi/symmetry.hpp"
#include "SimpleImage.h"
#include "ReferenceMtzInfo.h"
#include "../common/MovingAverage.h"

Q_DECLARE_METATYPE(std::shared_ptr<const JFJochReaderDataset>)
Q_DECLARE_METATYPE(std::shared_ptr<const JFJochReaderImage>)

Q_DECLARE_METATYPE(DiffractionExperiment)
Q_DECLARE_METATYPE(SpotFindingSettings)
Q_DECLARE_METATYPE(IndexingSettings)
Q_DECLARE_METATYPE(AzimuthalIntegrationSettings)
Q_DECLARE_METATYPE(BraggIntegrationSettings)
Q_DECLARE_METATYPE(ScalingSettings)

Q_DECLARE_METATYPE(UnitCell)
Q_DECLARE_METATYPE(std::shared_ptr<const SimpleImage>)
Q_DECLARE_METATYPE(BrokerStatus)
Q_DECLARE_METATYPE(ROIDefinition)

// Everything a reprocessing job needs, captured from the worker's current state in one locked read.
struct ReprocessingInputs {
    bool valid = false;
    QString file;
    DiffractionExperiment experiment;
    PixelMask pixel_mask;
    SpotFindingSettings spot_finding;
    std::vector<MergedReflection> reference_data;   // empty unless a reference (MTZ / SF-mmCIF) is loaded
    // The rest of what the reference fixes, as `rugnux -z` does: its free-R partition, its cell and
    // the space group in the setting it is written in. Empty/unset without a reference.
    bool reference_has_free_flags = false;
    std::optional<UnitCell> reference_cell;
    std::optional<gemmi::SpaceGroup> reference_space_group;
    QString reference_path;
    QString reference_column;
    QString model_path;   // atomic model for post-merge validation (`rugnux --model`); empty = none
    // The rotation axis the file was opened with, before the settings panel had a say. `experiment`
    // carries the panel's, which may have cleared it; this is what a fresh `rugnux` run would see.
    std::optional<GoniometerAxis> file_goniometer;
    // Whether the file states ice-ring flagging; unset, a run decides it as `rugnux` does.
    std::optional<bool> file_detect_ice_rings;
};

#include "RunData.h"

class JFJochImageReadingWorker : public QObject {
    Q_OBJECT

public:
    // HTTPSync: follow latest image + dataset. HTTPSyncDataset: follow dataset/plots/image count
    // but keep the displayed image frozen (entered by manually selecting an image while following).
    enum class AutoloadMode {HTTPSync, HTTPSyncDataset, Movie, None};
    Q_ENUM(AutoloadMode)
private:
    mutable QMutex m;
    bool http_mode = false;

    AutoloadMode autoload_mode = AutoloadMode::None;

    // A stored file is read by one of these, whichever its format needs; file_reader points at the
    // one in use. The HDF5-only extras - calibration images and reprocessing snapshots - always go
    // through hdf5_reader. A raw image sweep has no calibration images and its runs are files of their
    // own (raw_runs_ below), so cbf_mode covers all three of the raw formats rather than CBF alone.
    JFJochHDF5Reader hdf5_reader;
    JFJochCBFReader cbf_reader;
    JFJochMarCCDReader marccd_reader;
    JFJochSMVReader smv_reader;
    JFJochXFELReader xfel_reader;
    JFJochXTCReader xtc_reader;
    QStringList pending_files_;
    JFJochReader *file_reader = &hdf5_reader;
    bool cbf_mode = false;      // "a raw sweep is open", whichever of the three formats it is

    // The runs of a raw sweep. A run of an HDF5 file is a snapshot - its metadata read back over the
    // file's own images - but a raw sweep has no HDF5 to hang one on. Its run's _process.h5 carries
    // the images as well, so showing that run means reading that file (with hdf5_reader), and showing
    // "Original" means reading the sweep again (raw_reader_).
    struct RawRun {
        std::string path;
        std::shared_ptr<const JFJochReaderDataset> dataset;
    };
    JFJochReader *raw_reader_ = nullptr;
    std::map<std::string, RawRun> raw_runs_;
    std::string raw_active_ = "Original";

    JFJochHttpReader http_reader;

    QString current_file;

    DiffractionExperiment curr_experiment;
    // Once the user edits ROIs they override whatever the file carried, for this and
    // every subsequently loaded image, until a new file is opened.
    std::optional<ROIDefinition> roi_override_;
    // The rotation axis of the open file, captured before the settings panel can change it — see
    // ReprocessingInputs::file_goniometer.
    std::optional<GoniometerAxis> file_goniometer_;
    std::optional<bool> file_detect_ice_rings_;
    std::map<std::string, QString> run_labels_;   // snapshot id -> editable display label
    IndexingSettings indexing_settings;
    AzimuthalIntegrationSettings azint_settings;
    BraggIntegrationSettings bragg_settings;
    ScalingSettings scaling_settings;
    // Reference reflections (from a reference MTZ), kept independent of the loaded dataset - they
    // survive file switches and feed reference-based scaling / CCref reprocessing jobs.
    std::optional<ReferenceMtzData> reference_;
    QString reference_path_;   // where reference_ was loaded from (for the command line)
    QString model_path_;
    std::unique_ptr<IndexerThreadPool> indexing;

    std::shared_ptr<JFJochReaderImage> current_image_ptr;
    std::unique_ptr<AzimuthalIntegrationMapping> azint_mapping;
    std::unique_ptr<MXAnalysisWithoutFPGA> image_analysis;
    std::unique_ptr<IndexAndRefine> index_and_refine;

    // Azimuthal profile buffer of the most recently analyzed image (filled by Analyze).
    std::unique_ptr<AzimuthalIntegrationProfile> last_profile_;

    SpotFindingSettings spot_finding_settings;

    std::optional<int64_t> current_image;
    int64_t current_summation = 1;
    int64_t total_images = 0;
    int64_t jump_value = 1;
    Logger logger{"jfjoch_viewer"};

    bool auto_reanalyze = false;

    QTimer *autoload_timer;
    int autoload_interval = 500; // milliseconds

    // Broker status polling (independent of image sync, runs whenever connected over HTTP)
    QTimer *status_timer = nullptr;
    int status_interval = 1000; // milliseconds
    bool http_connected = false;

    // Adaptive autoload interval based on recent load+analysis time
    MovingAverage autoload_ms_ma{8};          // window size (tune as needed)
    int autoload_interval_min_ms = 50;        // 20 Hz is the top performance!
    int autoload_interval_max_ms = 2000;      // 0.5 Hz as bottom limit
    float autoload_safety_factor = 2.0f;

    // Backpressure: cap how many frames may be "in flight" to the GUI (emitted but not yet rendered
    // and acked via ImageConsumed). The autoload timer won't produce another frame while the cap is
    // reached, so a fast producer can't grow the queued-event backlog without bound.
    int images_in_flight = 0;
    int max_images_in_flight = 4;

    // Same backpressure for the dataset stream (see EmitDatasetLoaded_i / DatasetConsumed): each
    // queued datasetLoaded event pins a full dataset (mask + per-image plots), so cap how many may
    // be in flight to the GUI. This also coalesces dataset-only follow (HTTPSyncDataset), which emits
    // no image and so was never gated by images_in_flight - stale datasets are dropped, the next tick
    // sends the latest. Without this the datasetLoaded backlog grew without bound (100 GB OOM).
    int datasets_in_flight = 0;
    int max_datasets_in_flight = 2;

    // File open retry/back-off (GPFS via NFSv4 visibility lag)
    struct PendingLoadRequest {
        QString filename;
        qint64 image_number = 0;
        qint64 summation = 1;
    };

    QTimer *file_open_retry_timer = nullptr;
    PendingLoadRequest pending_load;
    bool file_open_retry_active = false;
    bool file_open_retry_warned = false;
    int file_open_retry_attempts = 0;
    int file_open_retry_delay_ms = 100;
    int file_open_retry_delay_max_ms = 10000;
    QElapsedTimer file_open_retry_elapsed;

    void ResetFileOpenRetry_i();
    void ScheduleFileOpenRetry_i(const QString& reason);

    void LoadFile_i(const QString &filename, qint64 image_number, qint64 summation, bool retry);
    void LoadImage_i(int64_t image_number, int64_t summation);
    void EmitImageLoaded_i(const std::shared_ptr<const JFJochReaderImage>& image);
    void EmitDatasetLoaded_i(const std::shared_ptr<const JFJochReaderDataset>& dataset);
    void RefreshDatasetOnly_i();
    void SetHttpConnected_i(bool connected, const QString &addr);
    // Assumes m locked. A 401 from the broker: clear the display, stop following, tell the status bar.
    void HandleUnauthorized_i(const QString &filename);
    void ReanalyzeImage_i();
    void UpdateDataset_i(const std::optional<DiffractionExperiment>& experiment);
    void UpdateAzint_i(const JFJochReaderDataset *dataset);
    void UpdateUserMask_i(const std::vector<uint32_t> &mask);
    void SetROIDefinition_i(const ROIDefinition &rois);
    void ApplyROIOverrideToImage_i();   // rewrite current_image_ptr's dataset with roi_override_
    void setAutoLoadMode_i(AutoloadMode mode);
    void ActivateSnapshot_i(const QString &name);
    void EmitRuns_i();
    std::string ActiveRun_i() const;
    QStringList RunNames_i() const;   // emit runsChanged with every snapshot dataset + its label
    void EmitReferenceInfo_i();   // emit referenceMtzChanged from reference_ + current data

signals:
    void datasetLoaded(std::shared_ptr<const JFJochReaderDataset>);
    void imageLoaded(std::shared_ptr<const JFJochReaderImage>);
    void imageStatsUpdated(std::shared_ptr<const JFJochReaderImage>);
    void imageNumberChanged(int64_t total_images, int64_t current_image);
    void setRings(const QVector<float> &v);
    void simpleImageLoaded(std::shared_ptr<const SimpleImage> image);
    void autoloadChanged(AutoloadMode mode);
    void fileLoadError(QString title, QString message);
    void fileLoadRetryStatus(bool active, QString message);
    void brokerStatusUpdated(BrokerStatus status);
    void httpConnectionChanged(bool connected, QString addr);
    // The broker refused the dataset (401): the display was cleared, nothing else is raised.
    void httpAuthRequired(QString message);
    void liveRateChanged(double hz);
    void snapshotsChanged(QStringList names, QString active);
    void runsChanged(QVector<RunData> runs, QString active_id);
    void fileOpened();   // a new file/stream was loaded - listeners should reset per-file state
    void referenceMtzChanged(ReferenceMtzInfo info);   // a reference MTZ was (un)loaded - update the dock

public:
    JFJochImageReadingWorker(const SpotFindingSettings &settings, const DiffractionExperiment& experiment, QObject *parent = nullptr);
    ~JFJochImageReadingWorker() override = default;

    // Captures the current file + experiment + mask + spot-finding settings so a processing job
    // uses exactly the same settings as interactive analysis. Thread-safe (locks the worker mutex).
    ReprocessingInputs GetReprocessingInputs() const;

private slots:
    void AutoLoadTimerExpired();
    void StatusTimerExpired();
    void FileOpenRetryTimerExpired();

public slots:
    void LoadFile(const QString &filename, qint64 image_number, qint64 summation, bool retry);
    // The files of one dataset the next LoadFile opens together (the open dialog's selection); used
    // only if that LoadFile names one of them.
    void SetFileList(const QStringList &files);
    void CloseFile();
    void LoadImage(int64_t image_number, int64_t summation);

    void UpdateDataset(const DiffractionExperiment& experiment);

    // The calibrant as its ring q values (2*pi/d, 1/A) rather than a cell, so a standard whose rings
    // are measured rather than enumerated - hexagonal ice - can be calibrated against too.
    void FindCenter(QVector<float> ring_q, bool guess, bool refine_tilt);

    void Analyze();

    void UpdateSpotFindingSettings(const SpotFindingSettings &settings, const IndexingSettings &indexing, int64_t max_spots);
    void ReanalyzeImages(bool input);
    void UpdateAzintSettings(const AzimuthalIntegrationSettings& settings);
    void UpdateBraggIntegrationSettings(BraggIntegrationSettings settings);
    void UpdateScalingSettings(ScalingSettings settings);

    // Load (column empty = auto-select) or clear a reference MTZ; kept independent of the dataset.
    void SetReferenceMtz(QString path, QString column);
    void ClearReferenceMtz();
    // Bearer token sent with every broker request (empty = none).
    void SetHttpToken(QString token);
    // Atomic model a processing job validates its merged data against (empty = none).
    void SetModelPath(QString path);

    void SetROIDefinition(const ROIDefinition &rois);
    void DownloadROIsFromServer();
    void UploadROIsToServer();
    void MaskFromSelectedROI(QString name, bool add);   // rasterise an ROI into the user mask

    void SaveUserMaskTIFF(QString filename);
    void LoadUserMaskTIFF(QString filename, bool replace);
    void UploadUserMask();
    void ClearUserMask();

    void LoadCalibration(QString dataset);
    void setAutoLoadMode(AutoloadMode mode);
    void setAutoLoadJump(int64_t val);

    // GUI acks that it has consumed (rendered) one frame, freeing an in-flight slot. Drives the
    // autoload backpressure so the reader can't outrun the viewer (see images_in_flight).
    void ImageConsumed();

    // GUI ack that it has consumed (fanned out + rendered) one dataset, freeing an in-flight slot.
    // Drives the datasetLoaded backpressure so the reader can't outrun the viewer's dataset drain.
    void DatasetConsumed();

    // Register a reprocessing result (_process.h5) as a metadata snapshot over the same images and
    // make it the active dataset; SetActiveSnapshot switches between registered snapshots.
    void RegisterProcessingSnapshot(QString id, QString label, QString master_path);
    void SetActiveSnapshot(QString name);
    void RenameRun(QString id, QString label);   // change a run's display label (legend)
    void RemoveRun(QString id);                  // drop a reprocessing snapshot (not "Original")
};