// ==========================================================
// Google WebP Loader & Writer
//
// Design and implementation by
// - Herve Drolon (drolon@infonie.fr)
//
// This file is part of FreeImage 3
//
// COVERED CODE IS PROVIDED UNDER THIS LICENSE ON AN "AS IS" BASIS, WITHOUT WARRANTY
// OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, WITHOUT LIMITATION, WARRANTIES
// THAT THE COVERED CODE IS FREE OF DEFECTS, MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE
// OR NON-INFRINGING. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE COVERED
// CODE IS WITH YOU. SHOULD ANY COVERED CODE PROVE DEFECTIVE IN ANY RESPECT, YOU (NOT
// THE INITIAL DEVELOPER OR ANY OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY
// SERVICING, REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL
// PART OF THIS LICENSE. NO USE OF ANY COVERED CODE IS AUTHORIZED HEREUNDER EXCEPT UNDER
// THIS DISCLAIMER.
//
// Use at your own risk!
// ==========================================================

#include "FreeImage.h"
#include "Utilities.h"

#include "../Metadata/FreeImageTag.h"

#include "webp/decode.h"
#include "webp/encode.h"
#include "webp/mux.h"

// ==========================================================
// Plugin Interface
// ==========================================================

static int s_format_id;

// ----------------------------------------------------------
//   Helpers for the load function
// ----------------------------------------------------------

/**
Read the whole file into memory
*/
static BOOL
ReadFileToWebPData(FreeImageIO *io, fi_handle handle, WebPData * const bitstream) {
  uint8_t *raw_data = NULL;

  try {
	  // Read the input file and put it in memory
	  long start_pos = io->tell_proc(handle);
	  io->seek_proc(handle, 0, SEEK_END);
	  size_t file_length = (size_t)(io->tell_proc(handle) - start_pos);
	  io->seek_proc(handle, start_pos, SEEK_SET);
	  raw_data = (uint8_t*)malloc(file_length * sizeof(uint8_t));
	  if(!raw_data) {
		  throw FI_MSG_ERROR_MEMORY;
	  }
	  if(io->read_proc(raw_data, 1, (unsigned)file_length, handle) != file_length) {
		  throw "Error while reading input stream";
	  }
	  
	  // copy pointers (must be released later using free)
	  bitstream->bytes = raw_data;
	  bitstream->size = file_length;

	  return TRUE;

  } catch(const char *text) {
	  if(raw_data) {
		  free(raw_data);
	  }
	  memset(bitstream, 0, sizeof(WebPData));
	  if(NULL != text) {
		  FreeImage_OutputMessageProc(s_format_id, text);
	  }
	  return FALSE;
  }
}

// ----------------------------------------------------------
//   Helpers for the save function
// ----------------------------------------------------------

/**
Output function. Should return 1 if writing was successful.
data/data_size is the segment of data to write, and 'picture' is for
reference (and so one can make use of picture->custom_ptr).
*/
static int 
WebP_MemoryWriter(const BYTE *data, size_t data_size, const WebPPicture* const picture) {
	FIMEMORY *hmem = (FIMEMORY*)picture->custom_ptr;
	return data_size ? (FreeImage_WriteMemory(data, 1, (unsigned)data_size, hmem) == data_size) : 0;
}

// ==========================================================
// Plugin Implementation
// ==========================================================

static const char * DLL_CALLCONV
Format() {
	return "WebP";
}

static const char * DLL_CALLCONV
Description() {
	return "Google WebP image format";
}

static const char * DLL_CALLCONV
Extension() {
	return "webp";
}

static const char * DLL_CALLCONV
RegExpr() {
	return NULL;
}

static const char * DLL_CALLCONV
MimeType() {
	return "image/webp";
}

static BOOL DLL_CALLCONV
Validate(FreeImageIO *io, fi_handle handle) {
	BYTE riff_signature[4] = { 0x52, 0x49, 0x46, 0x46 };
	BYTE webp_signature[4] = { 0x57, 0x45, 0x42, 0x50 };
	BYTE signature[12] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

	io->read_proc(signature, 1, 12, handle);

	if(memcmp(riff_signature, signature, 4) == 0) {
		if(memcmp(webp_signature, signature + 8, 4) == 0) {
			return TRUE;
		}
	}

	return FALSE;
}

static BOOL DLL_CALLCONV
SupportsExportDepth(int depth) {
	return (
		(depth == 24) || 
		(depth == 32)
		);
}

static BOOL DLL_CALLCONV 
SupportsExportType(FREE_IMAGE_TYPE type) {
	return (type == FIT_BITMAP) ? TRUE : FALSE;
}

static BOOL DLL_CALLCONV
SupportsICCProfiles() {
	return TRUE;
}

static BOOL DLL_CALLCONV
SupportsNoPixels() {
	return TRUE;
}

// ----------------------------------------------------------

static void * DLL_CALLCONV
Open(FreeImageIO *io, fi_handle handle, BOOL read) {
	WebPMux *mux = NULL;
	int copy_data = 1;	// 1 : copy data into the mux, 0 : keep a link to local data

	if(read) {
		// create the MUX object from the input stream
		WebPData bitstream;
		// read the input file and put it in memory
		if(!ReadFileToWebPData(io, handle, &bitstream)) {
			return NULL;
		}
		// create the MUX object
		mux = WebPMuxCreate(&bitstream, copy_data);
		// no longer needed since copy_data == 1
		free((void*)bitstream.bytes);
		if(mux == NULL) {
			FreeImage_OutputMessageProc(s_format_id, "Failed to create mux object from file");
			return NULL;
		}
	} else {
		// creates an empty mux object
		mux = WebPMuxNew();
		if(mux == NULL) {
			FreeImage_OutputMessageProc(s_format_id, "Failed to create empty mux object");
			return NULL;
		}
	}
	
	return mux;
}

static void DLL_CALLCONV
Close(FreeImageIO *io, fi_handle handle, void *data) {
	WebPMux *mux = (WebPMux*)data;
	if(mux != NULL) {
		// free the MUX object
		WebPMuxDelete(mux);
	}
}

// ----------------------------------------------------------

/**
Decode a WebP image and returns a FIBITMAP image
@param webp_image Raw WebP image
@param flags FreeImage load flags
@return Returns a dib if successfull, returns NULL otherwise
*/
static FIBITMAP *
DecodeImage(WebPData *webp_image, int flags) {
	FIBITMAP *dib = NULL;

	const uint8_t* data = webp_image->bytes;	// raw image data
	const size_t data_size = webp_image->size;	// raw image size

    VP8StatusCode webp_status = VP8_STATUS_OK;

	BOOL header_only = (flags & FIF_LOAD_NOPIXELS) == FIF_LOAD_NOPIXELS;

	// Main object storing the configuration for advanced decoding
	WebPDecoderConfig decoder_config;
	// Output buffer
	WebPDecBuffer* const output_buffer = &decoder_config.output;
	// Features gathered from the bitstream
	WebPBitstreamFeatures* const bitstream = &decoder_config.input;

	try {
		// Initialize the configuration as empty
		// This function must always be called first, unless WebPGetFeatures() is to be called
		if(!WebPInitDecoderConfig(&decoder_config)) {
			throw "Library version mismatch";
		}

		// Retrieve features from the bitstream
		// The bitstream structure is filled with information gathered from the bitstream
		webp_status = WebPGetFeatures(data, data_size, bitstream);
		if(webp_status != VP8_STATUS_OK) {
			throw FI_MSG_ERROR_PARSING;
		}

		// Allocate output dib

		unsigned bpp = bitstream->has_alpha ? 32 : 24;	
		unsigned width = (unsigned)bitstream->width;
		unsigned height = (unsigned)bitstream->height;

		dib = FreeImage_AllocateHeader(header_only, width, height, bpp, FI_RGBA_RED_MASK, FI_RGBA_GREEN_MASK, FI_RGBA_BLUE_MASK);
		if(!dib) {
			throw FI_MSG_ERROR_DIB_MEMORY;
		}

		if(header_only) {
			WebPFreeDecBuffer(output_buffer);
			return dib;
		}

		// --- Set decoding options ---

		// use multi-threaded decoding
		decoder_config.options.use_threads = 1;
		// set output color space
		output_buffer->colorspace = bitstream->has_alpha ? MODE_BGRA : MODE_BGR;

		// ---

		// decode the input stream, taking 'config' into account. 
		
		webp_status = WebPDecode(data, data_size, &decoder_config);
		if(webp_status != VP8_STATUS_OK) {
			throw FI_MSG_ERROR_PARSING;
		}

		// fill the dib with the decoded data

		const BYTE *src_bitmap = output_buffer->u.RGBA.rgba;
		const unsigned src_pitch = (unsigned)output_buffer->u.RGBA.stride;

		switch(bpp) {
			case 24:
				for(unsigned y = 0; y < height; y++) {
					const BYTE *src_bits = src_bitmap + y * src_pitch;						
					BYTE *dst_bits = (BYTE*)FreeImage_GetScanLine(dib, height-1-y);
					for(unsigned x = 0; x < width; x++) {
						dst_bits[FI_RGBA_BLUE]	= src_bits[0];	// B
						dst_bits[FI_RGBA_GREEN]	= src_bits[1];	// G
						dst_bits[FI_RGBA_RED]	= src_bits[2];	// R
						src_bits += 3;
						dst_bits += 3;
					}
				}
				break;
			case 32:
				for(unsigned y = 0; y < height; y++) {
					const BYTE *src_bits = src_bitmap + y * src_pitch;						
					BYTE *dst_bits = (BYTE*)FreeImage_GetScanLine(dib, height-1-y);
					for(unsigned x = 0; x < width; x++) {
						dst_bits[FI_RGBA_BLUE]	= src_bits[0];	// B
						dst_bits[FI_RGBA_GREEN]	= src_bits[1];	// G
						dst_bits[FI_RGBA_RED]	= src_bits[2];	// R
						dst_bits[FI_RGBA_ALPHA]	= src_bits[3];	// A
						src_bits += 4;
						dst_bits += 4;
					}
				}
				break;
		}

		// Free the decoder
		WebPFreeDecBuffer(output_buffer);

		return dib;

	} catch (const char *text) {
		if(dib) {
			FreeImage_Unload(dib);
		}
		WebPFreeDecBuffer(output_buffer);

		if(NULL != text) {
			FreeImage_OutputMessageProc(s_format_id, text);
		}

		return NULL;
	}
}

static FIBITMAP * DLL_CALLCONV
Load(FreeImageIO *io, fi_handle handle, int page, int flags, void *data) {
	WebPMux *mux = NULL;
	WebPMuxFrameInfo webp_frame = { 0 };	// raw image
	WebPData color_profile;	// ICC raw data
	WebPData xmp_metadata;	// XMP raw data
	WebPData exif_metadata;	// EXIF raw data
	FIBITMAP *dib = NULL;
	WebPMuxError error_status;

	if(!handle) {
		return NULL;
	}

	try {
		// get the MUX object
		mux = (WebPMux*)data;
		if(!mux) {
			throw (1);
		}
		
		// gets the feature flags from the mux object
		uint32_t webp_flags = 0;
		error_status = WebPMuxGetFeatures(mux, &webp_flags);
		if(error_status != WEBP_MUX_OK) {
			throw (1);
		}

		// get image data
		error_status = WebPMuxGetFrame(mux, 1, &webp_frame);

		if(error_status == WEBP_MUX_OK) {
			// decode the data (can be limited to the header if flags uses FIF_LOAD_NOPIXELS)
			dib = DecodeImage(&webp_frame.bitstream, flags);
			if(!dib) {
				throw (1);
			}
			
			// get ICC profile
			if(webp_flags & ICCP_FLAG) {
				error_status = WebPMuxGetChunk(mux, "ICCP", &color_profile);
				if(error_status == WEBP_MUX_OK) {
					FreeImage_CreateICCProfile(dib, (void*)color_profile.bytes, (long)color_profile.size);
				}
			}

			// get XMP metadata
			if(webp_flags & XMP_FLAG) {
				error_status = WebPMuxGetChunk(mux, "XMP ", &xmp_metadata);
				if(error_status == WEBP_MUX_OK) {
					// create a tag
					FITAG *tag = FreeImage_CreateTag();
					if(tag) {
						FreeImage_SetTagKey(tag, g_TagLib_XMPFieldName);
						FreeImage_SetTagLength(tag, (DWORD)xmp_metadata.size);
						FreeImage_SetTagCount(tag, (DWORD)xmp_metadata.size);
						FreeImage_SetTagType(tag, FIDT_ASCII);
						FreeImage_SetTagValue(tag, xmp_metadata.bytes);
						
						// store the tag
						FreeImage_SetMetadata(FIMD_XMP, dib, FreeImage_GetTagKey(tag), tag);

						// destroy the tag
						FreeImage_DeleteTag(tag);
					}
				}
			}

			// get Exif metadata
			if(webp_flags & EXIF_FLAG) {
				error_status = WebPMuxGetChunk(mux, "EXIF", &exif_metadata);
				if(error_status == WEBP_MUX_OK) {
					// read the Exif raw data as a blob
					jpeg_read_exif_profile_raw(dib, exif_metadata.bytes, (unsigned)exif_metadata.size);
					// read and decode the Exif data
					jpeg_read_exif_profile(dib, exif_metadata.bytes, (unsigned)exif_metadata.size);
				}
			}
		}

		WebPDataClear(&webp_frame.bitstream);

		return dib;

	} catch(int) {
		WebPDataClear(&webp_frame.bitstream);
		return NULL;
	}
}

// --------------------------------------------------------------------------

/**
Encode a FIBITMAP to a WebP image
@param hmem Memory output stream, containing on return the encoded image
@param dib The FIBITMAP to encode
@param flags FreeImage save flags
@return Returns TRUE if successfull, returns FALSE otherwise
*/
static BOOL
EncodeImage(FIMEMORY *hmem, FIBITMAP *dib, int flags) {
	WebPPicture picture;	// Input buffer
	WebPConfig config;		// Coding parameters

	BOOL bIsFlipped = FALSE;

	try {
		const unsigned width = FreeImage_GetWidth(dib);
		const unsigned height = FreeImage_GetHeight(dib);
		const unsigned bpp = FreeImage_GetBPP(dib);
		const unsigned pitch = FreeImage_GetPitch(dib);

		// check image type
		FREE_IMAGE_TYPE image_type = FreeImage_GetImageType(dib);

		if( !((image_type == FIT_BITMAP) && ((bpp == 24) || (bpp == 32))) )  {
			throw FI_MSG_ERROR_UNSUPPORTED_FORMAT;
		}

		// check format limits
		if(MAX(width, height) > WEBP_MAX_DIMENSION) {
			FreeImage_OutputMessageProc(s_format_id, "Unsupported image size: width x height = %d x %d", width, height);
			return FALSE;
		}

		// Initialize output I/O
		if(WebPPictureInit(&picture) == 1) {
			picture.writer = WebP_MemoryWriter;
			picture.custom_ptr = hmem;
			picture.width = (int)width;
			picture.height = (int)height;
		} else {
			throw "Couldn't initialize WebPPicture";
		}

		// --- Set encoding parameters ---

		// Initialize encoding parameters to default values
		WebPConfigInit(&config);

		// quality/speed trade-off (0=fast, 6=slower-better)
		config.method = 6;

		if((flags & WEBP_LOSSLESS) == WEBP_LOSSLESS) {
			// lossless encoding
			config.lossless = 1;
			picture.use_argb = 1;
		} else if((flags & 0x7F) > 0) {
			// lossy encoding
			config.lossless = 0;
			// quality is between 1 (smallest file) and 100 (biggest) - default to 75
			config.quality = (float)(flags & 0x7F);
			if(config.quality > 100) {
				config.quality = 100;
			}
		}

		// validate encoding parameters
		if(WebPValidateConfig(&config) == 0) {
			throw "Failed to initialize encoder";
		}

		// --- Perform encoding ---
		
		// Invert dib scanlines
		bIsFlipped = FreeImage_FlipVertical(dib);


		// convert dib buffer to output stream

		const BYTE *bits = FreeImage_GetBits(dib);

#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_BGR
		switch(bpp) {
			case 24:
				WebPPictureImportBGR(&picture, bits, pitch);
				break;
			case 32:
				WebPPictureImportBGRA(&picture, bits, pitch);
				break;
		}
#else
		switch(bpp) {
			case 24:
				WebPPictureImportRGB(&picture, bits, pitch);
				break;
			case 32:
				WebPPictureImportRGBA(&picture, bits, pitch);
				break;
		}

#endif // FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_BGR

		if(!WebPEncode(&config, &picture)) {
			throw "Failed to encode image";
		}

		WebPPictureFree(&picture);

		if(bIsFlipped) {
			// invert dib scanlines
			FreeImage_FlipVertical(dib);
		}

		return TRUE;

	} catch (const char* text) {

		WebPPictureFree(&picture);

		if(bIsFlipped) {
			// invert dib scanlines
			FreeImage_FlipVertical(dib);
		}

		if(NULL != text) {
			FreeImage_OutputMessageProc(s_format_id, text);
		}
	}

	return FALSE;
}

static BOOL DLL_CALLCONV
Save(FreeImageIO *io, FIBITMAP *dib, fi_handle handle, int page, int flags, void *data) {
	WebPMux *mux = NULL;
	FIMEMORY *hmem = NULL;
	WebPData webp_image;
	WebPData output_data = { 0 };
	WebPMuxError error_status;

	int copy_data = 1;	// 1 : copy data into the mux, 0 : keep a link to local data

	if(!dib || !handle || !data) {
		return FALSE;
	}

	try {

		// get the MUX object
		mux = (WebPMux*)data;
		if(!mux) {
			return FALSE;
		}

		// --- prepare image data ---

		// encode image as a WebP blob
		hmem = FreeImage_OpenMemory();
		if(!hmem || !EncodeImage(hmem, dib, flags)) {
			throw (1);
		}
		// store the blob into the mux
		BYTE *data = NULL;
		DWORD data_size = 0;
		FreeImage_AcquireMemory(hmem, &data, &data_size);
		webp_image.bytes = data;
		webp_image.size = data_size;
		error_status = WebPMuxSetImage(mux, &webp_image, copy_data);
		// no longer needed since copy_data == 1
		FreeImage_CloseMemory(hmem);
		hmem = NULL;
		if(error_status != WEBP_MUX_OK) {
			throw (1);
		}

		// --- set metadata ---
		
		// set ICC color profile
		{
			FIICCPROFILE *iccProfile = FreeImage_GetICCProfile(dib);
			if (iccProfile->size && iccProfile->data) {
				WebPData icc_profile;
				icc_profile.bytes = (uint8_t*)iccProfile->data;
				icc_profile.size = (size_t)iccProfile->size;
				error_status = WebPMuxSetChunk(mux, "ICCP", &icc_profile, copy_data);
				if(error_status != WEBP_MUX_OK) {
					throw (1);
				}
			}
		}

		// set XMP metadata
		{
			FITAG *tag = NULL;
			if(FreeImage_GetMetadata(FIMD_XMP, dib, g_TagLib_XMPFieldName, &tag)) {
				WebPData xmp_profile;
				xmp_profile.bytes = (uint8_t*)FreeImage_GetTagValue(tag);
				xmp_profile.size = (size_t)FreeImage_GetTagLength(tag);
				error_status = WebPMuxSetChunk(mux, "XMP ", &xmp_profile, copy_data);
				if(error_status != WEBP_MUX_OK) {
					throw (1);
				}
			}
		}

		// set Exif metadata
		{
			FITAG *tag = NULL;
			if(FreeImage_GetMetadata(FIMD_EXIF_RAW, dib, g_TagLib_ExifRawFieldName, &tag)) {
				WebPData exif_profile;
				exif_profile.bytes = (uint8_t*)FreeImage_GetTagValue(tag);
				exif_profile.size = (size_t)FreeImage_GetTagLength(tag);
				error_status = WebPMuxSetChunk(mux, "EXIF", &exif_profile, copy_data);
				if(error_status != WEBP_MUX_OK) {
					throw (1);
				}
			}
		}
		
		// get data from mux in WebP RIFF format
		error_status = WebPMuxAssemble(mux, &output_data);
		if(error_status != WEBP_MUX_OK) {
			FreeImage_OutputMessageProc(s_format_id, "Failed to create webp output file");
			throw (1);
		}

		// write the file to the output stream
		if(io->write_proc((void*)output_data.bytes, 1, (unsigned)output_data.size, handle) != output_data.size) {
			FreeImage_OutputMessageProc(s_format_id, "Failed to write webp output file");
			throw (1);
		}

		// free WebP output file
		WebPDataClear(&output_data);

		return TRUE;

	} catch(int) {
		if(hmem) {
			FreeImage_CloseMemory(hmem);
		}
		
		WebPDataClear(&output_data);

		return FALSE;
	}
}

// ==========================================================
//	 Init
// ==========================================================

void DLL_CALLCONV
InitWEBP(Plugin *plugin, int format_id) {
	s_format_id = format_id;

	plugin->format_proc = Format;
	plugin->description_proc = Description;
	plugin->extension_proc = Extension;
	plugin->regexpr_proc = RegExpr;
	plugin->open_proc = Open;
	plugin->close_proc = Close;
	plugin->pagecount_proc = NULL;
	plugin->pagecapability_proc = NULL;
	plugin->load_proc = Load;
	plugin->save_proc = Save;
	plugin->validate_proc = Validate;
	plugin->mime_proc = MimeType;
	plugin->supports_export_bpp_proc = SupportsExportDepth;
	plugin->supports_export_type_proc = SupportsExportType;
	plugin->supports_icc_profiles_proc = SupportsICCProfiles;
	plugin->supports_no_pixels_proc = SupportsNoPixels;
}