/* BSD-2-Clause license
 * 
 * Copyright (c) 2018-2023 NST <www.newinfosec.ru>, sss <sss at dark-alexandr dot net>. 
 * 
 */

#include "rdp_png.h"
#include <stdlib.h>
#include <string.h>

static void
cbPngError(png_structp png, png_const_charp msg)
{
	printf("Png error: %s\n", (msg ? msg : "unknown error"));
}

static void
cbPngWarn(png_structp png, png_const_charp msg)
{
	printf("Png warning: %s\n", (msg ? msg : "unknown error"));
}

static void
cbPngWrite(png_structp png, png_bytep data, png_size_t len)
{
	rdp_png_buf *out_buf = (rdp_png_buf *)png_get_io_ptr(png);
	uint8_t *ptr = out_buf->buf;
	if (out_buf->written + len > out_buf->buf_size)
	{
		printf("error: rdp_png_buf is too small\n");
		return;
	}
	ptr += out_buf->written;
	memcpy(ptr, data, len);
	out_buf->written += len;
}

static void
cbPngFlush(png_structp png)
{
}

bool
png_generate_from_argb(
    int width, int height, uint8_t *data, rdp_png_buf *out_buf)
{
	png_infop info_ptr;
	png_structp png_ptr = png_create_write_struct(
	    PNG_LIBPNG_VER_STRING, 0, cbPngError, cbPngWarn);
	if (!png_ptr)
	{
		printf("Could not allocate png_struct\n");
		return false;
	}
	info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
		printf("Could not allocate png_info_struct\n");
		return false;
	}
	png_set_write_fn(png_ptr, out_buf, cbPngWrite, cbPngFlush);
	png_set_IHDR(png_ptr, info_ptr, width, height, 8,
	    PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE,
	    PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
	png_bytep *rows = malloc(sizeof(png_bytep) * height);
	if (!rows)
	{
		perror("malloc");
		return false;
	}
	for (int i = 0; i < height; ++i)
	{
		rows[i] = (png_bytep)data;
		data += width * 4;
	}
	png_set_rows(png_ptr, info_ptr, rows);
	png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

	free(rows);

	return true;
}

void
png_destroy(png_structp *p, png_infop *i)
{
	png_destroy_write_struct(p, i);
}