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

#pragma once

#include <cstdint>
#include <memory>
#include <vector>

#include "../../common/CompressedImage.h"
#include "../image_preprocessing/BSLZ4DecoderGPU.h"
#include "../indexing/CUDAMemHelpers.h"
#include "ShadowMaskGPU.h"

// The beam-stop projection accumulated on the device: only the compressed chunk crosses PCIe, and
// both the decode and the per-pixel maximum / sum / count run on the GPU. The projection comes back
// once, when the mask is read.
//
// This accumulates the RAW STORED COUNTS, with the pixel type's sentinel skipped, exactly as the
// host ShadowFinder does - NOT the preprocessed image. The two are not the same thing and the
// shadow is defined on the raw one. Sums and counts are integers, so the device result is identical
// to the host's rather than merely close.
//
// Only bitshuffle+LZ4 images are handled; anything else is left to the host path.
class ShadowAccumulatorGPU {
    const size_t npixels;
    int device = 0;                    // the card it was built on, where everything below lives
    std::shared_ptr<CudaStream> stream;
    std::unique_ptr<BSLZ4DecoderGPU> decoder;

    // Frames are decoded into a batch and folded into the projection together. The fold reads and
    // writes 20 bytes per pixel whatever the batch holds, so doing it once per frame would spend
    // most of the bandwidth on the accumulator rather than on the data. Four is where that stops
    // mattering - the fold is then a few percent of the phase - and every frame beyond it is a
    // further full frame of device memory, which costs more in cudaMalloc (device-synchronizing)
    // than it saves.
    static constexpr int BATCH = 4;
    CudaDevicePtr<uint8_t> raw;        // BATCH decoded frames, still in their stored type
    CudaDevicePtr<int64_t> gpu_max;
    CudaDevicePtr<int64_t> gpu_sum;
    CudaDevicePtr<uint32_t> gpu_count;
    size_t frame_bytes = 0;            // stride of one frame within `raw`
    size_t raw_capacity = 0;           // bytes actually allocated for `raw`
    int pending = 0;                   // frames decoded but not yet folded in
    CompressedImageMode pending_mode = CompressedImageMode::Uint32;
    uint32_t frames = 0;
    int blocks = 0;

    void EnsureRawCapacity(size_t bytes_per_frame);
    void FoldPending();

public:
    explicit ShadowAccumulatorGPU(size_t npixels);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~ShadowAccumulatorGPU() { if (stream) cudaStreamSynchronize(*stream); }

    // True when this image can be decoded and accumulated on the device.
    static bool Supports(const CompressedImage &image);

    // Decode and accumulate one frame. Throws if the image is not one this can take (ask Supports
    // first) or if the chunk does not decode.
    void Add(const CompressedImage &image);

    [[nodiscard]] uint32_t GetFrameCount() const { return frames; }
    [[nodiscard]] int GetDevice() const { return device; }

    // Add the projection another card accumulated to this one. Integer sums and a maximum taken as
    // Add takes it, so the result does not depend on how the frames were split between the cards or
    // on the order the cards are folded in. Called with this accumulator's card current.
    void Fold(ShadowAccumulatorGPU &other);

    // The beam-stop mask and the mean projection, made where the projection is (ShadowMaskGPU.h), so
    // that it does not have to come back at all.
    std::vector<uint32_t> Mask(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask);
    std::vector<float> MeanProjection(const std::vector<uint32_t> &pixel_mask);

    // Bring the projection back to the host, folding in whatever the last batch still holds. Cheap
    // to call once; it moves 20 bytes per pixel.
    void Download(std::vector<int64_t> &max_value, std::vector<int64_t> &sum_value,
                  std::vector<uint32_t> &valid_count);
};