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


#include <vector>
#include <cstdint>
#include <cmath>
#include <png.h>
#include <string>
#include <cstring>

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

#include "JFJochPNG.h"

struct PNGMemoryWriter {
    std::string buffer;
};

void pngWriteCallback(png_structp png_ptr, png_bytep data, png_size_t length) {
    auto writer = static_cast<PNGMemoryWriter*>(png_get_io_ptr(png_ptr));
    size_t old_size = writer->buffer.size();
    writer->buffer.resize(old_size + length);
    std::memcpy(writer->buffer.data() + old_size, data, length);
}

std::string WritePNGToMem(const CompressedImage& image, int compression_level) {
    if (image.GetMode() != CompressedImageMode::RGB)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Only RGB images allowed for PNG saving");

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

    int width = image.GetWidth();
    int height = image.GetHeight();
    int color_type = PNG_COLOR_TYPE_RGB;
    int bit_depth = 8;

    // Initialize libpng structures
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,
                                                  NULL, NULL, NULL);
    if (!png_ptr) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error initializing PNG write structure");
    }

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_write_struct(&png_ptr, NULL);
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error initializing PNG info structure");
    }

    // Set up error handling
    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_write_struct(&png_ptr, &info_ptr);
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Error during PNG creation");
    }

    PNGMemoryWriter writer;

    writer.buffer.reserve(width * height * 3);

    png_set_write_fn(png_ptr, &writer, pngWriteCallback, NULL);

    // Set image information
    png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, color_type,
                 PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

    // Set compression level (0-9, where 0 is no compression and 9 is maximum)
    png_set_compression_level(png_ptr, compression_level);

    // Write PNG header
    png_write_info(png_ptr, info_ptr);

    // Allocate row pointers
    std::vector<png_bytep> row_pointers(height);
    for (int y = 0; y < height; y++) {
        row_pointers[y] = buffer.data() + y * width * 3;
    }

    // Write PNG image data
    png_write_image(png_ptr, row_pointers.data());

    // Finish writing
    png_write_end(png_ptr, NULL);
    png_destroy_write_struct(&png_ptr, &info_ptr);

    return writer.buffer;
}