/*
    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.
*/

struct input_t {
    const uint8_t *ptr;
    size_t size, pos;
};

// feeds libpng more data from the memory stream
// (can be customized to read from any arbitrary source, not just memory)
void user_read_data(png_structp png_ptr, png_bytep data, png_size_t length)
{
    input_t *in = (input_t *)png_get_io_ptr(png_ptr);

    if (in->pos + length > in->size) {
        // Now we're in big trouble. There doesn't seem to be any way
        // to signal libpng that we can't read from the source.
        return;
    }

    memcpy(data, in->ptr + in->pos, length);
    in->pos += length;
}

/*
    Loads a 24/32-bit PNG image from a memory buffer. Returns true if ok, false if
    something went wrong. If successfull, 'w' and 'h' contains the image size in
    pixels, 'bpp' is the bytes per pixel (3 or 4) and 'data' contains the raw pixel
    data. 'data' must be NULL before calling this.
*/
bool loadMemoryPNG(const uint8_t *input, const size_t size, uint32_t &w, uint32_t &h, uint8_t &bpp, uint8_t *&data)
{
    if (!input || size < 8 || data)
        return false;

    // check the signature
    if (png_sig_cmp(input, 0, 8))
        return false;

    // initialize reading structures
    png_structp png_ptr = png_create_read_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_read_struct(&png_ptr, NULL, NULL);
        return false;
    }

    // setup error handling
    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
        delete [] data;
        data = NULL;
        return false;
    }

    // initialize input
    input_t in;

    in.ptr = input;
    in.size = size;
    in.pos = 8;

    png_set_read_fn(png_ptr, &in, user_read_data);
    png_set_sig_bytes(png_ptr, 8);

    // read the image information
    int depth, colors, interlace, compression, filter;

    png_read_info(png_ptr, info_ptr);

    if (!png_get_IHDR(png_ptr, info_ptr, &w, &h, &depth, &colors, &interlace, &compression, &filter)) {
        png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
        return false;
    }

    // can we load this image?
    if (depth != 8 ||
        (colors != PNG_COLOR_TYPE_RGB && colors != PNG_COLOR_TYPE_RGBA) ||
        interlace != PNG_INTERLACE_NONE ||
        compression != PNG_COMPRESSION_TYPE_BASE ||
        filter != PNG_FILTER_TYPE_BASE) {

        png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
        return false;
    }

    if (colors == PNG_COLOR_TYPE_RGB)
        bpp = 3;
    else bpp = 4;

    // allocate memory (no exceptions here)
    data = new(std::nothrow) uint8_t[w * h * bpp];

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

    // Read the image, one row at a time (allows customization and no row pointer mess).
    // The downside of this is that we can't load interlaced images.
    const uint32_t stride = w * bpp;
    png_byte *row = (png_byte *)data;

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

    // done
    png_read_end(png_ptr, info_ptr);
    png_destroy_read_struct(&png_ptr, &info_ptr, NULL);

    return true;
}