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

#pragma once

#include <cstdint>
#include <cstddef>
#include <vector>

// Produces a STANDARD Zstandard frame from bitshuffled data, decodable by stock ZSTD_decompress:
//   - zero / 0xFF runs        -> RLE_Blocks (cheap, like JFJochZstdCompressor)
//   - literal regions         -> Compressed_Blocks with Huffman literals (no sequences); short
//                                runs are absorbed into the literal stream
//   - incompressible literals -> Raw_Blocks (bounded worst case)
// Faster than full ZSTD (no match search) and better ratio than the plain RLE compressor (it
// entropy-codes the literals). The Huffman table is built/reused per block from that block's own
// literals via zstd's HUF_compress*X_repeat, so it is robust on any input (random, Poisson, masks,
// zeros) with no trained tables.
class JFJochZstdHuffCompressor {
    std::vector<uint8_t>  out;        // assembled frame
    std::vector<uint8_t>  literals;   // literal bytes in stream order (incl. absorbed short runs)
    std::vector<uint8_t>  hufbuf;     // scratch for one Huffman-coded literal chunk
    std::vector<size_t>   ctable;     // HUF_CElt[] (size_t-aligned)
    std::vector<uint64_t> entwksp;    // HUF compression workspace
    unsigned repeat_state = 0;        // HUF_repeat across literal blocks within the current frame

    struct Seg { uint8_t type; size_t bytes; size_t lit_off; }; // type 0=run0, 1=runFF, 2=literals
    std::vector<Seg> segs;

    void   put_le(uint64_t v, int nbytes);
    size_t blk_hdr(uint32_t type, uint32_t size);
    void   emit_run(uint8_t value, size_t nbytes, size_t &last_off);
    void   emit_lit_chunk(const uint8_t *lits, size_t n, size_t &last_off);
public:
    JFJochZstdHuffCompressor();
    // src = bitshuffled block (src_size bytes, a multiple of 8). Writes one zstd frame to dst and
    // returns its size. dst must hold at least ZSTD_compressBound(src_size) + 12 bytes.
    size_t Compress(uint8_t *dst, const uint64_t *src, size_t src_size);
};