#pragma once
#include <iostream>
#include <cstdio>  // for fopen, fclose
#include <vector>
#include <jxl/codestream_header.h>
#include <jxl/color_encoding.h>
#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>
class JPEGXL
{
public:
    JPEGXL(int width, int height, const unsigned char* img, int num_channels, float q, int e, int th = 4) :m_width(width), m_height(height), m_num_channels(num_channels), m_quality(q), m_efford(e)
    {
        auto enc = JxlEncoderMake(/*memory_manager=*/nullptr);

        auto runner = JxlThreadParallelRunnerMake(nullptr, th);
        if (JXL_ENC_SUCCESS != JxlEncoderSetParallelRunner(enc.get(),
            JxlThreadParallelRunner,
            runner.get())) {
            fprintf(stderr, "JxlEncoderSetParallelRunner failed\n");
        }
        
        JxlPixelFormat pixel_format = { m_num_channels, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0 };
        JxlBasicInfo basic_info;
        JxlEncoderInitBasicInfo(&basic_info);
        basic_info.xsize = m_width;
        basic_info.ysize = m_height;
        basic_info.alpha_bits = 0;
        basic_info.bits_per_sample = 8;
        basic_info.num_color_channels = m_num_channels;
        basic_info.uses_original_profile = false;

        if (JXL_ENC_SUCCESS != JxlEncoderSetBasicInfo(enc.get(), &basic_info)) {
            fprintf(stderr, "JxlEncoderSetBasicInfo failed\n");
            return;
        }
        JxlColorEncoding color_encoding = {};
        JxlColorEncodingSetToSRGB(&color_encoding,/*is_gray=*/pixel_format.num_channels < 3);
        if (JXL_ENC_SUCCESS != JxlEncoderSetColorEncoding(enc.get(), &color_encoding)) {
            fprintf(stderr, "JxlEncoderSetColorEncoding failed\n");
            return;
        }
        JxlEncoderFrameSettings* frame_settings = JxlEncoderFrameSettingsCreate(enc.get(), nullptr);
        JxlEncoderFrameSettingsSetOption(frame_settings, JXL_ENC_FRAME_SETTING_EFFORT, e);
        JxlEncoderFrameSettingsSetOption(frame_settings, JXL_ENC_FRAME_SETTING_DECODING_SPEED, 0);

        if (q == 0) {//if lossless
            JxlEncoderSetFrameDistance(frame_settings, 0);
            JxlEncoderSetFrameLossless(frame_settings, true);
        }
        else {
            JxlEncoderSetFrameLossless(frame_settings, false);
            JxlEncoderSetFrameDistance(frame_settings, q);
        }
        if (JXL_ENC_SUCCESS != JxlEncoderAddImageFrame(frame_settings, &pixel_format, (void*)img, sizeof(uint8_t) * m_width * m_height * m_num_channels)) {
            fprintf(stderr, "JxlEncoderAddImageFrame failed\n");
            return;
        }
        JxlEncoderCloseInput(enc.get());
        std::vector<uint8_t>* compressed = &compressed_data;
        compressed->resize(64);
        uint8_t* next_out = compressed->data();
        size_t avail_out = compressed->size() - (next_out - compressed->data());
        JxlEncoderStatus process_result = JXL_ENC_NEED_MORE_OUTPUT;
        while (process_result == JXL_ENC_NEED_MORE_OUTPUT) {
            process_result = JxlEncoderProcessOutput(enc.get(), &next_out, &avail_out);
            if (process_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;
            }
        }
        compressed->resize(next_out - compressed->data());
        if (JXL_ENC_SUCCESS != process_result) {
            fprintf(stderr, "JxlEncoderProcessOutput failed\n");
            return;
        }
    }
    ~JPEGXL() {}

    bool WriteFile(const char* filename) {
        FILE* file = fopen(filename, "wb");
        if (!file) {
            fprintf(stderr, "Could not open %s for writing\n", filename);
            return false;
        }
        if (fwrite(compressed_data.data(), sizeof(uint8_t), compressed_data.size(), file) !=
            compressed_data.size()) {
            fprintf(stderr, "Could not write bytes to %s\n", filename);
            fclose(file);
            return false;
        }
        if (fclose(file) != 0) {
            fprintf(stderr, "Could not close %s\n", filename);
            return false;
        }
        return true;
    }
private:
    int m_width;
    int m_height;
    int m_num_channels;
    float m_quality;
    int m_efford;
    std::vector<uint8_t> compressed_data;
};