// 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"

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 UINT16_MAX or more strong pixels (the connected-component search rejects it anyway), so that
    // is the largest gather that can be asked for.
    static constexpr size_t MAX_GATHER = UINT16_MAX;
    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;

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;
};