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

#ifdef WINDOWS

#include "StreamWrapper.h"
#include <wincodec.h>
#pragma comment (lib, "Windowscodecs.lib")

namespace graphics {

	namespace wic {

		// Property for IPropBag2.
		class Property {
		public:
			const wchar *property;
			VARIANT value;

			Property() : property(null) {}

			Property(const wchar *str, float value) : property(str) {
				// See WTypes.idl for valid types.
				this->value.vt = VT_R4;
				// See OAIdl.h for valid enum members.
				this->value.fltVal = value;
			}

			Bool empty() {
				return property == null;
			}
		};

		// Swap R and B channels in an Image.
		static void swapRB(Image *image) {
			for (Nat y = 0; y < image->height(); y++) {
				for (Nat x = 0; x < image->width(); x++) {
					byte *px = image->buffer(x, y);
					swap(px[0], px[2]);
				}
			}
		}

		HRESULT convertTexture(IWICImagingFactory *factory, IWICBitmapSource *src, IWICBitmapSource **dest, const wchar *&error) {
			IWICFormatConverter *converter = null;
			HRESULT r = factory->CreateFormatConverter(&converter);
			const wchar_t *err = S("Failed to create a format converter");

			if (SUCCEEDED(r)) {
				r = converter->Initialize(src, GUID_WICPixelFormat32bppBGRA, WICBitmapDitherTypeNone, NULL, 0.0f, WICBitmapPaletteTypeCustom);
				err = S("Failed to convert the picture");
			}

			if (SUCCEEDED(r)) {
				r = converter->QueryInterface(IID_PPV_ARGS(dest));
				err = S("Failed to get the converted picture");
			}

			release(converter);

			if (FAILED(r))
				error = err;

			return r;
		}

		Image *createTexture(Engine &e, IWICImagingFactory *factory, IWICBitmapSource *from, const wchar *&error) {
			IWICBitmapSource *converted = null;
			HRESULT r = convertTexture(factory, from, &converted, error);
			if (FAILED(r))
				return null;

			//NOTE: From here it is assumed that the format of "converted" is 32bpp, BGRA

			nat width, height;
			converted->GetSize(&width, &height);

			Image *image = new (e) Image(width, height);
			r = converted->CopyPixels(NULL, image->stride(), image->bufferSize(), image->buffer());

			// Swap red and blue... We could convert to RGBA, but that symbol is not present in the
			// shared library for some reason...
			swapRB(image);

			if (FAILED(r))
				error = S("Failed to copy bitmap data");

			release(converted);
			return image;
		}

		Image *loadImage(storm::IStream *from, const wchar *&err) {
			IWICImagingFactory *wicFactory = null;
			HRESULT r = CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&wicFactory));
			err = S("Failed to initialize Windows Imaging Components");

			IWICStream *wicStream = null;
			storm::RIStream *src = from->randomAccess();
			IStreamWrapper *fromStream = IStreamWrapper::create(&src);
			if (SUCCEEDED(r)) {
				r = wicFactory->CreateStream(&wicStream);
				err = S("Failed to create a WIC stream");
			}

			if (SUCCEEDED(r)) {
				r = wicStream->InitializeFromIStream(fromStream);
				err = S("Failed to initialize the WIC stream");
			}

			IWICBitmapDecoder *decoder = null;
			if (SUCCEEDED(r)) {
				r = wicFactory->CreateDecoderFromStream(wicStream, NULL, WICDecodeMetadataCacheOnLoad, &decoder);
				err = S("Failed to create a WIC decoder (possibly corrupt image)");
			}

			IWICBitmapFrameDecode *frame = null;
			if (SUCCEEDED(r)) {
				r = decoder->GetFrame(0, &frame);
				err = S("Failed to get bitmap data");
			}

			IWICBitmapSource *bitmapSource = null;
			if (SUCCEEDED(r)) {
				r = frame->QueryInterface(IID_PPV_ARGS(&bitmapSource));
				err = S("The object was not a BitmapSource object");
			}

			Image *result = null;
			// Convert the data and add it to the texture
			if (SUCCEEDED(r)) {
				result = createTexture(from->engine(), wicFactory, bitmapSource, err);
			}

			release(bitmapSource);
			release(frame);
			release(decoder);
			release(wicStream);
			release(fromStream);
			release(wicFactory);

			return result;
		}

		class StormImageSource : public IWICBitmapSource {
		public:
			// Create. Make sure that 'src' is properly pinned!
			StormImageSource(Image *src) : image(src), refcount(1) {}

			// Destroy.
			~StormImageSource() {}

			virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, void **ppvObject) {
				return E_NOINTERFACE;
			}

			virtual ULONG STDMETHODCALLTYPE AddRef(void) {
				return (ULONG)InterlockedIncrement(&refcount);
			}

			virtual ULONG STDMETHODCALLTYPE Release(void) {
				ULONG res = (ULONG)InterlockedDecrement(&refcount);
				if (res == 0)
					delete this;
				return res;
			}

			virtual HRESULT STDMETHODCALLTYPE CopyPalette(IWICPalette *) {
				return WINCODEC_ERR_PALETTEUNAVAILABLE;
			}

			virtual HRESULT STDMETHODCALLTYPE CopyPixels(const WICRect *rect, UINT stride, UINT bufferSize, BYTE *buffer) {
				if (rect->Height < 0 || rect->Width < 0)
					return S_OK;
				if (rect->X < 0 || rect->Y < 0)
					return E_INVALIDARG;

				for (Nat y = 0; y < Nat(rect->Height); y++) {
					byte *src = image->buffer(rect->X, y + rect->Y);
					byte *dest = buffer + y*stride;

					memcpy(dest, src, Nat(rect->Width) * 4);
					for (Nat x = 0; x < Nat(rect->Width); x++) {
						std::swap(dest[x*4], dest[x*4 + 2]);
					}
					// 	byte *src = image->buffer(x + rect->X, y + rect->Y);
					// 	BYTE *dst = buffer + y*stride + x*4;
					// 	dst[0] = src[2];
					// 	dst[1] = src[1];
					// 	dst[2] = src[0];
					// 	dst[3] = src[3];
					// }
				}

				return S_OK;
			}

			virtual HRESULT STDMETHODCALLTYPE GetPixelFormat(WICPixelFormatGUID *format) {
				*format = GUID_WICPixelFormat32bppBGRA;
				return S_OK;
			}

			virtual HRESULT STDMETHODCALLTYPE GetResolution(double *x, double *y) {
				*x = *y = 96.0;
				return S_OK;
			}

			virtual HRESULT STDMETHODCALLTYPE GetSize(UINT *w, UINT *h) {
				*w = image->width();
				*h = image->height();
				return S_OK;
			}

		private:
			Image *image;
			LONG refcount;
		};

		static void saveImage(Image *image, OStream *to, REFGUID type, Property *properties = null) {
			Engine &e = image->engine();

			// Make sure to pin 'image'.
			Image *volatile pinned;
			atomicWrite(pinned, image);

			const wchar *error = S("Failed to initialize Windows Imaging Components");
			IWICImagingFactory *wicFactory = null;
			HRESULT r = CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&wicFactory));

			StormImageSource *imageSource = new StormImageSource(image);

			IWICBitmapEncoder *encoder = null;
			if (SUCCEEDED(r)) {
				error = S("Failed to create an encoder.");
				r = wicFactory->CreateEncoder(type, NULL, &encoder);
			}

			OStreamWrapper *ostreamWrap = OStreamWrapper::create(&to);
			if (SUCCEEDED(r)) {
				error = S("Failed to initialize the encoder.");
				r = encoder->Initialize(ostreamWrap, WICBitmapEncoderNoCache);
			}

			IPropertyBag2 *propertyBag = null;
			IWICBitmapFrameEncode *frame = null;
			if (SUCCEEDED(r)) {
				error = S("Failed to create a new frame.");
				r = encoder->CreateNewFrame(&frame, &propertyBag);
			}

			// To get available properties:
			// if (SUCCEEDED(r)) {
			// 	unsigned long propCount = 0;
			// 	propertyBag->CountProperties(&propCount);
			// 	PVAR(propCount);
			// 	for (unsigned long i = 0; i < propCount; i++) {
			// 		PROPBAG2 p;
			// 		ULONG pcount;

			// 		propertyBag->GetPropertyInfo(i, 1, &p, &pcount);
			// 		PVAR(i);
			// 		PVAR(p.pstrName);
			// 	}
			// }

			// Set properties.
			if (SUCCEEDED(r) && properties) {
				error = S("Failed to set properties.");
				for (Property *p = properties; SUCCEEDED(r) && !p->empty(); p++) {
					PROPBAG2 prop = { 0 };
					prop.pstrName = (wchar *)p->property;
					r = propertyBag->Write(1, &prop, &p->value);
				}
			}

			if (SUCCEEDED(r)) {
				error = S("Failed to initialize the frame.");
				r = frame->Initialize(propertyBag);
			}

			GUID pixelFormat = GUID_WICPixelFormat32bppBGRA;
			if (SUCCEEDED(r)) {
				error = S("Failed to set pixel format.");
				r = frame->SetPixelFormat(&pixelFormat);
			}

			IWICFormatConverter *converter = null;
			if (SUCCEEDED(r)) {
				error = S("Failed to create a format converter.");
				r = wicFactory->CreateFormatConverter(&converter);
			}

			if (SUCCEEDED(r)) {
				error = S("Failed to initialize the format converter.");
				r = converter->Initialize(imageSource, pixelFormat, WICBitmapDitherTypeNone, NULL, 0.0, WICBitmapPaletteTypeCustom);
			}

			IWICBitmapSource *converted = null;
			if (SUCCEEDED(r)) {
				error = S("Failed to extract the converted image.");
				r = converter->QueryInterface(__uuidof(IWICBitmapSource), (void **)&converted);
			}

			if (SUCCEEDED(r)) {
				error = S("Failed to write pixels.");
				WICRect rect = { 0, 0, image->width(), image->height() };
				// r = frame->WriteSource(converted, &rect);
				r = frame->WriteSource(imageSource, &rect);
			}

			if (SUCCEEDED(r)) {
				error = S("Failed to commit image data.");
				frame->Commit();
			}
			if (SUCCEEDED(r)) {
				encoder->Commit();
			}

			release(converted);
			release(converter);
			release(propertyBag);
			release(frame);
			release(ostreamWrap);
			release(encoder);
			release(imageSource);
			release(wicFactory);

			if (FAILED(r))
				throw new (e) ImageSaveError(error);

			to->flush();
		}

	}


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

	void PNGOptions::save(Image *image, OStream *to) {
		// We can set InterlaceOption and FilterOption.
		wic::saveImage(image, to, GUID_ContainerFormatPng, null);
	}

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

	void JPEGOptions::save(Image *image, OStream *to) {
		wic::Property props[] = {
			wic::Property(L"ImageQuality", float(quality) / 100.0f),
			wic::Property(),
		};

		wic::saveImage(image, to, GUID_ContainerFormatJpeg, props);
	}

}

#endif