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

#pragma once

#include <memory>

#include "ImagePreprocessorBuffer.h"
#include "../indexing/CUDAMemHelpers.h"
#include "../spot_finding/SpotFindingSettings.h"

class ImagePreprocessorBufferGPU : public ImagePreprocessorBuffer {
    CudaDevicePtr<int32_t> gpu_image;
    CudaRegisteredVector<int32_t> buffer_reg;

    // Staging for Gather(). Its only caller is ImageSpotFinder::ExtractSpots, which gives up on a frame
    // with StrongPixelLimit or more strong pixels (the connected-component search rejects it anyway),
    // so that is the largest gather that can be asked for. It follows the detector, so this has to as
    // well - sized to a constant while the caller's bar was raised, the gather would run off the end
    // of the buffer rather than merely lose the frame.
    const size_t max_gather;
    CudaDevicePtr<uint32_t> gpu_gather_index;
    CudaDevicePtr<int32_t> gpu_gather_value;
    // Own stream: every analysis engine synchronises its own stream before it returns, so the device
    // image is final by the time a gather is asked for. The NULL stream would serialise all workers.
    CudaStream gather_stream;

    // GatherTiles: the squares' centres, the squares on the card, and the page-locked landing for them.
    // Grown to the largest frame's spots.
    size_t tile_capacity = 0;   // pixels of gpu_tiles and host_tiles
    CudaDevicePtr<int32_t> gpu_tile_centre;
    CudaDevicePtr<int32_t> gpu_tiles;
    PinnedVector<int32_t> tile_centre;
    PinnedVector<int32_t> host_tiles;

public:
    // host_mirror = false skips the host copy of the preprocessed image entirely (and with it the
    // page-locking): pass it when every engine reading this buffer runs on the device.
    explicit ImagePreprocessorBufferGPU(size_t npixel, bool host_mirror = true);

    int32_t *getGPUBuffer() override;
    const int32_t *getGPUBuffer() const override;
    void Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const override;
    const int32_t *GatherTiles(const std::vector<int32_t> &x, const std::vector<int32_t> &y,
                               int width, int half) override;
    const int32_t *HostImage() override;
};