#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>

#include <png.h>

#include "write_png.h"


#ifndef PNG_UINT_32_MAX
#  define PNG_UINT_32_MAX 	0x7fffffffUL
#  define png_sizeof 		sizeof
#endif

#ifndef png_infopp_NULL
#  define png_infopp_NULL 	(png_infopp)NULL
#endif

/* State vars used to control some save defaults */
static int zi_use_BGR = 0; /* If > 0, assume data is BGR rather than RGB packed */
static int zi_mono_inverted = 0; /* If > 0, b/w invert greyscale images */

/********************************************************************************/
void set_png_bgr_state(int iv)
{
	zi_use_BGR = iv;
}

/********************************************************************************/
 int get_png_bgr_state(void)
{
	return zi_use_BGR;
}

/********************************************************************************/
void set_png_mono_invert_state(int iv)
{
	zi_mono_inverted = iv;
}

/********************************************************************************/
 int get_png_mono_invert_state(void)
{
	return zi_mono_inverted;
}

/********************************************************************************/
/* Throw away warnings (die on errors as usual) */
static void user_warning_fn (png_structp png_ptr, png_const_charp warning_msg)
{
	puts("PNG writer warning");
	fflush(stdout);
}

/********************************************************************************/
/* write a png file */
int write_png(const char *file_name, const unsigned char *datap, int width, int height, int depth,
			  const char *keyword, const char *description)
{
	FILE *fp;
	png_structp png_ptr;
	png_infop   info_ptr;
	int bit_depth = 8; /* Assume all our data is based on 8-bit chars */
	int idx;
	png_bytep row_pointers[height];
	png_text text_chunk[1];
	png_color_8 sig_bit;

	/* open the file */
	fp = fopen(file_name, "wb");
	if (fp == NULL)
	{
		return (WRITE_DATA_ERROR);
	}

   /* Create and initialize the png_struct with the desired error handler
    * functions.  If you want to use the default stderr and longjump method,
    * you can supply NULL for the last three parameters.  We also check that
    * the library version is compatible with the one used at compile time,
    * in case we are using dynamically linked libraries.  REQUIRED.
    */
	png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, (void*)NULL, NULL, user_warning_fn);

	if (png_ptr == NULL)
	{
		fclose(fp);
		return (WRITE_DATA_ERROR);
	}

   /* Allocate/initialize the image information data.  REQUIRED */
	info_ptr = png_create_info_struct(png_ptr);
	if (info_ptr == NULL)
	{
		fclose(fp);
		png_destroy_write_struct(&png_ptr,  png_infopp_NULL);
		return (WRITE_DATA_ERROR);
	}

   /* Set error handling.  REQUIRED if you aren't supplying your own
    * error handling functions in the png_create_write_struct() call.
    */
	if (setjmp(png_jmpbuf(png_ptr)))
	{
		/* If we get here, we had a problem accessing the file */
		fclose(fp);
		png_destroy_write_struct(&png_ptr, &info_ptr);
		return (WRITE_DATA_ERROR);
	}

	/* set up the output control if you are using standard C streams */
	png_init_io(png_ptr, fp);

   /* Set the image information here.  Width and height are up to 2^31,
    * bit_depth is one of 1, 2, 4, 8, or 16, but valid values also depend on
    * the color_type selected. color_type is one of PNG_COLOR_TYPE_GRAY,
    * PNG_COLOR_TYPE_GRAY_ALPHA, PNG_COLOR_TYPE_PALETTE, PNG_COLOR_TYPE_RGB,
    * or PNG_COLOR_TYPE_RGB_ALPHA.  interlace is either PNG_INTERLACE_NONE or
    * PNG_INTERLACE_ADAM7, and the compression_type and filter_type MUST
    * currently be PNG_COMPRESSION_TYPE_BASE and PNG_FILTER_TYPE_BASE. REQUIRED
    */

	if (depth == 1)
	{
		png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, PNG_COLOR_TYPE_GRAY,
		   PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	}
	else if (depth == 2) /* Not sure we want to support this mode... */
	{
		png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, PNG_COLOR_TYPE_GRAY_ALPHA,
		   PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	}
	else if (depth == 3)
	{
		png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, PNG_COLOR_TYPE_RGB,
		   PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	}
	else if (depth == 4)
	{
		png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth, PNG_COLOR_TYPE_RGB_ALPHA,
		   PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
	}
	else
	{
		/* If we get here, the wrong data depth was passed */
		fclose(fp);
		png_destroy_write_struct(&png_ptr, &info_ptr);
		return (WRITE_DATA_ERROR);
	}

	/* Configure the optional "significant bits" chunk */
	if (depth == 1)
	{
		/* if we are dealing with a grayscale image */
		sig_bit.gray = bit_depth;
   	}
	else if (depth >= 3)
	{
		/* if we are dealing with a color image */
		sig_bit.red   = bit_depth;
		sig_bit.green = bit_depth;
		sig_bit.blue  = bit_depth;
	}
	else if (depth == 4)
	{
		/* if the image has an alpha channel */
		sig_bit.alpha = bit_depth;
	}
	png_set_sBIT(png_ptr, info_ptr, &sig_bit);


   /* Optional gamma chunk is strongly suggested if you have any guess
    * as to the correct gamma of the image. */
	png_set_gAMA(png_ptr, info_ptr, 1.0);

	/* Optionally write comments into the image */
	text_chunk[0].key = keyword;
	text_chunk[0].text = description;
	text_chunk[0].compression = PNG_TEXT_COMPRESSION_NONE;
#ifdef PNG_iTXt_SUPPORTED
	text_chunk[0].lang = NULL;
#endif
	png_set_text(png_ptr, info_ptr, text_chunk, 1);

	/* Write the file header information.  REQUIRED */
	png_write_info(png_ptr, info_ptr);

   /* Set up any transformations you want. These are
    * optional. Only set them if you need them. */

	/* invert monochrome pixels */
	if ((zi_mono_inverted > 0) && (depth == 1))
	{
		png_set_invert_mono(png_ptr);
	}

	/* flip BGR pixels to RGB */
	if ((zi_use_BGR > 0) && (depth == 3))
	{
		png_set_bgr(png_ptr);
	}

	if (height > PNG_UINT_32_MAX / png_sizeof(png_bytep))
	{
		png_error (png_ptr, "Image is too tall to process in memory");
	}

	for (idx = 0; idx < height; idx++)
	{
		row_pointers[idx] = &datap[idx * width * depth];
	}

	/* write out the entire image data in one call */
	png_write_image(png_ptr, row_pointers);


	/* It is REQUIRED to call this to finish writing the rest of the file */
	png_write_end(png_ptr, info_ptr);

	/* clean up after the write, and free any memory allocated */
	png_destroy_write_struct(&png_ptr, &info_ptr);

	/* close the file */
	fclose(fp);

	/* that's it */
	return (WRITE_DATA_OK);
}

/********************************************************************************/
#ifdef STANDALONE_WRITER_TEST
int main (void)
{
	int w, h, d;
	unsigned char * datap;
	int idx, jdx;
	int res;
	char desc[512];

	w = h = 256;
	d = 3;

	datap = malloc(w * h * d);

	for (idx = 0; idx < h; idx ++)
	{
		for (jdx = 0; jdx < w; jdx++)
		{
			int cell = (idx * h * d) + (jdx * d);
			datap[cell] = jdx;
			datap[cell + 1] = (idx + jdx) & 0xFF;
			datap[cell + 2] = idx;
		}
	}

	sprintf(desc, "Test Image - version 0\nsize %d %d %d", w, h, d);
	res =  write_png("Test0.png", datap, w, h, d, "SS7K-imageV1", desc);

	set_png_bgr_state(1);
	res =  write_png("Test1.png", datap, w, h, d, "SS7K-imageV2", "Fixed Text String");

	if (res) puts("T1 failed");

	free(datap);

	d = 1;
	datap = malloc(w * h);

	for (idx = 0; idx < h; idx ++)
	{
		for (jdx = 0; jdx < w; jdx++)
		{
			int cell = (idx * h) + jdx;
			datap[cell] = jdx;
		}
	}
	sprintf(desc, "Test Image - version 2\nsize %d %d %d\nThis is Flipped Mono\ndone!", w, h, d);
	set_png_mono_invert_state(1);
	res =  write_png("Test2.png", datap, w, h, d, "SS7K-image", desc);

	sprintf(desc, "Test Image - version 3\nsize %d %d %d\nThis is Mono\ndone!", w, h, d);
	set_png_mono_invert_state(0);
	res =  write_png("Test3.png", datap, w, h, d, "SS7K-image", desc);

	if (res) puts("T2 failed");

	free(datap);

	if(!res)
		puts("PNG test OK");

	return 0;
}
#endif
/********************************************************************************/
/* End of File */