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

#include <vector>
#include <csetjmp>
#include <cstdio>
#include <cstdint>
#include <cmath>
#include <new>
#include <jpeglib.h>
#include <jerror.h>

#include "../common/JFJochException.h"
#include "../common/CompressedImage.h"
#include "JFJochJPEG.h"

namespace {
    // libjpeg's own error handler prints its message and calls exit(), from whichever thread was
    // encoding - out of memory included. This one jumps back to WriteJPEGToMem, which throws instead.
    struct JPEGErrorManager {
        jpeg_error_mgr mgr;
        std::jmp_buf jump;
    };

    void JPEGErrorExit(j_common_ptr cinfo) {
        std::longjmp(reinterpret_cast<JPEGErrorManager *>(cinfo->err)->jump, 1);
    }
}


std::string WriteJPEGToMem(const CompressedImage& image, int quality) {
    unsigned char *buf = nullptr;
    unsigned long buf_size = 0;

    if (image.GetMode() != CompressedImageMode::RGB)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Only RGB images allowed for JPEG saving");

    std::vector<uint8_t> buffer;
    image.GetUncompressed(buffer);

    struct jpeg_compress_struct cinfo;
    JPEGErrorManager jerr;

    cinfo.err = jpeg_std_error(&jerr.mgr);
    jerr.mgr.error_exit = JPEGErrorExit;
    if (setjmp(jerr.jump)) {
        const bool out_of_memory = (jerr.mgr.msg_code == JERR_OUT_OF_MEMORY);
        char message[JMSG_LENGTH_MAX];
        jerr.mgr.format_message(reinterpret_cast<j_common_ptr>(&cinfo), message);
        jpeg_destroy_compress(&cinfo);
        free(buf);
        if (out_of_memory)
            throw std::bad_alloc();
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              std::string("JPEG encoding failed: ") + message);
    }
    jpeg_create_compress(&cinfo);
    jpeg_mem_dest(&cinfo, &buf, &buf_size);

    cinfo.image_width      = image.GetWidth();
    cinfo.image_height     = image.GetHeight();
    cinfo.input_components = 3;
    cinfo.in_color_space   = JCS_RGB;

    jpeg_set_defaults(&cinfo);
/*set the quality [0..100]  */
    jpeg_set_quality (&cinfo, quality, true);
    jpeg_simple_progression(&cinfo);
    jpeg_start_compress(&cinfo, true);

    JSAMPROW row_pointer;          /* pointer to a single row */

    while (cinfo.next_scanline < cinfo.image_height) {
        row_pointer = (JSAMPROW) (buffer.data() + cinfo.next_scanline*(sizeof(char)*3)*image.GetWidth());
        jpeg_write_scanlines(&cinfo, &row_pointer, 1);
    }

    jpeg_finish_compress(&cinfo);
    jpeg_destroy_compress(&cinfo);

    std::string ret;
    if (buf_size > 0)
        ret = std::string((char *)buf, buf_size);
    free(buf);

    return ret;
}