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

#pragma once

#include <bitshuffle/bitshuffle.h>

#include <vector>
#include <cstdint>
#include <cstddef>
#include "CompressionAlgorithmEnum.h"
#include "MaxCompressedSize.h"

#include "JFJochZstdCompressor.h"
#include "JFJochZstdHuffCompressor.h"

class JFJochBitShuffleCompressor {
    JFJochZstdCompressor zstd_compressor;
    JFJochZstdHuffCompressor huff_compressor;
    CompressionAlgorithm algorithm;
    std::vector<char> tmp_space;
    std::vector<char> scratch;

    size_t CompressBlock(char *dest, const char * source, size_t nelements, size_t elem_size);
public:
    // The bitshuffle block size is chosen as a byte target rather than a fixed element count, so
    // the per-block working set - and thus the cache behaviour - stays constant across pixel bit
    // depths. The target is per-algorithm: LZ4 favours a small, cache-resident block (throughput),
    // ZSTD/RLE a large block (ratio). The element block size scales inversely with elem_size,
    // which is always a power of two here, so the result is a power of two >= 8 (a valid bitshuffle
    // block, multiple of 8). E.g. ZSTD 128 kB -> 65536 elem at 16-bit, 32768 at 32-bit.
    constexpr static size_t DefaultBlockSizeBytes(CompressionAlgorithm algorithm) {
        return algorithm == CompressionAlgorithm::BSHUF_LZ4 ? 16384 : 131072;
    }
    constexpr static size_t BlockSize(CompressionAlgorithm algorithm, size_t elem_size) {
        return DefaultBlockSizeBytes(algorithm) / elem_size;
    }

    explicit JFJochBitShuffleCompressor(CompressionAlgorithm algorithm);

    template<class T>
    size_t Compress(void *dest, size_t dest_size, const std::vector<T> &src) {
        return Compress(dest, dest_size, src.data(), src.size(), sizeof(T));
    };

    template<class T>
    std::vector<uint8_t> Compress(const std::vector<T> &src) {
        return Compress(src.data(), src.size(), sizeof(T));
    }
    std::vector<uint8_t> Compress(const void* source, size_t nelements, size_t elem_size);
    // Throws CompressionBufferTooSmallException if the compressed output would not fit dest_size.
    size_t Compress(void *dest, size_t dest_size, const void* source, size_t nelements, size_t elem_size);
};

template <class T> std::vector<T> bitshuffle(const std::vector<T> &input, size_t block_size) {
    std::vector<T> ret(input.size());
    bshuf_bitshuffle(input.data(), ret.data(), input.size(), sizeof(T), block_size);
    return ret;
}