// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <random>
#include <cstring>
#include <limits>

#include "../image_analysis/image_preprocessing/BSLZ4DecoderGPU.h"
#include "../compression/JFJochCompressor.h"
#include "../compression/JFJochDecompress.h"

// The GPU decoder must agree with the CPU one BYTE FOR BYTE, on data produced by our own compressor,
// for every element size the detectors emit - including the 8-bit DECTRIS modes, which take a
// different branch in bitshuf_decode_block (bit un-transpose only, no byte interleave).
//
// Images are built to exercise what the LZ4 format actually does on detector data: long runs of a
// repeated byte (offset == 1 matches, the overlapping-match path), isolated bright pixels
// (literals), and a noisy region (short matches at assorted offsets). A uniformly random image
// would be almost incompressible and would never reach the match code at all.
namespace {

template <class T>
std::vector<T> MakeDetectorLikeImage(size_t npixels, uint32_t seed) {
    std::mt19937 rng(seed);
    std::vector<T> img(npixels, 0);          // sparse background: long zero runs

    // A band of low-level noise, so matches are short and offsets vary.
    for (size_t i = npixels / 4; i < npixels / 2; i++)
        img[i] = static_cast<T>(rng() % 7);

    // Bright, isolated spots - these become literals.
    for (size_t s = 0; s < 64; s++) {
        const size_t c = rng() % npixels;
        for (size_t d = 0; d < 9 && c + d < npixels; d++)
            img[c + d] = static_cast<T>(std::numeric_limits<T>::max() / (2 + (d % 3)));
    }
    // A run of one repeated non-zero value, the classic offset==1 match.
    for (size_t i = npixels * 3 / 4; i < npixels * 3 / 4 + 5000 && i < npixels; i++)
        img[i] = static_cast<T>(42);
    return img;
}

template <class T>
void RoundTrip(CompressedImageMode mode, size_t width, size_t height, uint32_t seed) {
    const size_t npixels = width * height;
    const auto original = MakeDetectorLikeImage<T>(npixels, seed);

    // Compress with the production compressor, so the container is exactly what the pipeline reads.
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(original);
    REQUIRE(!compressed.empty());

    const CompressedImage image(compressed.data(), compressed.size(), width, height, mode,
                                CompressionAlgorithm::BSHUF_LZ4);
    REQUIRE(BSLZ4DecoderGPU::Supports(image));
    REQUIRE(image.GetUncompressedSize() == npixels * sizeof(T));

    // CPU reference: the same call the host path makes.
    std::vector<uint8_t> cpu_buffer;
    const uint8_t *cpu_out = image.GetUncompressedPtr(cpu_buffer);
    REQUIRE(std::memcmp(cpu_out, original.data(), npixels * sizeof(T)) == 0);

    auto stream = std::make_shared<CudaStream>();
    BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream);
    CudaDevicePtr<uint8_t> gpu_out(npixels * sizeof(T));
    decoder.Decode(image, gpu_out.get());
    REQUIRE(cudaStreamSynchronize(*stream) == cudaSuccess);

    std::vector<T> gpu_result(npixels);
    REQUIRE(cudaMemcpy(gpu_result.data(), gpu_out.get(), npixels * sizeof(T),
                       cudaMemcpyDeviceToHost) == cudaSuccess);

    REQUIRE(std::memcmp(gpu_result.data(), original.data(), npixels * sizeof(T)) == 0);
}

} // namespace

TEST_CASE("BSLZ4DecoderGPU_MatchesCPU_AllElementSizes", "[BSLZ4DecoderGPU]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    // Sizes chosen so the last block is partial and the leftover tail (the elements bitshuffle
    // leaves uncompressed because they do not fill a multiple of 8) is non-empty on some of them.
    RoundTrip<uint8_t>(CompressedImageMode::Uint8, 1030, 517, 1);
    RoundTrip<int8_t>(CompressedImageMode::Int8, 1030, 517, 2);
    RoundTrip<uint16_t>(CompressedImageMode::Uint16, 1030, 517, 3);
    RoundTrip<int16_t>(CompressedImageMode::Int16, 1030, 517, 4);
    RoundTrip<uint32_t>(CompressedImageMode::Uint32, 1030, 517, 5);
    RoundTrip<int32_t>(CompressedImageMode::Int32, 1030, 517, 6);
}

TEST_CASE("BSLZ4DecoderGPU_MatchesCPU_LargeFrame", "[BSLZ4DecoderGPU]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    // Many blocks, so the per-block descriptor scan and the one-warp-per-block launch are exercised
    // at a realistic scale rather than on a handful of blocks.
    RoundTrip<uint32_t>(CompressedImageMode::Uint32, 2068, 2162, 7);
}

// A decoder that cannot handle an image must SAY so rather than produce something wrong: the caller
// relies on Supports() to decide whether the host route is needed.
TEST_CASE("BSLZ4DecoderGPU_DeclinesWhatItCannotDecode", "[BSLZ4DecoderGPU]") {
    std::vector<uint8_t> dummy(1024, 0);
    const size_t w = 16, h = 16;

    CHECK_FALSE(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32,
                        CompressionAlgorithm::BSHUF_ZSTD)));
    CHECK_FALSE(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32,
                        CompressionAlgorithm::BSHUF_ZSTD_RLE)));
    CHECK_FALSE(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32,
                        CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF)));
    CHECK_FALSE(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32,
                        CompressionAlgorithm::NO_COMPRESSION)));
    CHECK_FALSE(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Float32,
                        CompressionAlgorithm::BSHUF_LZ4)));

    CHECK(BSLZ4DecoderGPU::Supports(
        CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32,
                        CompressionAlgorithm::BSHUF_LZ4)));
}

// A malformed container comes off the network or off disk, so it must throw rather than run off
// the end of a buffer on the device.
TEST_CASE("BSLZ4DecoderGPU_RejectsMalformed", "[BSLZ4DecoderGPU]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    const size_t width = 128, height = 128, npixels = width * height;
    const auto original = MakeDetectorLikeImage<uint32_t>(npixels, 11);
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(original);

    auto stream = std::make_shared<CudaStream>();
    BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream);
    CudaDevicePtr<uint8_t> gpu_out(npixels * sizeof(uint32_t));

    // Truncated mid-stream: the block header promises more than is there.
    const CompressedImage truncated(compressed.data(), compressed.size() / 2, width, height,
                                    CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4);
    CHECK_THROWS(decoder.Decode(truncated, gpu_out.get()));

    // Shorter than the 12-byte container header.
    const CompressedImage stub(compressed.data(), 8, width, height,
                               CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4);
    CHECK_THROWS(decoder.Decode(stub, gpu_out.get()));
}

#endif