#include "cngdecoder.hpp"

#ifdef USE_TURBOJPEG
	#include "../deps_include/turbojpeg.h"
#else
	#define _NJ_INCLUDE_HEADER_ONLY
	#include "nanojpeg.cpp"
#endif

// {A6B7D610-C2FB-4EA5-8271-EA0B963321D2} WAS 67b4...
static const GUID CLSID_CNG_Container =
{ 0xa6b7d610, 0xc2fb, 0x4ea5, { 0x82, 0x71, 0xea, 0xb, 0x96, 0x33, 0x21, 0xd2 } };

// {70B4C80E-312B-48D9-A454-BD1AEB5A9E81} WAS ebf7...
static const GUID CLSID_CNG_Decoder =
{ 0x70b4c80e, 0x312b, 0x48d9, { 0xa4, 0x54, 0xbd, 0x1a, 0xeb, 0x5a, 0x9e, 0x81 } };

// {A3738EE8-6A83-4163-8719-C372A851B18E}
static const GUID CLSID_CNG_PropertyStore =
{ 0xa3738ee8, 0x6a83, 0x4163, { 0x87, 0x19, 0xc3, 0x72, 0xa8, 0x51, 0xb1, 0x8e } };


namespace cngx
{
	template< typename T >
	static HRESULT ReadFromIStream( IStream *stream, T &val )
	{
		HRESULT result = S_OK;
		if ( NULL == stream ) result = E_INVALIDARG;
		ULONG numRead = 0;
		if ( SUCCEEDED( result )) result = stream->Read( &val, sizeof(T), &numRead );
		if ( SUCCEEDED( result )) result = ( sizeof(T) == numRead ) ? S_OK : E_UNEXPECTED;
		return result;
	}

	static HRESULT ReadBytesFromIStream( IStream *stream, unsigned char *bytes, unsigned count )
	{
		HRESULT result = S_OK;
		if ( NULL == stream ) result = E_INVALIDARG;
		ULONG numRead = 0;
		if ( SUCCEEDED( result )) result = stream->Read( bytes, count, &numRead );
		if ( SUCCEEDED( result )) result = ( count == numRead ) ? S_OK : E_UNEXPECTED;
		return result;
	}

	static HRESULT GetIStreamLength(IStream *stream, ULARGE_INTEGER* pOut) {

		STATSTG ssg    = { 0 };
		HRESULT result = S_OK;
		if (NULL == stream) result = E_INVALIDARG;

		if (SUCCEEDED(result)) result = stream->Stat(&ssg, STATFLAG_NONAME);

		*pOut = ssg.cbSize;

		return result;
	}

	//----------------------------------------------------------------------------------------
	// CNG_FrameDecode implementation
	//----------------------------------------------------------------------------------------

	CNG_FrameDecode::CNG_FrameDecode(IWICImagingFactory *pIFactory, UINT num)
		: BaseFrameDecode( pIFactory, num )
	{
	}

	HRESULT CNG_FrameDecode::LoadImage(IStream *pIStream)
	{
		HRESULT result = S_OK;

		// Get data size

		ULARGE_INTEGER streamsize = { 0 };
		GetIStreamLength(pIStream, &streamsize);

		if (!SUCCEEDED(result)) return result;
		
		// Read all data

		std::vector<unsigned char> data;
		data.resize(streamsize.QuadPart);
		ReadBytesFromIStream(pIStream, &data[0], streamsize.QuadPart);
		
		if (!SUCCEEDED(result)) return result;

		// Decrypt

		for (size_t i = 0, e = streamsize.QuadPart; i != e; ++i) {
			data[i] ^= 0xEF;
		}

		// Load JPEG header

		std::vector<unsigned char> rawbmp;

#ifdef USE_TURBOJPEG

		tjhandle jpeg = tjInitDecompress();
		int width, height, subsamp;
		int tjret = tjDecompressHeader2(jpeg, &data[0], data.size(), &width, &height, &subsamp);

		if (tjret != 0) return E_INVALID_PROTOCOL_FORMAT;

		// Load JPEG data

		rawbmp.resize( TJBUFSIZE(width, height) );
		tjret = tjDecompress(jpeg, &data[0], data.size(), &rawbmp[0], width, width * 4, height, TJPF_RGBX, 0);

		if (tjret != 0) return E_INVALID_PROTOCOL_FORMAT;

		result = FillBitmapSource(width, height, GUID_WICPixelFormat32bpp3ChannelsAlpha, width * 4, rawbmp.size(), &rawbmp[0]);

#else

		njInit();

		if (njDecode(&data[0], data.size()) != NJ_OK) {
			njDone();
			return E_INVALID_PROTOCOL_FORMAT;
		}

		rawbmp.resize(njGetImageSize());
		memcpy(&rawbmp[0], njGetImage(), rawbmp.size());

		int width = njGetWidth();

		result = FillBitmapSource(width, njGetHeight(), GUID_WICPixelFormat24bppRGB, width * 3, rawbmp.size(), &rawbmp[0]);

		njDone();
#endif

		return result;
	}

	HRESULT CNG_FrameDecode::FillBitmapSource(UINT width, UINT height,
		REFWICPixelFormatGUID pixelFormat, UINT cbStride, UINT cbBufferSize, BYTE *pbBuffer)
	{
		HRESULT result = S_OK;

		IWICBitmap* ptr;

		result = m_factory->CreateBitmapFromMemory(width, height, pixelFormat, cbStride, cbBufferSize, pbBuffer, &ptr);

		//reinterpret_cast<IWICBitmap **>(&(m_bitmapSource)));
		if (SUCCEEDED(result)) {
			m_bitmapSource = reinterpret_cast<IWICBitmapSource*>(ptr);
		}

//		m_thumbnail = m_bitmapSource;

		return result;
	}

	//----------------------------------------------------------------------------------------
	// CNG_Decoder implementation
	//----------------------------------------------------------------------------------------

	CNG_Decoder::CNG_Decoder()
		: BaseDecoder( CLSID_CNG_Decoder, CLSID_CNG_Container )
	{
	}

	CNG_Decoder::~CNG_Decoder()
	{
	}

	// TODO: implement real query capability
	STDMETHODIMP CNG_Decoder::QueryCapability(IStream *pIStream, DWORD *pCapability)
	{
		UNREFERENCED_PARAMETER( pIStream );

		HRESULT result = S_OK;

		*pCapability =
			WICBitmapDecoderCapabilityCanDecodeAllImages;

		return result;
	}


	STDMETHODIMP CNG_Decoder::Initialize(IStream *pIStream, WICDecodeOptions cacheOptions)
	{
		UNREFERENCED_PARAMETER( cacheOptions );

		HRESULT result = E_INVALIDARG;

		ReleaseMembers( true );

		if ( pIStream )
		{
			VerifyFactory();
			
			// There's only one frame
			
			CNG_FrameDecode* frame = CreateNewDecoderFrame(m_factory, 0);

			result = frame->LoadImage( pIStream );
				
			if ( SUCCEEDED( result ))
				AddDecoderFrame( frame );
			else
				delete frame;

		}
		else
			result = E_INVALIDARG;

		return result;
	}

	CNG_FrameDecode* CNG_Decoder::CreateNewDecoderFrame(IWICImagingFactory* factory, UINT i)
	{
		return new CNG_FrameDecode(factory, i);
	}

	void CNG_Decoder::Register(RegMan &regMan)
	{
		HMODULE curModule = GetModuleHandle( L"natgeo-wic.dll" );
		wchar_t tempFileName[MAX_PATH];
		if ( curModule != NULL ) GetModuleFileName( curModule, tempFileName, MAX_PATH );
		
		/*                     {CATID_WICBitmapDecoders             } */
		regMan.SetSZ( L"CLSID\\{7ED96837-96F0-4812-B211-F13C24117ED3}\\Instance\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"CLSID", L"{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}" );
		regMan.SetSZ( L"CLSID\\{7ED96837-96F0-4812-B211-F13C24117ED3}\\Instance\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"FriendlyName", L"CNG Decoder" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"Version", L"1.0.0.1" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"Date", _T(__DATE__) );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"SpecVersion", L"1.0.0.0" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"ColorManagementVersion", L"1.0.0.0" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"MimeTypes", L"x-image/cng" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"FileExtensions", L".cng" );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"SupportsAnimation", 0 );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"SupportChromakey", 1 );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"SupportLossless", 1 );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"SupportMultiframe", 1 );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"ContainerFormat", L"{A6B7D610-C2FB-4EA5-8271-EA0B963321D2}" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"Author", L"natgeo-wic authors" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"Description", L"CNG Format Decoder" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}", L"FriendlyName", L"CNG Decoder" );

		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Formats", L"", L"" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Formats\\{6FDDC324-4E03-4BFE-B185-3D77768DC90F}", L"", L"" );
		/*                                                                      {GUID_WICPixelFormat32bppBGRA        } */

		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\InprocServer32", L"", tempFileName );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\InprocServer32", L"ThreadingModel", L"Apartment" );
		regMan.SetSZ( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Patterns", L"", L"" );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Patterns\\0", L"Position", 0 );
		regMan.SetDW( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Patterns\\0", L"Length", 4 );

		BYTE bytes[8] = { 0 };
		bytes[0] = 0x10; bytes[1] = 0x37; bytes[2] = 0x10; bytes[3] = 0x0F;
		regMan.SetBytes( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Patterns\\0", L"Pattern", bytes, 4 );
		bytes[0] = bytes[1] = bytes[2] = bytes[3] = 0xFF;
		regMan.SetBytes( L"CLSID\\{70B4C80E-312B-48D9-A454-BD1AEB5A9E81}\\Patterns\\0", L"Mask", bytes, 4 );

		// Register the filetype too

	}
}
