#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 void* img, int width, int height,
                                 int channels, int bits, float distance, int effort,
                                 int threads) {
    if (!img || width <= 0 || height <= 0 || (channels != 1 && channels != 3) || bits < 8 ||
        bits > 16 || distance < 0.0f) {
        LOG_ERROR << "JpegXlEncoder: bad arguments (" << width << 'x' << height << ", "
                  << channels << " channels, " << bits << " bits, distance " << distance << ')';
        return false;
    }
    const bool lossless = distance == 0.0f;
    const bool wide     = bits > 8;

    auto enc = JxlEncoderMake(nullptr);
    JxlThreadParallelRunnerPtr runner;
    if (threads > 0) {
        runner = JxlThreadParallelRunnerMake(nullptr, static_cast<size_t>(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),
                                   wide ? JXL_TYPE_UINT16 : 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    = static_cast<uint32_t>(bits);
    basic_info.num_color_channels = channels;
    // Lossless requires the original colour space; with it off, libjxl encodes
    // through its lossy XYB space, reports success, and the pixels come back
    // changed (known-issues #20). Lossy keeps XYB, as the program always has.
    basic_info.uses_original_profile = lossless ? JXL_TRUE : 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 (lossless) {
        JxlEncoderSetFrameDistance(fs, 0);
        JxlEncoderSetFrameLossless(fs, JXL_TRUE);
    } else {
        JxlEncoderSetFrameLossless(fs, JXL_FALSE);
        JxlEncoderSetFrameDistance(fs, distance);
    }
    if (wide) {
        // The 16-bit words hold `bits`-bit values: read them at the header's
        // depth, not as the full 0..65535 range.
        JxlBitDepth depth = {JXL_BIT_DEPTH_FROM_CODESTREAM, 0, 0};
        if (JXL_ENC_SUCCESS != JxlEncoderSetFrameBitDepth(fs, &depth)) {
            LOG_ERROR << "JxlEncoderSetFrameBitDepth failed";
            return false;
        }
    }

    const size_t byte_size =
        static_cast<size_t>(width) * height * channels * (wide ? 2 : 1);
    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;
        std::remove(path.c_str());
    }
    return ok;
}

}  // namespace fgc