/*
    This file is placed into public domain. There is no warranty of any kind for any purpose.
    Read http://z0b.kapsi.fi/snippets.php before using this code. Thank you.
*/

// This example uses std::vector to store the output, but feel free to change that.

void user_write_data(png_structp png_ptr, png_bytep data, png_size_t length)
{
    std::vector<uint8_t> *output = (std::vector<uint8_t> *)png_get_io_ptr(png_ptr);
    output->insert(output->end(), data, data + length);
}

// needed, but not used
void user_flush_data(png_structp png_ptr)
{
}

/*
    Saves a 24/32-bit PNG image into a memory buffer, returns false if fails.
    'bpp' must be either 3 (for RGB) or 4 (for RGBA) image.
*/
bool saveMemoryPNG(std::vector<uint8_t> &output, const uint32_t w, const uint32_t h, const uint8_t bpp, const uint8_t *data)
{
    // sanity checks
    if (!data || w < 1 || h < 1 || (bpp != 3 && bpp != 4))
        return false;

    // initialize writing structures
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

    if (!png_ptr)
        return false;

    png_infop info_ptr = png_create_info_struct(png_ptr);

    if (!info_ptr) {
        png_destroy_write_struct(&png_ptr, NULL);
        return false;
    }

    // setup error handling
    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return false;
    }

    // initialize output
    png_set_write_fn(png_ptr, &output, user_write_data, user_flush_data);
    output.clear();

    // this is optional, set the zlib compression level
    png_set_compression_level(png_ptr, Z_BEST_COMPRESSION);

    // set the image properties
    png_set_IHDR(png_ptr, info_ptr, w, h, 8,
                 bpp == 3 ? PNG_COLOR_TYPE_RGB : PNG_COLOR_TYPE_RGBA,
                 PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT,
                 PNG_FILTER_TYPE_DEFAULT);

    // write the image data (one scanline at a time, it can be customized easily
    // and we don't have to mess with row pointers)
    png_write_info(png_ptr, info_ptr);

    const uint32_t stride = w * bpp;
    png_byte *row = (png_byte *)data;

    for (uint32_t y = 0; y < h; y++) {
        png_write_row(png_ptr, row);
        row += stride;
    }

    // done
    png_write_end(png_ptr, info_ptr);
    png_destroy_write_struct(&png_ptr, &info_ptr);

    return true;
}