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

#pragma once

#include <memory>

#include "../../common/CompressedImage.h"
#include "../indexing/CUDAMemHelpers.h"

// What one tile of the byte-offset stream does to a decode that enters it at each of the 15 byte
// positions a token can straddle the tile boundary by: where the decode leaves the tile (the offset
// into the next one), how many pixels it emitted, and what they added to the running value. Entry and
// exit 15 is "the stream has ended", which every later tile passes on.
struct CBFTileMap {
    static constexpr int DEAD = 15;
    uint32_t count[16];
    uint32_t sum[16];      // modulo 2^32, which is all an int32 pixel keeps of the int64 running value
    uint8_t exit[16];
};

// x-CBF_BYTE_OFFSET decoded ON THE DEVICE, exactly as DecodeCBFByteOffset does on the host: the same
// int32 pixels, and the same refusal of a stream that ends before the image does.
//
// The format is a running sum of deltas whose tokens are 1, 3, 7 or 15 bytes long, so where a token
// starts depends on every token before it. The stream is cut into short parts; for each part and each
// of the 15 offsets a decode can enter it at, one thread works out where it leaves, how many pixels it
// emits and what they sum to (the walks from different offsets meet within a few bytes, after which
// they are the same walk). Composing those maps is a scan - over tiles of a few parts, whose maps are
// the parts' composed - which gives every part its true entry offset, its first pixel and the running
// value there; a second walk of each part then writes the pixels.
class CBFDecoderGPU {
    std::shared_ptr<CudaStream> stream;

    // The reader's buffer is pageable; copied here first, the upload is one DMA (see BSLZ4DecoderGPU).
    CudaHostPtr<uint8_t> host_compressed;
    CudaDevicePtr<uint8_t> gpu_compressed;
    size_t compressed_capacity = 0;

    CudaDevicePtr<CBFTileMap> gpu_part_maps;
    CudaDevicePtr<CBFTileMap> gpu_maps;
    CudaDevicePtr<CBFTileMap> gpu_entries;   // exclusive scan of gpu_maps: the state each tile is entered in
    CudaDevicePtr<uint8_t> gpu_scan_temp;
    size_t scan_temp_bytes = 0;
    size_t tile_capacity = 0;

    CudaDevicePtr<uint32_t> gpu_total;   // pixels the whole stream decodes to
    CudaHostPtr<uint32_t> host_total;
    size_t last_npixels = 0;

    CudaEvent decode_start;
    CudaEvent decode_stop;

public:
    explicit CBFDecoderGPU(std::shared_ptr<CudaStream> stream);
    ~CBFDecoderGPU() { if (stream) cudaStreamSynchronize(*stream); }

    // Upload the stream and queue its decode into gpu_out (image width x height int32 pixels) on the
    // decoder's stream. The caller synchronises, then calls ThrowIfDecodeFailed().
    void Decode(const CompressedImage &image, int32_t *gpu_out);

    // Report a stream that ended before the image did. MUST be called after the caller has
    // synchronised the stream.
    void ThrowIfDecodeFailed() const;

    // Device time of the last upload and decode, in seconds. Valid after the caller synchronises.
    [[nodiscard]] float GetDecodeTime_s() const;
};