#pragma once
// Decoding JPEG XL files in tests (libjxl), shared by the image tests.

#include <jxl/decode.h>
#include <jxl/decode_cxx.h>

#include <cstdint>
#include <fstream>
#include <iterator>
#include <string>
#include <vector>

namespace fgc_test {

struct Decoded {
    std::vector<uint16_t> px;   // values at the file's own bit depth
    int                   bits = 0;
    int                   channels = 0;
    bool                  original_profile = false;
};

// Decode to 16-bit words at the file's own bit depth (12-bit files give
// 0..4095, 8-bit files 0..255).
inline Decoded decode(const std::string& path, int channels) {
    std::ifstream        f(path, std::ios::binary);
    std::vector<uint8_t> data{std::istreambuf_iterator<char>(f), {}};
    Decoded              out;
    auto dec = JxlDecoderMake(nullptr);
    JxlDecoderSubscribeEvents(dec.get(), JXL_DEC_BASIC_INFO | JXL_DEC_FULL_IMAGE);
    JxlDecoderSetInput(dec.get(), data.data(), data.size());
    JxlDecoderCloseInput(dec.get());
    JxlPixelFormat pf = {static_cast<uint32_t>(channels), JXL_TYPE_UINT16, JXL_NATIVE_ENDIAN, 0};
    for (;;) {
        auto s = JxlDecoderProcessInput(dec.get());
        if (s == JXL_DEC_BASIC_INFO) {
            JxlBasicInfo bi;
            JxlDecoderGetBasicInfo(dec.get(), &bi);
            out.bits             = static_cast<int>(bi.bits_per_sample);
            out.channels         = static_cast<int>(bi.num_color_channels);
            out.original_profile = bi.uses_original_profile;
        } else if (s == JXL_DEC_NEED_IMAGE_OUT_BUFFER) {
            size_t n = 0;
            JxlDecoderImageOutBufferSize(dec.get(), &pf, &n);
            out.px.resize(n / 2);
            JxlDecoderSetImageOutBuffer(dec.get(), &pf, out.px.data(), n);
            JxlBitDepth bd = {JXL_BIT_DEPTH_FROM_CODESTREAM, 0, 0};
            JxlDecoderSetImageOutBitDepth(dec.get(), &bd);
        } else if (s == JXL_DEC_FULL_IMAGE || s == JXL_DEC_SUCCESS) {
            return out;
        } else {
            out.px.clear();
            return out;
        }
    }
}

}  // namespace fgc_test