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

#include <algorithm>

#include "BSLZ4DecoderGPU.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 WARP per LZ4 block. Every lane runs the same sequence parser over the same bytes - a
    // broadcast read, so no divergence - and the literal and match copies are split across the 32
    // lanes so the stores coalesce. One thread per block instead has each thread streaming its own
    // 8 kB region, which coalesces not at all and measured 13x slower.
    __global__ void lz4_decode_blocks(const uint8_t *__restrict__ src,
                                      const BSLZ4BlockDesc *__restrict__ desc,
                                      uint8_t *__restrict__ dst,
                                      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 *ip = src + desc[b].in_off;
        const uint8_t *const iend = ip + desc[b].in_len;
        uint8_t *op = dst + desc[b].out_off;
        uint8_t *const oend = op + desc[b].nelem * elem_size;

        while (ip < iend) {
            const uint32_t token = *ip++;
            uint32_t litlen = token >> 4;
            if (litlen == 15) {
                uint32_t s;
                do { s = *ip++; litlen += s; } while (s == 255 && ip < iend);
            }
            if (litlen) {
                const uint32_t n = min(litlen, (uint32_t)(oend - op));
                for (uint32_t i = lane; i < n; i += 32) op[i] = ip[i];
                op += n; ip += litlen;
            }
            if (ip >= iend) break;   // last sequence carries literals only

            const uint32_t offset = (uint32_t)ip[0] | ((uint32_t)ip[1] << 8);
            ip += 2;
            uint32_t matchlen = token & 0x0F;
            if (matchlen == 15) {
                uint32_t s;
                do { s = *ip++; matchlen += s; } while (s == 255 && ip < iend);
            }
            matchlen += 4;   // minmatch

            if (offset == 0 || offset > (uint32_t)(op - (dst + desc[b].out_off))) break;  // malformed
            const uint8_t *mp = op - offset;
            const uint32_t n = min(matchlen, (uint32_t)(oend - op));
            if (offset >= matchlen) {
                for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i];
            } else {
                // An overlapping match is a pattern of period `offset`. mp[0..offset-1] all lie
                // before op and are already final, so each output byte can be sourced from them
                // independently - which keeps this parallel rather than a serial byte loop. Long
                // zero runs in sparse detector data arrive here with offset == 1.
                for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i % offset];
            }
            op += n;
        }
    }

    __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, one CUDA block per bitshuffle block, mirroring bitshuf_decode_block:
    // un-transpose the bits of each byte-plane, then interleave the planes back into elements. The
    // planes are staged in shared memory so the final store is coalesced.
    extern __shared__ uint8_t smem[];
    __global__ void bitshuffle_untranspose(const uint8_t *__restrict__ shuffled,
                                           const BSLZ4BlockDesc *__restrict__ desc,
                                           uint8_t *__restrict__ out,
                                           int nblocks, uint32_t elem_size) {
        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;

        for (uint32_t p = 0; p < elem_size; p++) {
            const uint8_t *pin = in + p * size;
            uint8_t *pout = smem + p * size;
            for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
                uint64_t a = 0;
                #pragma unroll
                for (int k = 0; k < 8; k++) a |= (uint64_t)pin[k * n + i] << (8 * k);
                const uint64_t x = transpose8(a);
                #pragma unroll
                for (int k = 0; k < 8; k++) pout[i * 8 + k] = (uint8_t)(x >> (8 * k));
            }
        }
        __syncthreads();

        for (uint32_t i = threadIdx.x; i < size; i += blockDim.x)
            for (uint32_t j = 0; j < elem_size; j++)
                dst[i * elem_size + j] = smem[j * size + i];
    }

    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
        }
    }
}

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) {
    // The compressed chunk is smaller than the image in every case worth having, but a pathological
    // frame can expand slightly, so allow headroom rather than risk a per-frame reallocation.
    max_compressed_bytes = max_uncompressed_bytes + max_uncompressed_bytes / 16 + 4096;
    // Enough descriptors for the block size the compressor actually uses (8 kB); Decode() grows
    // these if a stream turns up with smaller blocks. Sizing for the format's theoretical minimum
    // block instead would allocate megabytes of pinned memory that no real file needs.
    max_blocks = max_uncompressed_bytes / 8192 + 2;

    gpu_compressed = CudaDevicePtr<uint8_t>(max_compressed_bytes);
    gpu_shuffled = CudaDevicePtr<uint8_t>(max_uncompressed_bytes);
    gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(max_blocks);
    host_desc = CudaHostPtr<BSLZ4BlockDesc>(max_blocks);
}

void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) {
    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 (total_bytes > max_uncompressed_bytes || clen > max_compressed_bytes)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 image larger than the decoder was sized for");
    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;
    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.
    const size_t nblocks_needed = nfull + (last > 0 ? 1 : 0);
    if (nblocks_needed > max_blocks) {   // a stream with smaller blocks than the compressor emits
        max_blocks = nblocks_needed;
        gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(max_blocks);
        host_desc = CudaHostPtr<BSLZ4BlockDesc>(max_blocks);
    }

    size_t nblk = 0, off = 12, out_off = 0;
    for (size_t i = 0; i < nfull + (last > 0 ? 1 : 0); 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");
        if (nblk >= max_blocks)
            throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk has more blocks than expected");
        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;
    }
    if (nblk == 0)
        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk contains no blocks");

    cuda_err(cudaMemcpyAsync(gpu_compressed.get(), src, clen, cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaMemcpyAsync(gpu_desc.get(), host_desc.get(), nblk * sizeof(BSLZ4BlockDesc),
                             cudaMemcpyHostToDevice, *stream));

    const int nb = static_cast<int>(nblk);
    lz4_decode_blocks<<<(nb * 32 + 255) / 256, 256, 0, *stream>>>(
            gpu_compressed.get(), gpu_desc.get(),
            gpu_shuffled.get(), nb, static_cast<uint32_t>(elem_size));
    bitshuffle_untranspose<<<nb, 256, block_bytes, *stream>>>(
            gpu_shuffled.get(), gpu_desc.get(),
            gpu_out, nb, static_cast<uint32_t>(elem_size));
    cuda_err(cudaGetLastError());

    // The tail that bitshuffle leaves uncompressed and copies verbatim.
    if (leftover_bytes > 0) {
        if (off + leftover_bytes > clen)
            throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 leftover bytes");
        cuda_err(cudaMemcpyAsync(gpu_out + out_off, src + off, leftover_bytes,
                                 cudaMemcpyHostToDevice, *stream));
    }
}