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

#include <algorithm>
#include <type_traits>

#include "BSLZ4DecoderGPU.h"
#include "BSLZ4DecodeWarp.h"
#include "../../common/JFJochException.h"
#include "../../compression/JFJochDecompress.h"   // BSHUF_BLOCKED_MULT and the container layout

namespace {
    void cuda_err(cudaError_t val) {
        if (val != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
    }

    // One CUDA block per 256/32 = 8 LZ4 blocks; the parser itself is lz4_decode_block_warp, shared
    // with the fused decode+preprocess kernel so the two cannot decode a chunk differently.
    __global__ void lz4_decode_blocks(const uint8_t *__restrict__ src,
                                      const BSLZ4BlockDesc *__restrict__ desc,
                                      uint8_t *__restrict__ dst,
                                      uint32_t *__restrict__ status,
                                      int nblocks, uint32_t elem_size) {
        const int lane = threadIdx.x & 31;
        const int b = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
        if (b >= nblocks) return;

        const uint8_t *const ip = src + desc[b].in_off;
        uint8_t *const obase = dst + desc[b].out_off;
        const bool ok = lz4_decode_block_warp(ip, ip + desc[b].in_len,
                                              obase, obase + desc[b].nelem * elem_size, lane);
        if (lane == 0 && !ok)
            atomicExch(status, 1u);
    }

    __device__ __forceinline__ uint64_t transpose8(uint64_t x) {
        uint64_t t;
        t = (x ^ (x >> 7)) & 0x00aa00aa00aa00aaULL;  x = x ^ t ^ (t << 7);
        t = (x ^ (x >> 14)) & 0x0000cccc0000ccccULL; x = x ^ t ^ (t << 14);
        t = (x ^ (x >> 28)) & 0x00000000f0f0f0f0ULL; x = x ^ t ^ (t << 28);
        return x;
    }

    // The bitshuffle inverse, mirroring bitshuf_decode_block. One thread owns one group of 8 elements
    // across EVERY byte-plane, so after transposing its 8 bytes out of each plane it holds all bytes
    // of 8 complete elements and can write them straight out. That needs no staging buffer, which is
    // what keeps the kernel free of the 48 kB dynamic-shared-memory ceiling a block size taken from
    // the file header would otherwise run into.
    template<int ES>
    __global__ void bitshuffle_untranspose(const uint8_t *__restrict__ shuffled,
                                           const BSLZ4BlockDesc *__restrict__ desc,
                                           uint8_t *__restrict__ out, int nblocks) {
        const int b = blockIdx.x;
        if (b >= nblocks) return;

        const uint32_t size = desc[b].nelem;             // bytes per plane
        const uint8_t *in = shuffled + desc[b].out_off;
        uint8_t *dst = out + desc[b].out_off;
        const uint32_t n = size / 8;

        // The 8 elements a thread owns are contiguous and 8*ES-byte aligned, so they are assembled
        // whole and written through an element-typed pointer. Storing them byte by byte instead
        // costs about 4x on a full frame.
        using UT = typename std::conditional<ES == 1, uint8_t,
                   typename std::conditional<ES == 2, uint16_t, uint32_t>::type>::type;

        for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
            uint64_t x[ES];
            #pragma unroll
            for (int p = 0; p < ES; p++) {
                const uint8_t *pin = in + p * size;
                uint64_t a = 0;
                #pragma unroll
                for (int k = 0; k < 8; k++) a |= (uint64_t)pin[k * n + i] << (8 * k);
                x[p] = transpose8(a);
            }
            UT *dstT = reinterpret_cast<UT *>(dst) + i * 8;
            #pragma unroll
            for (int k = 0; k < 8; k++) {
                UT v = 0;
                #pragma unroll
                for (int p = 0; p < ES; p++) v |= (UT)((UT)((x[p] >> (8 * k)) & 0xff) << (8 * p));
                dstT[k] = v;
            }
        }
    }

    uint64_t be64(const uint8_t *p) { uint64_t v = 0; for (int i = 0; i < 8; i++) v = (v << 8) | p[i]; return v; }
    uint32_t be32(const uint8_t *p) { return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3]; }

    size_t elem_size_of(CompressedImageMode mode) {
        switch (mode) {
            // 8-bit is a real DECTRIS mode. bitshuf_decode_block takes a separate branch for
            // elem_size == 1 (bit un-transpose only, no byte interleave), and the kernel here
            // reproduces that for free: with one plane the interleave step degenerates to a copy.
            case CompressedImageMode::Int8:
            case CompressedImageMode::Uint8: return 1;
            case CompressedImageMode::Int16:
            case CompressedImageMode::Uint16: return 2;
            case CompressedImageMode::Int32:
            case CompressedImageMode::Uint32: return 4;
            default: return 0;   // float and RGB modes are never bitshuffled by this pipeline
        }
    }

    // The smallest a block can be on the wire: a 4-byte length plus at least one payload byte. Used
    // to reject a header whose declared block size implies more blocks than the chunk could hold,
    // before that count is turned into an allocation.
    constexpr size_t MIN_BLOCK_BYTES_ON_WIRE = 5;
}

bool BSLZ4DecoderGPU::Supports(const CompressedImage &image) {
    return image.GetCompressionAlgorithm() == CompressionAlgorithm::BSHUF_LZ4
           && elem_size_of(image.GetMode()) != 0;
}

BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_ptr<CudaStream> in_stream)
    : stream(std::move(in_stream)),
      max_uncompressed_bytes(in_max_uncompressed_bytes) {
    gpu_status = CudaDevicePtr<uint32_t>(1);
    host_status = CudaHostPtr<uint32_t>(1);
    // The compressed buffer and the descriptors are grown to fit the first image instead of being
    // sized for a worst case that no real frame reaches. A chunk is a few MB against an image of
    // tens; sizing this from the UNCOMPRESSED size cost ~73 MB per worker to hold ~4 MB.
}

// gpu_shuffled holds a whole uncompressed frame - 72 MB at 18 Mpx, per worker - and only the
// DecodeShuffled() route ever writes it. That route is taken when a bitshuffle block is too large for
// the fused kernel, which neither writer this pipeline reads produces, so on a real frame the buffer
// is never touched. So allocate it the first time it is actually asked for, at the size the caller
// declared, rather than in the constructor.
void BSLZ4DecoderGPU::EnsureUncompressedCapacity(size_t bytes) {
    if (gpu_shuffled.get() && bytes <= max_uncompressed_bytes)
        return;
    const size_t want = std::max(bytes, max_uncompressed_bytes);
    cuda_err(cudaStreamSynchronize(*stream));
    gpu_shuffled = CudaDevicePtr<uint8_t>(want);
    max_uncompressed_bytes = want;
}

void BSLZ4DecoderGPU::EnsureCompressedCapacity(size_t bytes) {
    if (bytes <= compressed_capacity)
        return;
    // A little slack, so a frame that compresses slightly worse than the last does not reallocate.
    const size_t want = bytes + bytes / 4;
    cuda_err(cudaStreamSynchronize(*stream));   // nothing may still be reading the old buffer
    gpu_compressed = CudaDevicePtr<uint8_t>(want);
    compressed_capacity = want;
}

void BSLZ4DecoderGPU::EnsureBlockCapacity(size_t nblocks) {
    if (nblocks <= max_blocks)
        return;
    const size_t want = nblocks + nblocks / 4 + 16;
    cuda_err(cudaStreamSynchronize(*stream));   // the previous descriptor upload must have landed
    gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(want);
    host_desc = CudaHostPtr<BSLZ4BlockDesc>(want);
    max_blocks = want;
}

uint32_t BSLZ4DecoderGPU::BlockBytes(const CompressedImage &image) {
    if (image.GetCompressedSize() < 12)
        return 0;   // not a chunk at all; the caller's decode reports it properly
    return be32(image.GetCompressed() + 8);
}

// Everything up to and including getting the chunk onto the device: the container scan, which is the
// only part that has to happen on the host, the capacity checks, and the uploads. What decodes the
// blocks is left to the caller of this - either the LZ4 kernel below or the fused kernel in the
// preprocessor - because that is the whole difference between the two routes.
BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) {
    const uint8_t *src = image.GetCompressed();
    const size_t clen = image.GetCompressedSize();
    const size_t elem_size = elem_size_of(image.GetMode());
    const size_t total_bytes = image.GetUncompressedSize();

    if (clen < 12)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk shorter than its header");
    if (be64(src) != total_bytes)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 header size does not match the image");

    const uint32_t block_bytes = be32(src + 8);
    if (block_bytes == 0 || block_bytes % elem_size != 0)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block size invalid");

    const size_t block_elems = block_bytes / elem_size;
    // bitshuffle transposes 8 elements at a time and refuses a block that is not a multiple of 8;
    // the host decoder rejects this too (JFJochDecompress.h). Without the check the un-transpose
    // would silently drop the last size % 8 elements of every block.
    if (block_elems % BSHUF_BLOCKED_MULT != 0)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block size is not a multiple of 8 elements");

    const size_t nelements = total_bytes / elem_size;
    const size_t nfull = nelements / block_elems;
    const size_t rem = nelements - nfull * block_elems;
    const size_t last = rem - rem % BSHUF_BLOCKED_MULT;
    const size_t leftover_bytes = (rem % BSHUF_BLOCKED_MULT) * elem_size;

    // Walk the container to locate the blocks. Lengths are only knowable in order, so this scan is
    // inherent to the format rather than an implementation choice.
    // An image of fewer than 8 elements has no bitshuffle block at all - it is entirely the verbatim
    // tail. The host decoder handles that, so handle it here rather than declining: nblocks is simply
    // zero and only the tail is copied.
    const size_t nblocks_needed = nfull + (last > 0 ? 1 : 0);
    // Bound the block count by what the chunk could actually hold BEFORE it becomes an allocation:
    // a header declaring a one-element block size would otherwise ask for hundreds of MB of pinned
    // memory, and only then fail on the first block header.
    if (nblocks_needed > (clen - 12) / MIN_BLOCK_BYTES_ON_WIRE)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk too short for the blocks its header implies");

    EnsureCompressedCapacity(clen);
    EnsureBlockCapacity(nblocks_needed);

    size_t nblk = 0, off = 12, out_off = 0;
    for (size_t i = 0; i < nblocks_needed; i++) {
        if (off + 4 > clen)
            throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 block header");
        const uint32_t block_clen = be32(src + off);
        off += 4;
        if (block_clen == 0 || off + block_clen > clen)
            throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block extends past the chunk");
        const uint32_t ne = (i < nfull) ? static_cast<uint32_t>(block_elems) : static_cast<uint32_t>(last);
        host_desc.get()[nblk++] = {static_cast<uint32_t>(off), block_clen, static_cast<uint32_t>(out_off), ne};
        off += block_clen;
        out_off += static_cast<size_t>(ne) * elem_size;
    }
    // The tail that bitshuffle leaves uncompressed and copies verbatim, and then nothing else: the
    // host decoder requires the chunk to be consumed exactly, so require it here too rather than
    // ignoring trailing bytes that indicate the container is not what it claims to be.
    if (off + leftover_bytes > clen)
        throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 leftover bytes");
    if (off + leftover_bytes != clen)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk has trailing bytes after the last block");

    // Everything that can be checked on the host has been checked; from here work is queued.
    cuda_err(cudaEventRecord(decode_start, *stream));
    host_status.get()[0] = 0;
    cuda_err(cudaMemcpyAsync(gpu_status.get(), host_status.get(), sizeof(uint32_t),
                             cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaMemcpyAsync(gpu_compressed.get(), src, clen, cudaMemcpyHostToDevice, *stream));

    const int nb = static_cast<int>(nblk);
    if (nb > 0)
        cuda_err(cudaMemcpyAsync(gpu_desc.get(), host_desc.get(), nblk * sizeof(BSLZ4BlockDesc),
                                 cudaMemcpyHostToDevice, *stream));

    BSLZ4ShuffledImage ret;
    ret.compressed = gpu_compressed.get();
    ret.desc = gpu_desc.get();
    ret.status = gpu_status.get();
    ret.nblocks = nb;
    ret.elem_size = static_cast<uint32_t>(elem_size);
    ret.block_bytes = block_bytes;
    ret.tail_elems = static_cast<uint32_t>(leftover_bytes / elem_size);
    ret.tail_elem0 = static_cast<uint32_t>(out_off / elem_size);
    ret.tail_src = leftover_bytes > 0 ? gpu_compressed.get() + off : nullptr;
    return ret;
}

BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(const CompressedImage &image) {
    EnsureUncompressedCapacity(image.GetUncompressedSize());
    BSLZ4ShuffledImage ret = PrepareChunk(image);
    if (ret.nblocks > 0) {
        lz4_decode_blocks<<<(ret.nblocks * 32 + 255) / 256, 256, 0, *stream>>>(
                ret.compressed, ret.desc, gpu_shuffled.get(), ret.status,
                ret.nblocks, ret.elem_size);
        cuda_err(cudaGetLastError());
    }
    ret.shuffled = gpu_shuffled.get();
    QueueDecodeStatus();
    // Stop the clock here rather than after the un-transpose: getting the chunk onto the device and
    // LZ4-decoding it is the part that replaced the host decompression.
    cuda_err(cudaEventRecord(decode_stop, *stream));
    decode_timed = true;
    return ret;
}

BSLZ4ShuffledImage BSLZ4DecoderGPU::UploadCompressed(const CompressedImage &image) {
    BSLZ4ShuffledImage ret = PrepareChunk(image);
    // All this route decodes on its own is the PCIe upload, so that is what the clock brackets. The
    // LZ4 pass happens inside the caller's kernel, inseparably from the un-transpose and the
    // preprocessing, and is reported with them.
    cuda_err(cudaEventRecord(decode_stop, *stream));
    decode_timed = true;
    return ret;
}

void BSLZ4DecoderGPU::QueueDecodeStatus() {
    cuda_err(cudaMemcpyAsync(host_status.get(), gpu_status.get(), sizeof(uint32_t),
                             cudaMemcpyDeviceToHost, *stream));
}

void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) {
    const BSLZ4ShuffledImage s = DecodeShuffled(image);

    if (s.nblocks > 0) {   // an image of fewer than 8 elements is all tail and has no block
        switch (s.elem_size) {
            case 1: bitshuffle_untranspose<1><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
            case 2: bitshuffle_untranspose<2><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
            default: bitshuffle_untranspose<4><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
        }
        cuda_err(cudaGetLastError());
    }

    // The verbatim tail is already on the device inside the uploaded chunk.
    if (s.tail_elems > 0)
        cuda_err(cudaMemcpyAsync(gpu_out + static_cast<size_t>(s.tail_elem0) * s.elem_size, s.tail_src,
                                 static_cast<size_t>(s.tail_elems) * s.elem_size,
                                 cudaMemcpyDeviceToDevice, *stream));
}

void BSLZ4DecoderGPU::ThrowIfDecodeFailed() {
    if (host_status.get()[0] != 0)
        throw JFJochException(JFJochExceptionCategory::Compression,
                              "bslz4 block did not decode to its declared length - the compressed data is corrupt");
}

float BSLZ4DecoderGPU::GetDecodeTime_s() const {
    if (!decode_timed)
        return 0.0f;
    float ms = 0.0f;
    if (cudaEventElapsedTime(&ms, decode_start, decode_stop) != cudaSuccess)
        return 0.0f;
    return ms * 1e-3f;
}