#include "stdafx.h"
#include "Jpeg.h"
#include "Exception.h"
#include "Core/Convert.h"

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

namespace graphics {

	static Bool CODECALL jpegApplicable(IStream *from) {
		return checkHeader(from, "\xff\xd8", false)
			|| checkHeader(from, "JFIF", false)
			|| checkHeader(from, "Exif", false);
	}

	static FormatOptions *CODECALL jpegCreate(ImageFormat *f) {
		return new (f) JPEGOptions();
	}

	ImageFormat *jpegFormat(Engine &e) {
		const wchar *exts[] = {
			S("jpg"),
			S("jpeg"),
			null
		};
		return new (e) ImageFormat(S("Joint Photographic Experts Group"), exts, &jpegApplicable, &jpegCreate);
	}

	JPEGOptions::JPEGOptions() : quality(90) {}

	JPEGOptions::JPEGOptions(Nat quality) : quality(quality) {}

	void JPEGOptions::toS(StrBuf *to) const {
		*to << S("JPEG: { quality: ") << quality << S(" }");
	}

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

	struct JpegError : jpeg_error_mgr {
		Engine *e;
	};

	void onJpegError(j_common_ptr info) {
		struct JpegError *me = (struct JpegError *)info;

		char buffer[JMSG_LENGTH_MAX];
		(*me->format_message)(info, buffer);

		Str *msg = new (*me->e) Str(toWChar(*me->e, buffer));
		throw new (*me->e) ImageLoadError(msg);
	}

	class JpegInput : public jpeg_source_mgr {
	public:
		JpegInput(IStream *src) : src(src) {
			next_input_byte = NULL;
			bytes_in_buffer = 0;
			init_source = &JpegInput::sInit;
			fill_input_buffer = &JpegInput::sFill;
			skip_input_data = &JpegInput::sSkip;
			resync_to_restart = &jpeg_resync_to_restart;
			term_source = &JpegInput::sClose;
		}

		void init() {
			// Read in 8k chunks.
			current = buffer(src->engine(), 1024*8);
		}

		bool fill() {
			current.filled(0);
			current = src->read(current);

			bytes_in_buffer = current.filled();
			next_input_byte = current.dataPtr();

			return !current.empty();
		}

		void skip(long count) {
			while (size_t(count) > bytes_in_buffer) {
				count -= bytes_in_buffer;
				fill();
			}

			next_input_byte += count;
			bytes_in_buffer -= count;
		}

	private:
		// Source.
		IStream *src;

		// Current buffer.
		Buffer current;

		// Wrapper functions.
		static void sInit(j_decompress_ptr info) {
			((JpegInput *)info->src)->init();
		}
		static boolean sFill(j_decompress_ptr info) {
			return ((JpegInput *)info->src)->fill() ? TRUE : FALSE;
		}
		static void sSkip(j_decompress_ptr info, long count) {
			((JpegInput *)info->src)->skip(count);
		}
		static void sClose(j_decompress_ptr info) {
			// Usually nothing is required here. Called when reading an image is finished.
		}
	};

	Image *JPEGOptions::load(IStream *from) {
		struct jpeg_decompress_struct decode;
		JpegError errorMgr;
		errorMgr.e = &from->engine();
		decode.err = jpeg_std_error(&errorMgr);
		decode.err->error_exit = &onJpegError;

		try {
			jpeg_create_decompress(&decode);
			JpegInput input(from);
			decode.src = &input;

			if (jpeg_read_header(&decode, TRUE) != JPEG_HEADER_OK)
				throw new (from) ImageLoadError(S("No JPEG header was found."));

			// Set up output format and start decompression.
			decode.out_color_space = JCS_EXT_RGBA;
			jpeg_start_decompress(&decode);

			// Create and fill target bitmap.
			Image *out = new (from) Image(decode.output_width, decode.output_height);
			while (decode.output_scanline < decode.output_height) {
				byte *buf = out->buffer(0, decode.output_scanline);
				jpeg_read_scanlines(&decode, &buf, 1);
			}

			// Clean up.
			jpeg_finish_decompress(&decode);
			jpeg_destroy_decompress(&decode);
			return out;
		} catch (...) {
			// Clean up!
			jpeg_destroy_decompress(&decode);
			throw;
		}
	}

	// Note: We assume that this struct is stack-allocated, since we contain pointers into the
	// middle of GC-allocated data.
	class JpegOutput : public jpeg_destination_mgr {
	public:
		JpegOutput(OStream *to) : to(to), error(false) {
			next_output_byte = NULL;
			free_in_buffer = 0;

			init_destination = &JpegOutput::sInit;
			empty_output_buffer = &JpegOutput::sFlush;
			term_destination = &JpegOutput::sFinish;
		}

		// I/O error?
		Bool error;

	private:
		// Stream.
		OStream *to;

		// Output buffer.
		Buffer buffer;

		// Initialize the buffer.
		static void sInit(j_compress_ptr info) {
			JpegOutput *me = fromInfo(info);
			me->buffer = storm::buffer(me->to->engine(), 1024);
			me->next_output_byte = me->buffer.dataPtr();
			me->free_in_buffer = me->buffer.count();
		}

		// Flush the entire buffer.
		static boolean sFlush(j_compress_ptr info) {
			// Note: Documentation says to ignore the current value of 'free_in_buffer' and just
			// flush the entire buffer.
			JpegOutput *me = fromInfo(info);
			me->buffer.filled(me->buffer.count());
			Nat written = me->to->write(me->buffer);
			if (written != me->buffer.count())
				me->error = true;

			me->buffer.filled(0);
			me->next_output_byte = me->buffer.dataPtr();
			me->free_in_buffer = me->buffer.count();
			return TRUE;
		}

		// Close. Flush any remaining data in the buffer.
		static void sFinish(j_compress_ptr info) {
			// Flush remaining data.
			JpegOutput *me = fromInfo(info);
			me->buffer.filled(me->buffer.count() - me->free_in_buffer);
			if (me->to->write(me->buffer) != me->buffer.filled())
				me->error = true;

			// Ask it to flush to the filesystem as well.
			me->to->flush();

			// Clear the buffer.
			me->buffer = Buffer();
		}

		// Get a pointer to this object.
		static JpegOutput *fromInfo(j_compress_ptr info) {
			return (JpegOutput *)info->dest;
		}
	};

	void JPEGOptions::save(Image *image, OStream *to) {
		struct jpeg_compress_struct encode;
		JpegError errorMgr;
		errorMgr.e = &image->engine();
		encode.err = jpeg_std_error(&errorMgr);
		encode.err->error_exit = &onJpegError;

		try {
			jpeg_create_compress(&encode);
			JpegOutput output(to);
			encode.dest = &output;

			encode.image_width = image->width();
			encode.image_height = image->height();
			encode.input_components = 4; // RGBA, so we don't have to convert ourselves.
			encode.in_color_space = JCS_EXT_RGBA;
			encode.data_precision = 8; // 8 bits per pixel?

			jpeg_set_defaults(&encode);

			jpeg_set_quality(&encode, int(quality), TRUE); // TRUE = limit to baseline JPEG values.

			jpeg_start_compress(&encode, TRUE);

			// Output one scanline at a time.
			while (encode.next_scanline < encode.image_height) {
				byte *scanline = image->buffer(0, encode.next_scanline);
				jpeg_write_scanlines(&encode, &scanline, 1);
			}

			// Finish up.
			jpeg_finish_compress(&encode);

			if (output.error)
				throw new (this) ImageSaveError(TO_S(to, S("I/O error while saving image: ") << to->error()));

			jpeg_destroy_compress(&encode);
		} catch (...) {
			jpeg_destroy_compress(&encode);
			throw;
		}
	}

#endif
}