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

#include "StillsPrePass.h"

#include "../common/JFJochException.h"
#include "../reader/JFJochReader.h"
#include "../reader/JFJochStartMessageReader.h"
#include "../rugnux/RugnuxLive.h"

namespace {
    // The held images as a collection of their own, numbered from 0 as the collection numbers them.
    class HeldImageReader : public JFJochReader {
        const std::map<int64_t, std::shared_ptr<HeldImage>> &images;
        uint64_t number_of_images;

        bool LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset, DataMessage &message,
                         std::vector<uint8_t> &, int64_t image_number, bool) override {
            const auto it = images.find(image_number);
            if (!dataset || it == images.end())
                return false;
            message = it->second->message;
            return true;
        }
    public:
        HeldImageReader(const StartMessage &start, const std::map<int64_t, std::shared_ptr<HeldImage>> &in_images)
            : images(in_images), number_of_images(in_images.empty() ? 0 : in_images.rbegin()->first + 1) {
            SetStartMessage(DatasetFromStartMessage(start, default_experiment));
        }

        [[nodiscard]] uint64_t GetNumberOfImages() const override { return number_of_images; }
        void Close() override {}

        bool ReadRawImage(int64_t image_number, JFJochReaderRawImage &ret) override {
            const auto it = images.find(image_number);
            if (it == images.end())
                return false;
            ret.image = it->second->message.image;
            return true;
        }

        std::vector<SpotToSave> ReadSpots(int64_t) const override { return {}; }
    };
}

StillsPrePassResult RunStillsPrePass(const StartMessage &start, const std::map<int64_t, std::shared_ptr<HeldImage>> &images,
                                     int nthreads) {
    HeldImageReader reader(start, images);
    const auto dataset = reader.GetDataset();
    const RugnuxLiveSetup setup = MakeRugnuxStillsPrePassSetup(*dataset, nthreads);

    Rugnux rugnux(reader, setup.experiment, *dataset->pixel_mask, setup.config);
    const ProcessResult result = rugnux.Run();
    Rugnux::ReleasePinnedPool();

    return StillsPrePassResult{
        // Where the beam lands: the receiver analyses with the detector square to it, as stated, and a tilt
        // the run refined moves its PONI, not this.
        .beam_x_pxl = static_cast<float>(result.used_direct_beam_x_pxl),
        .beam_y_pxl = static_cast<float>(result.used_direct_beam_y_pxl),
        .distance_mm = result.used_distance_mm,
        .mask = rugnux.GetPixelMask(),
        .images = images.size()
    };
}