#include "ImageDecoder.h"
#include "BoostLog.h"
#include <jpeglib.h>
#include <stdio.h>

std::optional<ImageDecoder::Image> ImageDecoder::extractJpegYComponent(const std::string &filename) {
    struct jpeg_decompress_struct cinfo;
    struct jpeg_error_mgr jerr;
    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_decompress(&cinfo);

    FILE *infile = fopen(filename.c_str(), "rb");
    if (infile == NULL) {
        LOG(error) << "cannot open " << filename;
        return std::nullopt;
    }
    ImageDecoder::Image ret;
    jpeg_stdio_src(&cinfo, infile);
    jpeg_read_header(&cinfo, TRUE);
    cinfo.out_color_space = JCS_YCbCr; // We want YCbCr color space
    jpeg_start_decompress(&cinfo);

    int row_stride = cinfo.output_width * cinfo.output_components;
    JSAMPARRAY buffer = (*cinfo.mem->alloc_sarray)((j_common_ptr)&cinfo, JPOOL_IMAGE, row_stride, 1);

    ret.width = cinfo.output_width;
    ret.height = cinfo.output_height;
    ret.data.resize(ret.width * ret.height);
    while (cinfo.output_scanline < cinfo.output_height) {
        jpeg_read_scanlines(&cinfo, buffer, 1);
        for (unsigned int i = 0; i < cinfo.output_width; ++i) {
            ret.data[(cinfo.output_scanline - 1) * cinfo.output_width + i] =
                buffer[0][i * 3]; // Y component is the first byte in YCbCr
        }
    }

    jpeg_finish_decompress(&cinfo);
    jpeg_destroy_decompress(&cinfo);
    fclose(infile);
    return std::make_optional(ret);
}