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

#include <algorithm>
#include <cstring>
#include <string>

#include <cub/cub.cuh>

#include "CBFDecoderGPU.h"
#include "../../common/JFJochException.h"

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

    // The stream is walked in parts of PART_BYTES, one thread each, so that a frame of a few MB is many
    // short walks. The scan that strings them together runs over tiles of PARTS_PER_TILE parts: each
    // of its elements is a whole 16-entry map, and on parts that scan cost more than the walks.
    constexpr size_t PART_BYTES = 68;
    constexpr size_t PARTS_PER_TILE = 8;
    constexpr size_t TILE_BYTES = PART_BYTES * PARTS_PER_TILE;
    // Threads per block of the part kernels. A block reads its parts' bytes into shared memory together,
    // coalesced, and walks them there; the write kernel stages its pixels the same way - at most one
    // pixel per byte, so PART_BYTES per thread - and writes them out together.
    constexpr int BLOCK_THREADS = 128;
    // A token starting in a block's last part can run 14 bytes past it.
    constexpr size_t BLOCK_BYTES = PART_BYTES * BLOCK_THREADS + 16;

    // The bytes of the stream from `first` on, as a block holds them in shared memory.
    struct StagedBytes {
        const uint8_t *bytes;
        size_t first;
        size_t size;   // of the whole stream
        __device__ uint32_t operator[](size_t p) const { return bytes[p - first]; }
    };

    __device__ inline StagedBytes StageBlock(const uint8_t *__restrict__ d, size_t size, uint8_t *staged) {
        const size_t first = blockIdx.x * PART_BYTES * BLOCK_THREADS;
        for (size_t i = threadIdx.x; i < BLOCK_BYTES; i += BLOCK_THREADS)
            staged[i] = first + i < size ? d[first + i] : 0;
        __syncthreads();
        return {staged, first, size};
    }

    // The token starting at byte p: its length, and its delta (low 32 bits) in `delta`. 0 when the
    // token runs past the end of the stream - where the host decoder stops too.
    // Following Bernstein & Hammersley (2006) Int. Tables Cryst. G, 37-43
    __device__ inline int TokenAt(const StagedBytes &d, size_t p, uint32_t &delta) {
        const size_t size = d.size;
        const int8_t b = static_cast<int8_t>(d[p]);
        if (b != -128) {
            delta = static_cast<uint32_t>(static_cast<int32_t>(b));
            return 1;
        }
        if (p + 3 > size) return 0;
        const int16_t v16 = static_cast<int16_t>(d[p + 1] | (d[p + 2] << 8));
        if (v16 != -32768) {
            delta = static_cast<uint32_t>(static_cast<int32_t>(v16));
            return 3;
        }
        if (p + 7 > size) return 0;
        const uint32_t v32 = d[p + 3] | (d[p + 4] << 8) | (d[p + 5] << 16) | (static_cast<uint32_t>(d[p + 6]) << 24);
        if (v32 != 0x80000000u) {
            delta = v32;
            return 7;
        }
        if (p + 15 > size) return 0;
        // An int64 delta: only its low 32 bits reach an int32 pixel.
        delta = d[p + 7] | (d[p + 8] << 8) | (d[p + 9] << 16) | (static_cast<uint32_t>(d[p + 10]) << 24);
        return 15;
    }

    // One step of a walk at position pos: false when the stream has ended there.
    __device__ inline bool Step(const StagedBytes &d, size_t &pos, uint32_t &count, uint32_t &sum) {
        uint32_t delta = 0;
        const int len = TokenAt(d, pos, delta);
        if (len == 0) return false;
        pos += len;
        count++;
        sum += delta;
        return true;
    }

    __device__ inline uint8_t ExitOf(size_t pos, size_t tile_end, size_t size, bool ended) {
        return (ended || pos >= size) ? CBFTileMap::DEAD : static_cast<uint8_t>(pos - tile_end);
    }

    // The map of one part. A part wholly past the end of the stream maps every entry to DEAD.
    __global__ void part_map_kernel(const uint8_t *__restrict__ stream, size_t size, size_t nparts,
                                    CBFTileMap *__restrict__ maps) {
        __shared__ uint8_t staged[BLOCK_BYTES];
        const StagedBytes d = StageBlock(stream, size, staged);
        const size_t t = blockIdx.x * static_cast<size_t>(BLOCK_THREADS) + threadIdx.x;
        if (t >= nparts) return;
        const size_t start = t * PART_BYTES;
        const size_t end = start + PART_BYTES;
        const size_t limit = min(end, size);
        CBFTileMap m;

        // The walk from the part's first byte, in full.
        size_t pos0 = start;
        uint32_t count0 = 0, sum0 = 0;
        bool ended0 = false;
        while (pos0 < limit && !ended0)
            ended0 = !Step(d, pos0, count0, sum0);
        m.count[0] = count0;
        m.sum[0] = sum0;
        m.exit[0] = ExitOf(pos0, end, size, ended0);

        // The walks from the other entries. Each is stepped beside the first walk, whichever is behind
        // moving, until both stand on the same byte - from there they are one walk, and the rest is
        // the first walk's remainder.
        for (int e = 1; e < CBFTileMap::DEAD; e++) {
            size_t a = start, b = start + e;
            uint32_t count_a = 0, sum_a = 0, count_b = 0, sum_b = 0;
            bool a_ended = false, b_ended = false, merged = false;
            while (b < limit && !b_ended) {
                if (a == b && !a_ended) {
                    merged = true;
                    break;
                }
                if (a < b && !a_ended)
                    a_ended = !Step(d, a, count_a, sum_a);
                else
                    b_ended = !Step(d, b, count_b, sum_b);
            }
            if (merged) {
                m.count[e] = count_b + (count0 - count_a);
                m.sum[e] = sum_b + (sum0 - sum_a);
                m.exit[e] = m.exit[0];
            } else {
                m.count[e] = count_b;
                m.sum[e] = sum_b;
                m.exit[e] = ExitOf(b, end, size, b_ended);
            }
        }
        m.count[CBFTileMap::DEAD] = 0;
        m.sum[CBFTileMap::DEAD] = 0;
        m.exit[CBFTileMap::DEAD] = CBFTileMap::DEAD;
        maps[t] = m;
    }

    // a, then b. The scan's look-back also composes the uninitialised state of tiles before the first
    // one and discards the result, so an exit is masked to the table rather than trusted to be in it.
    struct ComposeTileMaps {
        __host__ __device__ CBFTileMap operator()(const CBFTileMap &a, const CBFTileMap &b) const {
            CBFTileMap r;
            for (int e = 0; e < 16; e++) {
                const int mid = a.exit[e] & 15;
                r.exit[e] = b.exit[mid];
                r.count[e] = a.count[e] + b.count[mid];
                r.sum[e] = a.sum[e] + b.sum[mid];
            }
            return r;
        }
    };

    __global__ void tile_map_kernel(const CBFTileMap *__restrict__ part_maps, size_t ntiles,
                                    CBFTileMap *__restrict__ tile_maps) {
        const size_t t = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
        if (t >= ntiles) return;
        CBFTileMap m = part_maps[t * PARTS_PER_TILE];
        for (size_t k = 1; k < PARTS_PER_TILE; k++)
            m = ComposeTileMaps()(m, part_maps[t * PARTS_PER_TILE + k]);
        tile_maps[t] = m;
    }

    __global__ void part_write_kernel(const uint8_t *__restrict__ stream, size_t size, size_t nparts,
                                      const CBFTileMap *__restrict__ tile_entries,
                                      const CBFTileMap *__restrict__ part_maps,
                                      int32_t *__restrict__ out, uint32_t npixels, uint32_t *__restrict__ total) {
        __shared__ uint8_t staged_bytes[BLOCK_BYTES];
        __shared__ int32_t staged[PART_BYTES * BLOCK_THREADS];
        __shared__ uint32_t block_first;
        __shared__ uint32_t block_end;
        const StagedBytes d = StageBlock(stream, size, staged_bytes);
        const size_t p = blockIdx.x * static_cast<size_t>(BLOCK_THREADS) + threadIdx.x;

        // The state this part is entered in: its tile's, carried through the parts before it in the tile.
        uint32_t state = CBFTileMap::DEAD, pixel = 0, value = 0;
        if (p < nparts) {
            const size_t tile = p / PARTS_PER_TILE;
            state = tile_entries[tile].exit[0];
            pixel = tile_entries[tile].count[0];
            value = tile_entries[tile].sum[0];
            for (size_t k = tile * PARTS_PER_TILE; k < p; k++) {
                pixel += part_maps[k].count[state];
                value += part_maps[k].sum[state];
                state = part_maps[k].exit[state];
            }
            if (p == nparts - 1)
                *total = pixel + part_maps[p].count[state];
        }
        // The block's pixels are contiguous, from the first one of its first part.
        if (threadIdx.x == 0) {
            block_first = pixel;
            block_end = pixel;
        }
        __syncthreads();

        if (p < nparts && state != CBFTileMap::DEAD) {
            const size_t end = min((p + 1) * PART_BYTES, size);
            size_t pos = p * PART_BYTES + state;
            while (pos < end && pixel < npixels) {
                uint32_t delta = 0;
                const int len = TokenAt(d, pos, delta);
                if (len == 0) break;
                pos += len;
                value += delta;
                staged[pixel++ - block_first] = static_cast<int32_t>(value);
            }
            atomicMax(&block_end, min(pixel, npixels));
        }
        __syncthreads();
        for (uint32_t i = threadIdx.x; block_first + i < block_end; i += BLOCK_THREADS)
            out[block_first + i] = staged[i];
    }
}

CBFDecoderGPU::CBFDecoderGPU(std::shared_ptr<CudaStream> in_stream)
    : stream(std::move(in_stream)), gpu_total(1), host_total(1) {}

void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) {
    const size_t size = image.GetCompressedSize();
    last_npixels = image.GetWidth() * image.GetHeight();
    if (size == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "miniCBF byte-offset stream ended after 0 of " + std::to_string(last_npixels) + " pixels");
    const size_t ntiles = (size + TILE_BYTES - 1) / TILE_BYTES;
    const size_t nparts = ntiles * PARTS_PER_TILE;

    if (size > compressed_capacity || ntiles > tile_capacity) {
        // A frame's stream is a little longer or shorter than the last one's: grow with slack.
        cuda_err(cudaStreamSynchronize(*stream));
        compressed_capacity = size + size / 4;
        tile_capacity = (compressed_capacity + TILE_BYTES - 1) / TILE_BYTES;
        host_compressed = CudaHostPtr<uint8_t>(compressed_capacity);
        gpu_compressed = CudaDevicePtr<uint8_t>(compressed_capacity);
        gpu_part_maps = CudaDevicePtr<CBFTileMap>(tile_capacity * PARTS_PER_TILE);
        gpu_maps = CudaDevicePtr<CBFTileMap>(tile_capacity);
        gpu_entries = CudaDevicePtr<CBFTileMap>(tile_capacity);
        CBFTileMap identity{};
        cuda_err(cub::DeviceScan::ExclusiveScan(nullptr, scan_temp_bytes, gpu_maps.get(), gpu_entries.get(),
                                                ComposeTileMaps(), identity, tile_capacity, *stream));
        gpu_scan_temp = CudaDevicePtr<uint8_t>(scan_temp_bytes);
    }

    // The upload waits for the previous frame's to have left the staging buffer.
    cuda_err(cudaStreamSynchronize(*stream));
    std::memcpy(host_compressed.get(), image.GetCompressed(), size);
    cuda_err(cudaEventRecord(decode_start, *stream));
    cuda_err(cudaMemcpyAsync(gpu_compressed, host_compressed.get(), size, cudaMemcpyHostToDevice, *stream));

    const int threads = 128;
    part_map_kernel<<<(nparts + BLOCK_THREADS - 1) / BLOCK_THREADS, BLOCK_THREADS, 0, *stream>>>(
            gpu_compressed, size, nparts, gpu_part_maps);
    cuda_err(cudaGetLastError());
    tile_map_kernel<<<(ntiles + threads - 1) / threads, threads, 0, *stream>>>(gpu_part_maps, ntiles, gpu_maps);
    cuda_err(cudaGetLastError());

    // Entering tile 0 at its first byte: every entry maps to itself, nothing emitted yet.
    CBFTileMap identity{};
    for (int e = 0; e < 16; e++)
        identity.exit[e] = static_cast<uint8_t>(e);
    size_t temp_bytes = scan_temp_bytes;
    cuda_err(cub::DeviceScan::ExclusiveScan(gpu_scan_temp.get(), temp_bytes, gpu_maps.get(), gpu_entries.get(),
                                            ComposeTileMaps(), identity, ntiles, *stream));

    part_write_kernel<<<(nparts + BLOCK_THREADS - 1) / BLOCK_THREADS, BLOCK_THREADS, 0, *stream>>>(
            gpu_compressed, size, nparts, gpu_entries, gpu_part_maps, gpu_out,
            static_cast<uint32_t>(last_npixels), gpu_total);
    cuda_err(cudaGetLastError());
    cuda_err(cudaEventRecord(decode_stop, *stream));
    cuda_err(cudaMemcpyAsync(host_total.get(), gpu_total.get(), sizeof(uint32_t), cudaMemcpyDeviceToHost, *stream));
}

void CBFDecoderGPU::ThrowIfDecodeFailed() const {
    const size_t total = host_total.get()[0];
    if (total < last_npixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "miniCBF byte-offset stream ended after " + std::to_string(total)
                              + " of " + std::to_string(last_npixels) + " pixels");
}

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