/***************************************************************************

	png.c

    PSP PNG format image I/O functions. (based on M.A.M.E. PNG functions)

***************************************************************************/

#include "psp.h"
#include <math.h>
#include "zlib/zlib.h"

#define PNG_Signature       "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A"
#define MNG_Signature       "\x8A\x4D\x4E\x47\x0D\x0A\x1A\x0A"

#define PNG_CN_IHDR 0x49484452L     /* Chunk names */
#define PNG_CN_PLTE 0x504C5445L
#define PNG_CN_IDAT 0x49444154L
#define PNG_CN_IEND 0x49454E44L
#define PNG_CN_gAMA 0x67414D41L
#define PNG_CN_sBIT 0x73424954L
#define PNG_CN_cHRM 0x6348524DL
#define PNG_CN_tRNS 0x74524E53L
#define PNG_CN_bKGD 0x624B4744L
#define PNG_CN_hIST 0x68495354L
#define PNG_CN_tEXt 0x74455874L
#define PNG_CN_zTXt 0x7A545874L
#define PNG_CN_pHYs 0x70485973L
#define PNG_CN_oFFs 0x6F464673L
#define PNG_CN_tIME 0x74494D45L
#define PNG_CN_sCAL 0x7343414CL

#define PNG_PF_None     0   /* Prediction filters */
#define PNG_PF_Sub      1
#define PNG_PF_Up       2
#define PNG_PF_Average  3
#define PNG_PF_Paeth    4

#define MNG_CN_MHDR 0x4D484452L     /* MNG Chunk names */
#define MNG_CN_MEND 0x4D454E44L
#define MNG_CN_TERM 0x5445524DL
#define MNG_CN_BACK 0x4241434BL


/* PNG support */
struct png_info {
	u32 width, height;
	u32 xres, yres;
	struct rectangle screen;
	double xscale, yscale;
	double source_gamma;
	u32 chromaticities[8];
	u32 resolution_unit, offset_unit, scale_unit;
	u8 bit_depth;
	u32 significant_bits[4];
	u32 background_color[4];
	u8 color_type;
	u8 compression_method;
	u8 filter_method;
	u8 interlace_method;
	u32 num_palette;
	u8 *palette;
	u32 num_trans;
	u8 *trans;
	u8 *image;

	/* The rest is private and should not be used
	 * by the public functions
	 */
	u8 bpp;
	u32 rowbytes;
	u8 *zimage;
	u32 zlength;
	u8 *fimage;
};

/********************************************************************************

  PNG read functions

********************************************************************************/

/* convert_uint is here so we don't have to deal with byte-ordering issues */
static u32 convert_from_network_order(u8 *v)
{
	return (v[0] << 24) | (v[1] << 16) | (v[2] << 8) | (v[3]);
}

static int png_unfilter(struct png_info *p)
{
	int i, j, bpp, filter;
	s32 prediction, pA, pB, pC, dA, dB, dC;
	u8 *src, *dst;

	if ((p->image = (u8 *)malloc(p->height * p->rowbytes)) == NULL)
	{
		ErrorMsg("PNG: Out of memory");
		free(p->fimage);
		return 0;
	}

	src = p->fimage;
	dst = p->image;
	bpp = p->bpp;

	for (i = 0; i < p->height; i++)
	{
		filter = *src++;
		if (!filter)
		{
			memcpy(dst, src, p->rowbytes);
			src += p->rowbytes;
			dst += p->rowbytes;
		}
		else
		{
			for (j = 0; j < p->rowbytes; j++)
			{
				pA = (j < bpp) ? 0 : *(dst - bpp);
				pB = (i < 1) ? 0: *(dst - p->rowbytes);
				pC = (j < bpp || i < 1) ? 0: *(dst - p->rowbytes - bpp);

				switch (filter)
				{
				case PNG_PF_Sub:
					prediction = pA;
					break;

				case PNG_PF_Up:
					prediction = pB;
					break;

				case PNG_PF_Average:
					prediction = ((pA + pB) / 2);
					break;

				case PNG_PF_Paeth:
					prediction = pA + pB - pC;
					dA = abs(prediction - pA);
					dB = abs(prediction - pB);
					dC = abs(prediction - pC);
					if (dA <= dB && dA <= dC) prediction = pA;
					else if (dB <= dC) prediction = pB;
					else prediction = pC;
					break;

				default:
					ErrorMsg("PNG: Unknown filter type %i", filter);
					prediction = 0;
				}
				*dst = 0xff & (*src + prediction);
				dst++; src++;
			}
		}
	}

	free(p->fimage);
	return 1;
}

int png_verify_signature(FILE *fp)
{
	s8 signature[8];

	if (fread(signature, 1, 8, fp) != 8)
	{
		ErrorMsg("PNG: Unable to read PNG signature (EOF)");
		return 0;
	}

	if (memcmp(signature, PNG_Signature, 8))
	{
		ErrorMsg("PNG: PNG signature mismatch found: %s expected: %s", signature, PNG_Signature);
		return 0;
	}
	return 1;
}

int png_inflate_image(struct png_info *p)
{
	unsigned long fbuff_size;

	fbuff_size = p->height * (p->rowbytes + 1);

	if ((p->fimage = (u8 *)malloc(fbuff_size)) == NULL)
	{
		ErrorMsg("PNG: Out of memory");
		free(p->zimage);
		return 0;
	}

	if (uncompress(p->fimage, &fbuff_size, p->zimage, p->zlength) != Z_OK)
	{
		ErrorMsg("PNG: Error while inflating image");
		return 0;
	}

	free(p->zimage);
	return 1;
}

int png_read_file(FILE *fp, struct png_info *p)
{
	/* translates color_type to bytes per pixel */
	const int samples[] = {1, 0, 3, 1, 2, 0, 4};

	u32 chunk_length, chunk_type=0, chunk_crc, crc;
	u8 *chunk_data, *temp;
	u8 str_chunk_type[5], v[4];

	struct idat
	{
		struct idat *next;
		int length;
		u8 *data;
	} *ihead, *pidat;

	if ((ihead = malloc(sizeof(struct idat))) == 0)
		return 0;

	pidat = ihead;

	p->zlength = 0;
	p->num_palette = 0;
	p->num_trans = 0;
	p->trans = NULL;
	p->palette = NULL;

	if (png_verify_signature(fp) == 0)
		return 0;

	while (chunk_type != PNG_CN_IEND)
	{
		if (fread(v, 1, 4, fp) != 4)
			ErrorMsg("PNG: Unexpected EOF in PNG");
		chunk_length = convert_from_network_order(v);

		if (fread(str_chunk_type, 1, 4, fp) != 4)
			ErrorMsg("PNG: Unexpected EOF in PNG file");

		str_chunk_type[4] = 0; /* terminate string */

		crc = crc32(0, str_chunk_type, 4);
		chunk_type = convert_from_network_order(str_chunk_type);

		if (chunk_length)
		{
			if ((chunk_data = (u8 *)malloc(chunk_length + 1)) == NULL)
			{
				ErrorMsg("PNG: Out of memory");
				return 0;
			}
			if (fread(chunk_data, 1, chunk_length, fp) != chunk_length)
			{
				ErrorMsg("PNG: Unexpected EOF in PNG file");
				free(chunk_data);
				return 0;
			}

			crc = crc32(crc, chunk_data, chunk_length);
		}
		else
			chunk_data = NULL;

		if (fread(v, 1, 4, fp) != 4)
			ErrorMsg("PNG: Unexpected EOF in PNG");
		chunk_crc = convert_from_network_order(v);

		if (crc != chunk_crc)
		{
			ErrorMsg("PNG: CRC check failed while reading PNG chunk %s", str_chunk_type);
//			ErrorMsg("PNG: Found: %08X  Expected: %08X", crc, chunk_crc);
			return 0;
		}

		switch (chunk_type)
		{
		case PNG_CN_IHDR:
			p->width = convert_from_network_order(chunk_data);
			p->height = convert_from_network_order(chunk_data + 4);
			p->bit_depth = *(chunk_data + 8);
			p->color_type = *(chunk_data + 9);
			p->compression_method = *(chunk_data + 10);
			p->filter_method = *(chunk_data + 11);
			p->interlace_method = *(chunk_data + 12);
			free(chunk_data);
			break;

		case PNG_CN_PLTE:
			p->num_palette = chunk_length/3;
			p->palette = chunk_data;
			break;

		case PNG_CN_tRNS:
			p->num_trans = chunk_length;
			p->trans = chunk_data;
			break;

		case PNG_CN_IDAT:
			pidat->data = chunk_data;
			pidat->length = chunk_length;
			if ((pidat->next = malloc(sizeof(struct idat))) == NULL)
				return 0;
			pidat = pidat->next;
			pidat->next = 0;
			p->zlength += chunk_length;
			break;

		case PNG_CN_tEXt:
			{
				char *text = (char *)chunk_data;

				while(*text++) ;
				chunk_data[chunk_length]=0;
			}
			free(chunk_data);
			break;

		case PNG_CN_tIME:
			free(chunk_data);
			break;

		case PNG_CN_gAMA:
			p->source_gamma	 = convert_from_network_order(chunk_data) / 100000.0;
			free(chunk_data);
			break;

		case PNG_CN_pHYs:
			p->xres = convert_from_network_order(chunk_data);
			p->yres = convert_from_network_order(chunk_data + 4);
			p->resolution_unit = *(chunk_data + 8);
			free(chunk_data);
			break;

		case PNG_CN_IEND:
			break;

		default:
			if (chunk_data) free(chunk_data);
			break;
		}
	}

	if ((p->zimage = (u8 *)malloc(p->zlength)) == NULL)
	{
		ErrorMsg("PNG: Out of memory");
		return 0;
	}

	/* combine idat chunks to compressed image data */
	temp = p->zimage;
	while (ihead->next)
	{
		pidat = ihead;
		memcpy(temp, pidat->data, pidat->length);
		free(pidat->data);
		temp += pidat->length;
		ihead = pidat->next;
		free(pidat);
	}
	p->bpp = (samples[p->color_type] * p->bit_depth) / 8;
	p->rowbytes = ceil((p->width * p->bit_depth * samples[p->color_type]) / 8.0);

	if (png_inflate_image(p) == 0)
		return 0;

	if (png_unfilter(p) == 0)
		return 0;

	return 1;
}


/********************************************************************************

  PNG write functions (16bit color only)

********************************************************************************/

struct png_text
{
	char *data;
	int length;
	struct png_text *next;
};

static struct png_text *png_text_list = 0;

static void convert_to_network_order(u32 i, u8 *v)
{
	v[0] = (i >> 24) & 0xff;
	v[1] = (i >> 16) & 0xff;
	v[2] = (i >>  8) & 0xff;
	v[3] = (i >>  0) & 0xff;
}

static int png_add_text(const char *keyword, const char *text)
{
	struct png_text *pt;

	pt = malloc(sizeof(struct png_text));
	if (pt == 0)
		return 0;

	pt->length = strlen(keyword) + strlen(text) + 1;
	pt->data = malloc(pt->length + 1);
	if (pt->data == 0)
		return 0;

	strcpy(pt->data, keyword);
	strcpy(pt->data + strlen(keyword) + 1, text);
	pt->next = png_text_list;
	png_text_list = pt;

	return 1;
}

static int write_chunk(FILE *fp, u32 chunk_type, u8 *chunk_data, u32 chunk_length)
{
	u32 crc;
	u8 v[4];
	int written;

	/* write length */
	convert_to_network_order(chunk_length, v);
	written = fwrite(v, 1, 4, fp);

	/* write type */
	convert_to_network_order(chunk_type, v);
	written += fwrite(v, 1, 4, fp);

	/* calculate crc */
	crc = crc32(0, v, 4);
	if (chunk_length > 0)
	{
		/* write data */
		written += fwrite(chunk_data, 1, chunk_length, fp);
		crc = crc32(crc, chunk_data, chunk_length);
	}
	convert_to_network_order(crc, v);

	/* write crc */
	written += fwrite(v, 1, 4, fp);

	if (written != 3*4+chunk_length)
	{
		ErrorMsg("PNG: Chunk write failed");
		return 0;
	}
	return 1;
}

static int png_write_sig(FILE *fp)
{
	/* PNG Signature */
	if (fwrite(PNG_Signature, 1, 8, fp) != 8)
	{
		ErrorMsg("PNG: PNG sig write failed");
		return 0;
	}
	return 1;
}

static int png_write_datastream(FILE *fp, struct png_info *p)
{
	u8 ihdr[13];
	struct png_text *pt;

	/* IHDR */
	convert_to_network_order(p->width, ihdr);
	convert_to_network_order(p->height, ihdr + 4);
	*(ihdr +  8) = p->bit_depth;
	*(ihdr +  9) = p->color_type;
	*(ihdr + 10) = p->compression_method;
	*(ihdr + 11) = p->filter_method;
	*(ihdr + 12) = p->interlace_method;

	if (write_chunk(fp, PNG_CN_IHDR, ihdr, 13) == 0)
		return 0;

	/* PLTE */
	if (p->num_palette > 0)
		if (write_chunk(fp, PNG_CN_PLTE, p->palette, p->num_palette*3) == 0)
			return 0;

	/* IDAT */
	if (write_chunk(fp, PNG_CN_IDAT, p->zimage, p->zlength) == 0)
		return 0;

	/* tEXt */
	while (png_text_list)
	{
		pt = png_text_list;
		if (write_chunk(fp, PNG_CN_tEXt, (u8 *)pt->data, pt->length) == 0)
			return 0;
		free(pt->data);

		png_text_list = pt->next;
		free(pt);
	}

	/* IEND */
	if (write_chunk(fp, PNG_CN_IEND, NULL, 0) == 0)
		return 0;

	return 1;
}

static int png_filter(struct png_info *p)
{
	int i;
	u8 *src, *dst;

	if ((p->fimage = (u8 *)malloc(p->height * (p->rowbytes + 1))) == NULL)
	{
		ErrorMsg("PNG: Out of memory.");
		return 0;
	}

	dst = p->fimage;
	src = p->image;

	for (i = 0; i < p->height; i++)
	{
		*dst++ = 0; /* No filter */
		memcpy(dst, src, p->rowbytes);
		src += p->rowbytes;
		dst += p->rowbytes;
	}
	return 1;
}

static int png_deflate_image(struct png_info *p)
{
	unsigned long zbuff_size;

	zbuff_size = (p->height * (p->rowbytes + 1)) * 1.1 + 12;

	if ((p->zimage = (u8 *)malloc(zbuff_size)) == NULL)
	{
		ErrorMsg("PNG: Out of memory.");
		return 0;
	}

	if (compress(p->zimage, &zbuff_size, p->fimage, p->height * (p->rowbytes + 1)) != Z_OK)
	{
		ErrorMsg("PNG: Error while deflating image");
		return 0;
	}
	p->zlength = zbuff_size;

	return 1;
}

static int png_pack_buffer(struct png_info *p)
{
	u8 *outp, *inp;
	int i,j,k;

	outp = inp = p->image;

	if (p->bit_depth < 8)
	{
		for (i = 0; i < p->height; i++)
		{
			for (j=0; j<p->width/(8/p->bit_depth); j++)
			{
				for (k=8/p->bit_depth-1; k>=0; k--)
					*outp |= *inp++ << k * p->bit_depth;
				outp++;
				*outp = 0;
			}
			if (p->width % (8/p->bit_depth))
			{
				for (k=p->width%(8/p->bit_depth)-1; k>=0; k--)
					*outp |= *inp++ << k * p->bit_depth;
				outp++;
				*outp = 0;
			}
		}
	}
	return 1;
}

static int png_create_datastream(FILE *fp)
{
	int i, j;
	u8 *ip;
	struct png_info p;
	u16 *vptr, *src;

	memset(&p, 0, sizeof (struct png_info));
	p.xscale = p.yscale = p.source_gamma = 0.0;
	p.palette = p.trans = p.image = p.zimage = p.fimage = NULL;
	p.width = SCR_WIDTH;
	p.height = SCR_HEIGHT;
	p.color_type = 2;
	p.rowbytes = p.width * 3;
	p.bit_depth = 8;

	if ((p.image = (u8 *)malloc(p.height * p.rowbytes))==NULL)
	{
		ErrorMsg("PNG: Out of memory.");
		return 0;
	}

	ip = p.image;
	vptr = video_frame_addr(show_frame, 0, 0);

	for (i = 0; i < p.height; i++)
	{
		src = &vptr[i * BUF_WIDTH];

		for (j = 0; j < p.width; j++)
		{
			u16 color = src[j];
			*ip++ = (u8)GETR15(color);
			*ip++ = (u8)GETG15(color);
			*ip++ = (u8)GETB15(color);
		}
	}

	if (png_filter(&p) == 0)
		return 0;

	if (png_deflate_image(&p) == 0)
		return 0;

	if (png_write_datastream(fp, &p) == 0)
		return 0;

	if (p.palette) free(p.palette);
	if (p.image) free(p.image);
	if (p.zimage) free(p.zimage);
	if (p.fimage) free(p.fimage);
	return 1;
}


/********************************************************************************

  Interface for PSP

********************************************************************************/

/*--------------------------------------------------------
	Load PNG
 -------------------------------------------------------*/

int load_png(const char *name)
{
	struct png_info p;
	FILE *fp;
	int res;

	memset(&p, 0, sizeof(struct png_info));

	video_clear_frame(draw_frame);

	if ((fp = fopen(name, "rb")) == NULL)
		return 0;

	if ((res = png_read_file(fp, &p)))
	{
		int x, y, sx, sy, width, height;
		u8 *ip;
		u16 *vptr, *dst;

		width = p.width > SCR_WIDTH ? SCR_WIDTH : p.width;
		height = p.height > SCR_HEIGHT ? SCR_HEIGHT : p.height;

		sx = (SCR_WIDTH - width) >> 1;
		sy = (SCR_HEIGHT - height) >> 1;

		ip = p.image;
		vptr = video_frame_addr(draw_frame, sx, sy);

		switch (p.bpp * p.bit_depth)
		{
		case 8:
			for (y = 0; y < height; y++)
			{
				dst = &vptr[y * BUF_WIDTH];

				for (x = 0; x < width; x++)
				{
					int color = ip[x];
					int r = p.palette[color * 3 + 0];
					int g = p.palette[color * 3 + 1];
					int b = p.palette[color * 3 + 2];
					dst[x] = MAKECOL15(r, g, b);
				}
				ip += p.rowbytes;
			}
			break;

		case 24:
			for (y = 0; y < height; y++)
			{
				dst = &vptr[y * BUF_WIDTH];

				for (x = 0; x < width; x++)
				{
					int r = ip[x * 3 + 0];
					int g = ip[x * 3 + 1];
					int b = ip[x * 3 + 2];
					dst[x] = MAKECOL15(r, g, b);
				}
				ip += p.rowbytes;
			}
			break;

		default:
			ErrorMsg("PNG: %d bit color not supported.", p.bpp * p.bit_depth);
			break;
		}

		if (p.palette) free(p.palette);
		if (p.image) free(p.image);
	}

	fclose(fp);

	return res;
}


/*--------------------------------------------------------
	Save PNG
 -------------------------------------------------------*/

int save_png(const char *name)
{
	FILE *fp;
	int res;

	if ((fp = fopen(name, "wb")) == NULL)
		return 0;

	if ((res = png_add_text("Software", APPNAME_STR " " VERSION_STR)))
	{
		if ((res = png_add_text("System", "PSP")))
		{
			if ((res = png_write_sig(fp)))
			{
				res = png_create_datastream(fp);
			}
		}
	}

	fclose(fp);

	return res;
}
