// 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 "PreprocessedPixel.h"

class ImagePreprocessorBufferGPU : public ImagePreprocessorBuffer {
    CudaDevicePtr<int32_t> gpu_image;
    CudaRegisteredVector<int32_t> buffer_reg;
    // Set per image by the preprocessor: true once it has written this frame in its source's 16-bit
    // width rather than widening it, with the code that frame uses for a saturated pixel.
    bool narrow = false;
    uint16_t narrow_sat_code = 0;

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

    // The device image is one allocation either way - the narrow form is the same memory read two
    // bytes at a time, so a run that switches width allocates nothing and frees nothing.
    bool IsNarrow() const override { return narrow; }
    const uint16_t *getGPUBufferNarrow() const override;
    uint16_t NarrowSatCode() const override { return narrow_sat_code; }
    void SetNarrow(bool v, uint16_t sat_code) override { narrow = v; narrow_sat_code = sat_code; }
    void Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const override;
};