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

    constexpr size_t TILE_BYTES = 256;

    // 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 uint8_t *__restrict__ d, size_t size, size_t p, uint32_t &delta) {
        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 uint8_t *__restrict__ d, size_t size, size_t &pos,
                                uint32_t &count, uint32_t &sum) {
        uint32_t delta = 0;
        const int len = TokenAt(d, size, 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);
    }

    __global__ void tile_map_kernel(const uint8_t *__restrict__ d, size_t size, size_t ntiles,
                                    CBFTileMap *__restrict__ maps) {
        const size_t t = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
        if (t >= ntiles) return;
        const size_t start = t * TILE_BYTES;
        const size_t end = start + TILE_BYTES;
        const size_t limit = min(end, size);
        CBFTileMap m;

        // The walk from the tile'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, size, 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, size, a, count_a, sum_a);
                else
                    b_ended = !Step(d, size, 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_write_kernel(const uint8_t *__restrict__ d, size_t size, size_t ntiles,
                                      const CBFTileMap *__restrict__ entries, const CBFTileMap *__restrict__ maps,
                                      int32_t *__restrict__ out, uint32_t npixels, uint32_t *__restrict__ total) {
        const size_t t = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
        if (t >= ntiles) return;
        const CBFTileMap &in = entries[t];
        if (t == ntiles - 1)
            *total = in.count[0] + maps[t].count[in.exit[0]];
        if (in.exit[0] == CBFTileMap::DEAD)
            return;
        const size_t end = min((t + 1) * TILE_BYTES, size);
        size_t pos = t * TILE_BYTES + in.exit[0];
        uint32_t pixel = in.count[0];
        uint32_t value = in.sum[0];
        while (pos < end && pixel < npixels) {
            uint32_t delta = 0;
            const int len = TokenAt(d, size, pos, delta);
            if (len == 0) break;
            pos += len;
            value += delta;
            out[pixel++] = static_cast<int32_t>(value);
        }
    }
}

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;

    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_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;
    const int blocks = static_cast<int>((ntiles + threads - 1) / threads);
    tile_map_kernel<<<blocks, threads, 0, *stream>>>(gpu_compressed, size, 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));

    tile_write_kernel<<<blocks, threads, 0, *stream>>>(gpu_compressed, size, ntiles, gpu_entries, gpu_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;
}