/*
 * LWS PNG -- derived from uPNG -- derived from LodePNG version 20100808
 * Stateful, linewise PNG decode requiring ~36KB fixed heap
 *
 * Copyright (c) 2005-2010 Lode Vandevenne (LodePNG)
 * Copyright (c) 2010 Sean Middleditch (uPNG)
 * Copyright (c) 2021 Andy Green <andy@warmcat.com> (Stateful, incremental)
 *
 * This software is provided 'as-is', without any express or implied
 * warranty. In no event will the authors be held liable for any damages
 * arising from the use of this software.

 * Permission is granted to anyone to use this software for any purpose,
 * including commercial applications, and to alter it and redistribute it
 * freely, subject to the following restrictions:
 *
 *   1. The origin of this software must not be misrepresented; you must not
 *      claim that you wrote the original software. If you use this software
 *      in a product, an acknowledgment in the product documentation would be
 *	appreciated but is not required.
 *
 *   2. Altered source versions must be plainly marked as such, and must not be
 *	misrepresented as being the original software.
 *
 *   3. This notice may not be removed or altered from any source
 *	distribution.
 *
 *  AG: The above notice is the ZLIB license, libpng also uses it.
 *
 * This version was rewritten from the upng project's fork of lodepng and
 * adapted to be a stateful stream parser.  This rewrite retains the ZLIB
 * license of the source material for simplicity.
 *
 * That allows it to use a fixed 32KB ringbuffer to hold decodes, and
 * incrementally decode chunks into it as we want output lines that are not yet
 * present there.  The input png nor the output bitmap need to be all in one
 * place at one time.
 */

#include <private-lib-core.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>

typedef enum upng_color {
	LWS_UPNG_LUM		= 0,
	LWS_UPNG_RGB		= 2,
	LWS_UPNG_LUMA		= 4,
	LWS_UPNG_RGBA		= 6
} upng_color;

struct upng_unfline {
	uint8_t			*recon;
	const uint8_t		*scanline;
	const uint8_t		*precon;
	uint8_t			filterType;
	unsigned int		bypp;
	unsigned int		bypl;

	const uint8_t		*in;
	uint8_t			*lines;
	unsigned int		bpp;

	unsigned int		y;
	unsigned long		diff;
	unsigned long		ibp;
	unsigned long		sp;

	char			padded;
	char			alt;
};

typedef enum {
	UOF_MAGIC,
	UOF_SKIP,
	UOF_TYPE4,
	UOF_WIDTH4,
	UOF_HEIGHT4,
	UOF_CDEPTH,
	UOF_CTYPE,
	UOF_ONLY_ZERO3,
	UOF_SKIP4,

	UOF_CHUNK_LEN,
	UOF_CHUNK_TYPE,
	UOF_INSIDE,

	UOF_SKIP_CHUNK_LEN,
} upng_outer_framing_t;



struct lws_upng_t {
	struct upng_unfline	u;
	inflator_ctx_t		inf;

	unsigned int		width;
	unsigned int		height;

	upng_color		color_type;
	unsigned int		color_depth;
	lws_upng_format_t	format;

	const uint8_t		*chunk;

	int			sctr;
	uint32_t		acc;

	uint32_t		chunklen;
	uint32_t		ctype;

	upng_outer_framing_t	of;

	uint8_t			no_more_input;
	char			hold_at_metadata;
};

static lws_stateful_ret_t
lws_upng_decode(lws_upng_t *upng, const uint8_t **buf, size_t *size);

static int
paeth(int a, int b, int c)
{
	int p = a + b - c;
	int pa = p > a ? p - a : a - p;
	int pb = p > b ? p - b : b - p;
	int pc = p > c ? p - c : c - p;

	if (pa <= pb && pa <= pc)
		return a;

	if (pb <= pc)
		return b;

	return c;
}

static lws_stateful_ret_t
unfilter_scanline(lws_upng_t *u)
{
	struct upng_unfline *uf = &u->u;
	unsigned long i;

	switch (uf->filterType) {
	case 0: /* None */
		for (i = 0; i < uf->bypl; i++)
			uf->recon[i] = u->inf.out[(uf->sp + i) %
			                            u->inf.info_size];
		break;
	case 1: /* Sub */
		for (i = 0; i <  uf->bypp; i++)
			uf->recon[i] = u->inf.out[(uf->sp + i) %
			                            u->inf.info_size];

		for (i = uf->bypp; i < uf->bypl; i++)
			uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
			                            u->inf.info_size] +
				uf->recon[i - uf->bypp]);
		break;
	case 2: /* Up */
		if (uf->y)
			for (i = 0; i < uf->bypl; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				            u->inf.info_size] + uf->precon[i]);
		else
			for (i = 0; i < uf->bypl; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				                          u->inf.info_size]);
		break;
	case 3: /* Average */
		if (uf->y) {
			for (i = 0; i < uf->bypp; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				        u->inf.info_size] + uf->precon[i] / 2);
			for (i = uf->bypp; i < uf->bypl; i++)
				uf->recon[i] = (uint8_t)
					(u->inf.out[(uf->sp + i) %
					            u->inf.info_size] +
					((uf->recon[i - uf->bypp] +
							uf->precon[i]) / 2));
		} else {
			for (i = 0; i < uf->bypp; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				                         u->inf.info_size]);
			for (i = uf->bypp; i < uf->bypl; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				                         u->inf.info_size] +
					uf->recon[i - uf->bypp] / 2);
		}
		break;
	case 4: /* Paeth */
		if (uf->y) {
			for (i = 0; i < uf->bypp; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				                          u->inf.info_size] +
					paeth(0, uf->precon[i], 0));
			for (i = uf->bypp; i < uf->bypl; i++)
				uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
				                          u->inf.info_size] +
					paeth(uf->recon[i - uf->bypp],
							uf->precon[i],
							uf->precon[i - uf->bypp]));
			break;
		}

		for (i = 0; i < uf->bypp; i++)
			uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
			                            u->inf.info_size]);
		for (i = uf->bypp; i < uf->bypl; i++)
			uf->recon[i] = (uint8_t)(u->inf.out[(uf->sp + i) %
			                                u->inf.info_size] +
				paeth(uf->recon[i - uf->bypp], 0, 0));
		break;
	default:
		lwsl_err("%s: line start is broken %d\n", __func__,
				uf->filterType);
		return LWS_SRET_FATAL + 12;
	}

	u->inf.consumed_linear += uf->bypl;

	return LWS_SRET_OK;
}

lws_stateful_ret_t
lws_upng_emit_next_line(lws_upng_t *u, const uint8_t **ppix,
		    const uint8_t **pos, size_t *size, char hold_at_metadata)
{
	struct upng_unfline	*uf = &u->u;
	unsigned long		obp;
	lws_stateful_ret_t	ret = LWS_SRET_OK;

	*ppix = NULL;

	u->hold_at_metadata = hold_at_metadata;

	if (u->height && uf->y >= u->height)
		goto out;

	/*
	 * The decoder emits into the 32KB window ringbuffer, if we don't
	 * already have at least one line's worth of output in there, we'll
	 * have to do more inflation
	 */

	if (u->inf.outpos_linear - u->inf.consumed_linear < uf->bypl + 1) {
		ret = lws_upng_decode(u, pos, size);
		if ((!*size && ret == LWS_SRET_WANT_INPUT) ||
		    (ret & (LWS_SRET_FATAL | LWS_SRET_YIELD)) ||
		    !u->inf.outpos_linear)
			return ret;

		assert(u->inf.info_size);
		assert(uf->bypl + 1);
	}

	if (u->inf.outpos_linear - u->inf.consumed_linear < uf->bypl + 1)
		return ret;

	obp		= uf->alt ? uf->bypl : 0;
	uf->precon	= uf->alt ? uf->lines : uf->lines + uf->bypl;
	uf->recon	= &uf->lines[obp];
	*ppix		= uf->recon;
	uf->filterType	= uf->in[(u->inf.consumed_linear++) % u->inf.info_size];
	uf->sp		= u->inf.consumed_linear % u->inf.info_size;

	if (unfilter_scanline(u) != LWS_SRET_OK) {
		ret = LWS_SRET_FATAL + 13;

		goto out;
	}

	if (uf->padded) {
		unsigned long x;

		for (x = 0; x < (unsigned long)u->width * (unsigned long)uf->bpp; x++) {
			uint8_t bit = (uint8_t)((uf->in[((uf->ibp) >> 3) % u->inf.info_size] >>
						(7 - ((uf->ibp) & 7))) & 1);
			uf->ibp++;

			if (!bit)
				uf->lines[obp >> 3] &=
					(uint8_t)(~(1 << (7 - (obp & 7))));
			else
				uf->lines[obp >> 3] = (uint8_t)(uf->lines[obp >> 3] |
						(uint8_t)(1 << (7 - (obp & 7))));

			obp++;
		}

		uf->ibp += uf->diff;
	}

out:
	uf->alt ^= 1;
	uf->y++;

	return ret;
}

static lws_upng_format_t
determine_format(lws_upng_t* upng) {
	switch (upng->color_type) {
	case LWS_UPNG_LUM:
		switch (upng->color_depth) {
		case 1:
			return LWS_UPNG_LUMINANCE1;
		case 2:
			return LWS_UPNG_LUMINANCE2;
		case 4:
			return LWS_UPNG_LUMINANCE4;
		case 8:
			return LWS_UPNG_LUMINANCE8;
		default:
			return LWS_UPNG_BADFORMAT;
		}
	case LWS_UPNG_RGB:
		switch (upng->color_depth) {
		case 8:
			return LWS_UPNG_RGB8;
		case 16:
			return LWS_UPNG_RGB16;
		default:
			return LWS_UPNG_BADFORMAT;
		}
	case LWS_UPNG_LUMA:
		switch (upng->color_depth) {
		case 1:
			return LWS_UPNG_LUMINANCE_ALPHA1;
		case 2:
			return LWS_UPNG_LUMINANCE_ALPHA2;
		case 4:
			return LWS_UPNG_LUMINANCE_ALPHA4;
		case 8:
			return LWS_UPNG_LUMINANCE_ALPHA8;
		default:
			return LWS_UPNG_BADFORMAT;
		}
	case LWS_UPNG_RGBA:
		switch (upng->color_depth) {
		case 8:
			return LWS_UPNG_RGBA8;
		case 16:
			return LWS_UPNG_RGBA16;
		default:
			return LWS_UPNG_BADFORMAT;
		}
	default:
		return LWS_UPNG_BADFORMAT;
	}
}

static const uint8_t magic[] = { 137, 80, 78, 71, 13, 10, 26, 10 };

static lws_stateful_ret_t
lws_upng_decode(lws_upng_t* u, const uint8_t **_pos, size_t *_size)
{
	const uint8_t *pos = _pos ? *_pos : NULL, *end = pos + *_size;
	lws_stateful_ret_t r = LWS_SRET_FATAL + 60;
	size_t m;

	if (u->of == UOF_INSIDE && !u->inf.in) {
		u->inf.inpos = 0;
		u->inf.in = pos;
		u->inf.bp = 0;
		m = lws_ptr_diff_size_t(end, pos);
		if (m > u->chunklen)
			m = u->chunklen;
		u->inf.inlen = m;
	}

	while (!u->no_more_input &&
	       ((u->of == UOF_INSIDE && _pos == NULL) || pos < end)) {
		switch (u->of) {
		case UOF_MAGIC:
			if (*pos++ != magic[u->sctr++])
				return LWS_SRET_FATAL + 17;
			if (u->sctr == sizeof(magic)) {
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_SKIP:
			pos++;
			if (++u->sctr == 4) {
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_TYPE4:
			u->acc = (u->acc << 8) | *pos++;
			if (++u->sctr == 4) {
				if (u->acc != LWS_FOURCC('I','H','D','R'))
					return LWS_SRET_FATAL + 18;
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_WIDTH4:
			u->acc = (u->acc << 8) | *pos++;
			if (++u->sctr == 4) {
				u->width = u->acc;
				if (!u->acc)
					return LWS_SRET_FATAL + 18;
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_HEIGHT4:
			u->acc = (u->acc << 8) | *pos++;
			if (++u->sctr == 4) {
				u->height = u->acc;
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_CDEPTH:
			u->color_depth =*pos++;
			u->of++;
			break;

		case UOF_CTYPE:
			u->color_type = *pos++;
			//lwsl_notice("w %d, h %d, depth %d, type %d\n",
			//		u->width, u->height,
			//		u->color_depth, u->color_type);
			u->format = determine_format(u);
			if (u->format == LWS_UPNG_BADFORMAT)
				return LWS_SRET_FATAL + 19;
			u->of++;
			break;

		case UOF_ONLY_ZERO3:
			if (*pos++)
				return LWS_SRET_FATAL + 20;
			if (++u->sctr == 3) {
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_SKIP4:
			pos++;
			if (++u->sctr != 4)
				break;

			/* takes us to +33 */

			memset(&u->inf, 0, sizeof(u->inf));

			/* 32KB gz sliding window */
			u->inf.info_size = 32768 + 512;
			u->u.bpp	 = lws_upng_get_bpp(u);
			if (!u->u.bpp || u->width >= (UINT_MAX / u->u.bpp))
				return LWS_SRET_FATAL + 14;

			u->u.y		= 0;
			u->u.ibp	= 0;
			u->u.bypp	= (u->u.bpp + 7) / 8;
			u->inf.bypl	= u->u.bypl = u->width * u->u.bypp;

			u->inf.outlen	= u->inf.info_size;
			u->inf.outpos	= 0;
			u->inf.state	= UPNS_ID_BL_GB_DONE;
			u->inf.upng	= u;

			u->u.alt	= 0; /* which of the two lines to write to */
			u->u.padded	= u->u.bpp < 8 &&
					  u->width * u->u.bpp !=
					  ((u->width * u->u.bpp + 7) / 8) * 8;
			u->u.diff	= (((u->width * u->u.bpp + 7) / 8) * 8) -
						(u->width * u->u.bpp);

			u->of++;
			u->sctr = 0;
			break;

		case UOF_CHUNK_LEN:
			if (!u->inf.out) {
				size_t ims = (u->u.bypl * 2) + u->inf.info_size;

				if (u->u.bypl > UINT_MAX / 2 || u->inf.info_size > UINT_MAX - (u->u.bypl * 2)) {
					lwsl_err("%s: integer overflow occur in ims %llu",
						 __func__, (unsigned long long)ims);
					return LWS_SRET_FATAL + 27;
				}

				if (u->hold_at_metadata)
					return LWS_SRET_AWAIT_RETRY;

				u->inf.out = (uint8_t *)lws_malloc(ims, __func__);
				if (!u->inf.out) {
					lwsl_notice("%s: inf malloc %u OOM\n",
						__func__, (unsigned int)ims);

					return LWS_SRET_YIELD;
				}
				u->u.lines = u->inf.out + u->inf.info_size;
				u->u.in		= u->inf.out;
			}
			u->chunklen = (u->chunklen << 8) | *pos++;
			if (++u->sctr == 4) {
				u->of++;
				u->sctr = 0;
			}
			break;

		case UOF_CHUNK_TYPE:
			u->ctype = (u->ctype << 8) | *pos++;
			if (++u->sctr != 4)
				break;
			u->sctr = 0;
			if (u->ctype == LWS_FOURCC('I','E','N','D')) {
				u->no_more_input = 1;
				break;
			}
			if (u->ctype != LWS_FOURCC('I','D','A','T')) {
				if (!(u->ctype & (32 << 24)))
					/* says it is critical... */
					 return LWS_SRET_FATAL + 27;

				u->chunklen += 4; /* chunk-end CRC */

				/* noncritical, skip */
				u->of = UOF_SKIP_CHUNK_LEN;
				break;
			}

			if (u->chunklen < 2)
				return LWS_SRET_FATAL + 31;

			/* it's a usable IDAT */

			if (!u->inf.subsequent)
				u->inf.inpos = 2;
			else
				u->inf.inpos = 0;

			m = lws_ptr_diff_size_t(end, pos);
			if (m > u->chunklen)
				m = u->chunklen;

			u->inf.in = pos;
			u->inf.inlen = m;
			u->inf.bp = 0;
			u->of++;
			break;

		case UOF_INSIDE:
			if (!u->inf.subsequent) {

				switch (u->sctr) {
				case 0:
					u->acc = (uint32_t)((*pos++) << 8);
					u->sctr++;
					continue;

				case 1:
					u->acc |= *pos++;

					if (u->acc % 31)
						return LWS_SRET_FATAL + 31;

					if (((u->acc >> 8) & 15) != 8 ||
					    ((u->acc >> 12) & 15) > 7)
						return LWS_SRET_FATAL + 31;

					if ((u->acc >> 5) & 1)
						return LWS_SRET_FATAL + 31;

					u->inf.subsequent = 1;
					break;
				}
			}

			r = _lws_upng_inflate_data(&u->inf);
			switch (r) {

			case LWS_SRET_WANT_INPUT:

				/* indicate no existing to drain */
				u->inf.in = NULL;

				pos += u->inf.inlen - u->inf.inpos;
				u->chunklen = u->chunklen -
						(unsigned int)(u->inf.inlen);

				if (!u->chunklen) {
					u->chunklen = 4; /* skip the 32-bit CRC */

					u->of = UOF_SKIP_CHUNK_LEN;
					break;
				}
				if (pos != end) {
					u->inf.inpos = 0;
					u->inf.in = pos;
					m = lws_ptr_diff_size_t(end, pos);
					if (m > u->chunklen)
						m = u->chunklen;
					u->inf.inlen = m;
					continue;
				}
				goto bail;
			default:
				goto bail;
			}
			break;

		case UOF_SKIP_CHUNK_LEN:
			pos++;
			if (!--u->chunklen) {
				u->of = UOF_CHUNK_LEN;
				u->sctr = 0;
				break;
			}
			break;
		}
	}

	r = LWS_SRET_OK;
	if (!u->no_more_input)
		r = LWS_SRET_WANT_INPUT;

bail:
	*_pos = pos;
	*_size = lws_ptr_diff_size_t(end, pos);

	return r;
}

lws_upng_t *
lws_upng_new(void)
{
	lws_upng_t* upng;

	upng = (lws_upng_t*)lws_zalloc(sizeof(lws_upng_t), __func__);
	if (upng == NULL)
		return NULL;

	upng->color_type = LWS_UPNG_RGBA;
	upng->color_depth = 8;
	upng->format = LWS_UPNG_RGBA8;

	upng->of = UOF_MAGIC;
	upng->sctr = 0;

	upng->inf.upng = upng;

	return upng;
}

void
lws_upng_free(lws_upng_t** upng)
{
	if ((*upng)->inf.out)
		lws_free_set_NULL((*upng)->inf.out);

	lws_free_set_NULL(*upng);
}


unsigned int
lws_upng_get_width(const lws_upng_t* upng)
{
	return upng->width;
}

unsigned int
lws_upng_get_height(const lws_upng_t* upng)
{
	return upng->height;
}

unsigned int
lws_upng_get_bpp(const lws_upng_t* upng)
{
	return lws_upng_get_bitdepth(upng) *
			lws_upng_get_components(upng);
}

unsigned int
lws_upng_get_components(const lws_upng_t* upng)
{
	switch (upng->color_type) {
	case LWS_UPNG_LUM:
		return 1;
	case LWS_UPNG_RGB:
		return 3;
	case LWS_UPNG_LUMA:
		return 2;
	case LWS_UPNG_RGBA:
		return 4;
	default:
		return 0;
	}
}

unsigned int
lws_upng_get_bitdepth(const lws_upng_t* upng)
{
	return upng->color_depth;
}

unsigned int
lws_upng_get_pixelsize(const lws_upng_t* upng)
{
	unsigned bits = lws_upng_get_bitdepth(upng) *
				lws_upng_get_components(upng);

	bits += bits % 8;

	return bits;
}

lws_upng_format_t
lws_upng_get_format(const lws_upng_t *upng)
{
	return upng->format;
}