#include "fgc/JpegXlEncoder.h"
#include "fgc/Logger.h"

#include <cstdio>
#include <vector>

#include <jxl/encode.h>
#include <jxl/encode_cxx.h>
#include <jxl/thread_parallel_runner.h>
#include <jxl/thread_parallel_runner_cxx.h>
#include <jxl/types.h>

namespace fgc {

bool JpegXlEncoder::encodeToFile(const std::string& path, const uint8_t* img, int width,
                                 int height, int channels, float distance, int effort,
                                 int threads) {
    auto enc    = JxlEncoderMake(nullptr);
    auto runner = JxlThreadParallelRunnerMake(nullptr, threads);
    if (JXL_ENC_SUCCESS !=
        JxlEncoderSetParallelRunner(enc.get(), JxlThreadParallelRunner, runner.get())) {
        LOG_ERROR << "JxlEncoderSetParallelRunner failed";
        return false;
    }

    JxlPixelFormat pixel_format = {static_cast<uint32_t>(channels), JXL_TYPE_UINT8,
                                   JXL_NATIVE_ENDIAN, 0};

    JxlBasicInfo basic_info;
    JxlEncoderInitBasicInfo(&basic_info);
    basic_info.xsize              = width;
    basic_info.ysize              = height;
    basic_info.alpha_bits         = 0;
    basic_info.bits_per_sample    = 8;
    basic_info.num_color_channels = channels;
    basic_info.uses_original_profile = JXL_FALSE;
    if (JXL_ENC_SUCCESS != JxlEncoderSetBasicInfo(enc.get(), &basic_info)) {
        LOG_ERROR << "JxlEncoderSetBasicInfo failed";
        return false;
    }

    JxlColorEncoding color_encoding = {};
    JxlColorEncodingSetToSRGB(&color_encoding, /*is_gray=*/channels < 3);
    if (JXL_ENC_SUCCESS != JxlEncoderSetColorEncoding(enc.get(), &color_encoding)) {
        LOG_ERROR << "JxlEncoderSetColorEncoding failed";
        return false;
    }

    JxlEncoderFrameSettings* fs = JxlEncoderFrameSettingsCreate(enc.get(), nullptr);
    JxlEncoderFrameSettingsSetOption(fs, JXL_ENC_FRAME_SETTING_EFFORT, effort);
    JxlEncoderFrameSettingsSetOption(fs, JXL_ENC_FRAME_SETTING_DECODING_SPEED, 0);
    if (distance == 0.0f) {
        JxlEncoderSetFrameDistance(fs, 0);
        JxlEncoderSetFrameLossless(fs, JXL_TRUE);
    } else {
        JxlEncoderSetFrameLossless(fs, JXL_FALSE);
        JxlEncoderSetFrameDistance(fs, distance);
    }

    const size_t byte_size = static_cast<size_t>(width) * height * channels;
    if (JXL_ENC_SUCCESS !=
        JxlEncoderAddImageFrame(fs, &pixel_format, img, byte_size)) {
        LOG_ERROR << "JxlEncoderAddImageFrame failed";
        return false;
    }
    JxlEncoderCloseInput(enc.get());

    std::vector<uint8_t> compressed(64);
    uint8_t* next_out  = compressed.data();
    size_t   avail_out = compressed.size();
    JxlEncoderStatus result = JXL_ENC_NEED_MORE_OUTPUT;
    while (result == JXL_ENC_NEED_MORE_OUTPUT) {
        result = JxlEncoderProcessOutput(enc.get(), &next_out, &avail_out);
        if (result == JXL_ENC_NEED_MORE_OUTPUT) {
            size_t offset = next_out - compressed.data();
            compressed.resize(compressed.size() * 2);
            next_out  = compressed.data() + offset;
            avail_out = compressed.size() - offset;
        }
    }
    if (JXL_ENC_SUCCESS != result) {
        LOG_ERROR << "JxlEncoderProcessOutput failed";
        return false;
    }
    compressed.resize(next_out - compressed.data());

    FILE* file = std::fopen(path.c_str(), "wb");
    if (!file) {
        LOG_ERROR << "Could not open " << path << " for writing";
        return false;
    }
    bool ok = std::fwrite(compressed.data(), 1, compressed.size(), file) == compressed.size();
    if (std::fclose(file) != 0) ok = false;
    if (!ok) LOG_ERROR << "Could not write " << path;
    return ok;
}

}  // namespace fgc