#include "stdafx.h"
#include "Png.h"
#include "Exception.h"

#ifndef WINDOWS
#include <png.h>
#endif

namespace graphics {

	static Bool CODECALL pngApplicable(IStream *from) {
		return checkHeader(from, "\x89PNG", false);
	}

	static FormatOptions *CODECALL pngCreate(ImageFormat *f) {
		return new (f) PNGOptions();
	}

	ImageFormat *pngFormat(Engine &e) {
		const wchar *exts[] = { S("png"), null };
		return new (e) ImageFormat(S("Portable Network Graphics"), exts, &pngApplicable, &pngCreate);
	}

	PNGOptions::PNGOptions() {}

	// See ImageWin32.cpp for loading/saving on Windows.
#ifndef WINDOWS

	// IO callbacks from libpng.
	static void pngRead(png_structp png, byte *to, size_t length) {
		IStream *src = (IStream*)png_get_io_ptr(png);
		Buffer out = src->fill(length);
		if (out.count() > 0)
			memcpy(to, out.dataPtr(), out.count());
	}

	static Image *loadPng(PNGOptions *options, IStream *from, const wchar *&error) {
		// Anchor 'from' on the stack so that the GC does not move it.
		IStream *volatile anchor = null;
		atomicWrite(anchor, from);

		png_structp png = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
		png_infop info = null;
		error = S("Failed to initialize libpng");

		// Custom error handling:
		if (setjmp(png_jmpbuf(png))) {
			// Error. Clean up and exit.
			png_destroy_read_struct(&png, &info, NULL);
			error = S("Internal error in libpng.");
			return null;
		}

		png_set_read_fn(png, (void *)anchor, &pngRead);

		if (png) {
			info = png_create_info_struct(png);
			error = S("Failed to create info struct.");
		}

		png_uint_32 ok = 0;
		png_uint_32 width = 0, height = 0;
		int depth = 0, type = 0;
		if (info) {
			png_read_info(png, info);
			ok = png_get_IHDR(png, info, &width, &height, &depth, &type, NULL, NULL, NULL);
			error = S("Failed to read PNG header.");
		}

		Image *output = null;
		if (ok) {
			int typeNoAlpha = type & ~PNG_COLOR_MASK_ALPHA;

			// Fix the format.
			if (typeNoAlpha == PNG_COLOR_TYPE_PALETTE)
				png_set_palette_to_rgb(png);
			if (typeNoAlpha == PNG_COLOR_TYPE_GRAY)
				png_set_gray_to_rgb(png);
			if (png_get_valid(png, info, PNG_INFO_tRNS))
				png_set_tRNS_to_alpha(png);
			if (depth == 16)
				png_set_strip_16(png);
			if (depth < 8) {
				png_color_8p bit;
				if (png_get_sBIT(png, info, &bit))
					png_set_shift(png, bit);
			}
			if (!(type & PNG_COLOR_MASK_ALPHA))
				png_set_add_alpha(png, 0xFF, PNG_FILLER_AFTER);

			png_read_update_info(png, info);

			// Read the PNG file itself!
			output = new (from) Image(Nat(width), Nat(height));

			// Make sure to pin the internal array by storing a pointer to it on the stack.
			byte *volatile buffer = null;
			atomicWrite(buffer, output->buffer());
			Nat stride = output->stride();

			// Create a GcArray where we can store pointers to inside 'buffer'. Note: it is not
			// scanned, as that would lead to pointers to the middle of objects. Since the buffer is
			// pinned on the stack, it can not move so this should be fine anyway.
			GcArray<byte *> *rows = runtime::allocArray<byte *>(output->engine(), &sizeArrayType, height);
			for (Nat i = 0; i < height; i++)
				rows->v[i] = buffer + stride*i;

			png_read_image(png, rows->v);
		}

		// Clean up any structures we allocated.
		png_destroy_read_struct(&png, &info, NULL);

		return output;
	}

	Image *PNGOptions::load(IStream *from) {
		const wchar *error;
		Image *out = loadPng(this, from, error);
		if (!out)
			throw new (this) ImageLoadError(error);
		return out;
	}

	static void pngWrite(png_structp png, byte *from, size_t length) {
		OStream *dest = (OStream *)png_get_io_ptr(png);
		dest->write(buffer(dest->engine(), from, length));
	}

	static void pngFlush(png_structp png) {
		OStream *dest = (OStream *)png_get_io_ptr(png);
		dest->flush();
	}

	void PNGOptions::save(Image *image, OStream *to) {
		OStream *volatile anchor = null;
		atomicWrite(anchor, to);

		const wchar *error = S("Failed to initialize libpng.");
		png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
		png_infop info = null;
		if (!png) {
			throw new (image) ImageSaveError(error);
		}

		if (setjmp(png_jmpbuf(png))) {
			// Error. Clean up and exit.
			png_destroy_write_struct(&png, &info);
			throw new (image) ImageSaveError(S("Internal error in libpng."));
		}

		info = png_create_info_struct(png);
		if (!info)
			throw new (image) ImageSaveError(S("Failed to create png info struct."));

		png_set_write_fn(png, (void *)anchor, &pngWrite, &pngFlush);

		png_set_IHDR(png, info, image->width(), image->height(), 8,
					PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
		png_write_info(png, info);

		// Make sure to pin the internal array by storing a pointer to it on the stack.
		byte *volatile buffer = null;
		atomicWrite(buffer, image->buffer());
		Nat stride = image->stride();

		// Create a GcArray where we can store pointers to inside 'buffer'. Note: it is not
		// scanned, as that would lead to pointers to the middle of objects. Since the buffer is
		// pinned on the stack, it can not move so this should be fine anyway.
		GcArray<byte *> *rows = runtime::allocArray<byte *>(image->engine(), &sizeArrayType, image->height());
		for (Nat i = 0; i < image->height(); i++)
			rows->v[i] = buffer + stride*i;

		png_write_image(png, rows->v);
		png_write_end(png, NULL);

		png_destroy_write_struct(&png, &info);
	}

#endif

}
