// 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 "../../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.
    //
    // Because the lanes cooperate on the copies, a match can source bytes that OTHER lanes wrote in
    // an earlier sequence. Since Volta that needs an explicit __syncwarp() - implicit reconvergence
    // is not part of the programming model - so there is one after every copy loop. The full mask is
    // correct: the early return and every break test warp-uniform values, so lanes never diverge
    // permanently.
    //
    // Bounds: every read is clamped against iend and every write against oend, so a malformed or
    // corrupt payload cannot walk off either buffer. It can still stop early, which leaves the block
    // short; lane 0 flags that at the end and the host turns it into an exception.
    __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 *ip = src + desc[b].in_off;
        const uint8_t *const iend = ip + desc[b].in_len;
        uint8_t *const obase = dst + desc[b].out_off;
        uint8_t *op = obase;
        uint8_t *const oend = obase + desc[b].nelem * elem_size;
        bool malformed = false;

        while (ip < iend) {
            const uint32_t token = *ip++;
            uint32_t litlen = token >> 4;
            if (litlen == 15) {
                // read_variable_length(&ip, iend - RUN_MASK, initial_check=1) in the reference: the
                // chain may not start within, nor run into, the last RUN_MASK (15) input bytes. A
                // valid stream never does - the literals it counts have to follow it - so a chain
                // that reaches there is corruption, and this is the only place it shows up.
                if ((size_t)(iend - ip) <= 15) { malformed = true; break; }
                uint32_t s;
                do {
                    s = *ip++;
                    litlen += s;
                    if ((size_t)(iend - ip) < 15) { malformed = true; break; }
                } while (s == 255);
                if (malformed) break;
            }
            if (litlen) {
                // Clamped by the INPUT as well as the output: a corrupt litlen must not read past the
                // end of this block's payload or write past the end of the block. Clamping keeps the
                // kernel in bounds; needing to clamp at all means the stream is not decodable, which
                // is what the reference reports as an error, so record it.
                if (litlen > (uint32_t)(oend - op) || litlen > (uint32_t)(iend - ip))
                    malformed = true;
                const uint32_t n = min(min(litlen, (uint32_t)(oend - op)), (uint32_t)(iend - ip));
                for (uint32_t i = lane; i < n; i += 32) op[i] = ip[i];
                __syncwarp();
                op += n; ip += litlen;
            }

            // LZ4's parsing restrictions: an encoder may not leave a match within MFLIMIT (12) bytes
            // of the end of the block, nor fewer than 2+1+LASTLITERALS (8) input bytes after a
            // literal run that is not the last one. So once either limit is reached this can ONLY be
            // the final sequence, and the final sequence must consume the payload exactly. The
            // reference applies this whether or not the run was empty, which is why the test sits
            // outside the copy - a zero-length literal run near the end is just as illegal.
            if ((size_t)(oend - op) < 12 || (size_t)(iend - ip) < 8) {
                malformed = (ip != iend) || (op != oend);
                break;   // necessarily EOF
            }
            if (iend - ip < 2) 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) {
                // read_variable_length(&ip, iend - LASTLITERALS + 1, initial_check=0): bounded by the
                // last 4 input bytes rather than 15, and with no check before the first read.
                uint32_t s;
                do {
                    s = *ip++;
                    matchlen += s;
                    if ((size_t)(iend - ip) < 4) { malformed = true; break; }
                } while (s == 255);
                if (malformed) break;
            }
            matchlen += 4;   // minmatch

            if (offset == 0 || offset > (uint32_t)(op - obase)) { malformed = true; break; }
            const uint8_t *mp = op - offset;
            // A match may reach the end of the block but never past it - the reference treats an
            // overrun as an error rather than truncating, and so must this.
            if (matchlen > (uint32_t)(oend - op))
                malformed = true;
            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, and a runtime
                // modulo is an emulated division, so the two cheap cases are peeled off first.
                if (offset == 1) {
                    const uint8_t v = mp[0];
                    for (uint32_t i = lane; i < n; i += 32) op[i] = v;
                } else if ((offset & (offset - 1)) == 0) {
                    const uint32_t m = offset - 1;
                    for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i & m];
                } else {
                    for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i % offset];
                }
            }
            __syncwarp();
            op += n;
        }

        // A block must decode to exactly its declared length AND consume exactly its payload. Both
        // are conditions LZ4_decompress_safe reports to the host path, and both are needed: a corrupt
        // stream can land on the right output length while leaving input over, or run its input out
        // early. Either way the bytes are not the ones that were compressed.
        if (lane == 0 && (malformed || op != oend || ip != iend))
            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_shuffled = CudaDevicePtr<uint8_t>(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.
}

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

BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(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 (total_bytes > max_uncompressed_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;
    // 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));
        lz4_decode_blocks<<<(nb * 32 + 255) / 256, 256, 0, *stream>>>(
                gpu_compressed.get(), gpu_desc.get(), gpu_shuffled.get(), gpu_status.get(),
                nb, static_cast<uint32_t>(elem_size));
        cuda_err(cudaGetLastError());
    }
    cuda_err(cudaMemcpyAsync(host_status.get(), gpu_status.get(), sizeof(uint32_t),
                             cudaMemcpyDeviceToHost, *stream));
    // 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, and it is the same work on
    // both the raw-bytes path and the fused one, where the un-transpose is inseparable from
    // preprocessing and is reported with it.
    cuda_err(cudaEventRecord(decode_stop, *stream));
    decode_timed = true;

    BSLZ4ShuffledImage ret;
    ret.shuffled = gpu_shuffled.get();
    ret.desc = gpu_desc.get();
    ret.nblocks = nb;
    ret.elem_size = static_cast<uint32_t>(elem_size);
    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;
}

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