#include <png.h>
#include <string>
#include <vector>


// 指定したpngファイルをRGBA画像として読み込む。
// widthとheightに幅、高さが指定され、pixelにr,g,b,aの順でピクセルデータが格納される
bool LoadPNG(const std::string& filename, int& width, int& height, std::vector<unsigned char>& pixel);

// 指定したファイルにRGBA png画像として書き込む。
// 渡すデータ(pixel)のサイズはwindth*height*4である必要がある
bool SavePNG(const std::string& filename, int width, int height, std::vector<unsigned char>& pixel);





bool LoadPNG(const std::string& filename, int& width, int& height, std::vector<unsigned char>& pixel)
{
	FILE *fp = fopen(filename.c_str(), "rb");
	if(fp==0)
		return false;

	png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
	if(png_ptr==0) {
		fclose(fp);
		return false;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if(info_ptr==0) {
		fclose(fp);
		png_destroy_read_struct(&png_ptr, png_infopp_NULL, png_infopp_NULL);
		return false;
	}

	png_init_io(png_ptr, fp);

	png_uint_32 _width, _height;
	int bit_depth, color_type, interlace_type;

	png_read_info(png_ptr, info_ptr);
	png_get_IHDR(png_ptr, info_ptr, &_width, &_height, &bit_depth, &color_type, &interlace_type, int_p_NULL, int_p_NULL);
	width = _width;
	height = _height;
	pixel.resize(width*height*4);

	png_set_strip_16(png_ptr);
	png_set_packing(png_ptr);
	if(color_type==PNG_COLOR_TYPE_PALETTE)
		png_set_palette_to_rgb(png_ptr);
	if(png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
		png_set_tRNS_to_alpha(png_ptr);
	png_read_update_info(png_ptr, info_ptr);


	std::vector<png_bytep> row_pointers(height);
	for(int row=0; row<height; ++row)
		row_pointers[row] = (png_bytep)png_malloc(png_ptr, png_get_rowbytes(png_ptr, info_ptr));
	png_read_image(png_ptr, &row_pointers[0]);

	if(color_type==PNG_COLOR_TYPE_RGB_ALPHA) {
		for(int i=0; i<height; ++i) {
			for(int j=0; j<width; ++j) {
				pixel[i*width*4+j*4+0] = row_pointers[i][j*4+0];
				pixel[i*width*4+j*4+1] = row_pointers[i][j*4+1];
				pixel[i*width*4+j*4+2] = row_pointers[i][j*4+2];
				pixel[i*width*4+j*4+3] = row_pointers[i][j*4+3];
			}
		}
	}
	else if(color_type==PNG_COLOR_TYPE_RGB) {
		for(int i=0; i<height; ++i) {
			for(int j=0; j<width; ++j) {
				pixel[i*width*4+j*4+0] = row_pointers[i][j*3+0];
				pixel[i*width*4+j*4+1] = row_pointers[i][j*3+1];
				pixel[i*width*4+j*4+2] = row_pointers[i][j*3+2];
				pixel[i*width*4+j*4+3] = 255;
			}
		}
	}

	for(int row=0; row<height; ++row)
		png_free(png_ptr, row_pointers[row]);

	png_read_end(png_ptr, info_ptr);
	png_destroy_read_struct(&png_ptr, &info_ptr, png_infopp_NULL);
	fclose(fp);
	return true;
}




bool SavePNG(const std::string& filename, int width, int height, std::vector<unsigned char>& pixel)
{
	FILE *fp = fopen(filename.c_str(), "wb");
	if(fp==0)
	return false;

	png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
	if(png_ptr==0) {
		fclose(fp);
		return false;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if(info_ptr==0) {
		fclose(fp);
		png_destroy_write_struct(&png_ptr,  png_infopp_NULL);
		return false;
	}

	png_init_io(png_ptr, fp);
	png_set_IHDR(png_ptr, info_ptr, width, height, 8,
		PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	png_write_info(png_ptr, info_ptr);

	std::vector<png_bytep> row_pointers(height);
	for(int i=0; i<height; ++i)
		row_pointers[i] = &pixel[i*width*4];

	png_write_image(png_ptr, &row_pointers[0]);

	png_write_end(png_ptr, info_ptr);
	png_destroy_write_struct(&png_ptr, &info_ptr);
	fclose(fp);
	return true;
}