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

#pragma once

#include <vector>
#include <cstdint>
#include <cstddef>

class ImagePreprocessorBuffer {
protected:
    // The image size, kept apart from `buffer` so size() still answers when the host copy was not
    // allocated at all (see the host_mirror constructor).
    const size_t npixels;
    std::vector<int32_t> buffer;

    // host_mirror = false leaves `buffer` empty. On the GPU path the preprocessed image lives on the
    // device and no CPU engine reads it back, so the four-bytes-per-pixel host copy would be
    // allocated, zeroed and page-locked for nothing - 72 MB per worker on a 16M-pixel detector.
    ImagePreprocessorBuffer(size_t npixels, bool host_mirror);
public:
    explicit ImagePreprocessorBuffer(size_t npixels);
    virtual ~ImagePreprocessorBuffer() = default;

    // Standard CPU operation. Defined here rather than in the .cpp: these are called once per pixel by
    // the CPU preprocessor, the CPU azimuthal integrator and the CPU spot finder, and the build has no
    // link-time optimisation - so out of line they are a real call per pixel, and they stop the callers
    // vectorising at all. Inlining them costs about a third of the per-image CPU work on every path
    // that does not have a GPU, and leaves the results bit-identical.
    std::vector<int32_t> &getBuffer() { return buffer; }
    const std::vector<int32_t> &getBuffer() const { return buffer; }
    int32_t &operator[](size_t i) { return buffer[i]; }
    const int32_t &operator[](size_t i) const { return buffer[i]; }
    size_t size() const { return npixels; }

    int32_t *data() { return buffer.data(); }
    const int32_t *data() const { return buffer.data(); }

    // GPU operations (overriden in ImagePreprocessorBufferGPU
    virtual int32_t *getGPUBuffer();
    virtual const int32_t *getGPUBuffer() const;

    // Values of the pixels with the given flat indices. Not just operator[] in a loop: on the GPU the
    // preprocessed image can live in device memory only, and then this gathers them there.
    virtual void Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const;
};