/* -LICENSE-START-
** Copyright (c) 2011 Blackmagic Design
**
** Permission is hereby granted, free of charge, to any person or organization
** obtaining a copy of the software and accompanying documentation covered by
** this license (the "Software") to use, reproduce, display, distribute,
** execute, and transmit the Software, and to prepare derivative works of the
** Software, and to permit third-parties to whom the Software is furnished to
** do so, all subject to the following:
** 
** The copyright notices in the Software and this entire statement, including
** the above license grant, this restriction and the following disclaimer,
** must be included in all copies of the Software, in whole or in part, and
** all derivative works of the Software, unless such copies or derivative
** works are solely in the form of machine-executable object code generated by
** a source language processor.
** 
** THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
** IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
** FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
** SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
** FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
** ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
** DEALINGS IN THE SOFTWARE.
** -LICENSE-END-
*/

#import "QTDecodeSession.h"

#define BAIL_IF(cond, fmt, args...)                                        \
if (cond)                                                                  \
{                                                                          \
    NSLog(@"[%s:%3d] %s: " fmt, __FILE__, __LINE__, __FUNCTION__, ##args); \
    goto bail;                                                             \
}

static void StaticDecompressionComplete(void* session, OSStatus result,
										ICMDecompressionTrackingFlags flags,CVPixelBufferRef pixBuf,
										TimeValue64 time, TimeValue64 duration,
										ICMValidTimeFlags timeFlags, void* reserved, void* frameRef)
{
	QTDecodeSession* realSelf = (QTDecodeSession*)session;
	
	[realSelf decompressionCompleteWithResult:result
										flags:flags
									   pixBuf:pixBuf
										 time:time
									 duration:duration
									timeFlags:timeFlags];
}

@implementation QTDecodeSession

- (id)initWithMovieWidth:(int)width andHeight:(int)height
{
	self = [super init];
	if (! self)
		return self;
	
	m_width = width;
	m_height = height;
	
	m_nalParser = CreateBMDStreamingH264NALParser();
	
	if (m_nalParser == NULL)
	{
		NSLog(@"Error: couldn't load NAL parser\n");
		[self dealloc];
		return nil;
	}
	
	m_needsSPS = YES;
	m_needsPPS = YES;
	m_hasSentNALs = NO;
	
	return self;
}

- (void)dealloc
{
	if (m_nalParser)
		m_nalParser->Release();
	
	if (m_session != NULL)
		ICMDecompressionSessionRelease(m_session);
	
	[super dealloc];
}

- (ImageDescriptionHandle)createImageDescriptionHandle
{
	// This is whole lot of QuickTime stuff needed to set up the decoding session
	ImageDescriptionHandle	imageDescriptionHandle = NULL;
	char*					imgDescData	= NULL;
	ImageDescription*		imgDescPtr	= NULL;
	
	uint32_t spsLen = m_firstSPS->GetPayloadSize();
	const uint8_t* spsData;
	m_firstSPS->GetBytes((void**)&spsData);
	
	uint32_t ppsLen = m_firstPPS->GetPayloadSize();
	const uint8_t* ppsData;
	m_firstPPS->GetBytes((void**)&ppsData);
	
	uint32_t spsProfileIdc, spsProfileCompatability, spsLevelIdc;
	if (FAILED(m_nalParser->GetProfileAndLevelFromSPS(m_firstSPS, &spsProfileIdc, &spsProfileCompatability, &spsLevelIdc)))
	{
		NSLog(@"Failed to get profile and level from SPS!\n");
		return NULL;
	}
	
	imgDescData = (char*)malloc(1 << 20);
	bzero(imgDescData, (1 << 20));
	
	{
		// Write out the ImageDescription native-endian.
		imgDescPtr = (ImageDescription*)imgDescData;
		imgDescPtr->cType = 'avc1';
		imgDescPtr->dataRefIndex = 1;
		imgDescPtr->vendor = 'BMAG';
		imgDescPtr->width = m_width & ~0x0F;
		imgDescPtr->height = m_height;
		
		imgDescPtr->hRes = 72 << 16;
		imgDescPtr->vRes = 72 << 16;
		imgDescPtr->frameCount = 1;
		strcpy((char*)imgDescPtr->name + 1, "H.264/AVC");
		imgDescPtr->name[0] = strlen((char*)imgDescPtr->name + 1);
		imgDescPtr->depth = 24;
		imgDescPtr->clutID = -1;
		
		// Get a pointer to the first byte after the image description and
		// write the image description extension atoms out big-endian.
		char* p = (char*)(imgDescPtr + 1);
		
		// 'avcC' atom
		uint32_t	atomCode				= 'avcC';
		uint8_t		version					= 0x01;
		uint8_t		profile;
		uint8_t		level;
		uint8_t		profileCompatibility;
		uint8_t		reserved1				= 0x3f;
		uint8_t		length_size				= 0x04;
		uint8_t		reserved2				= 0x07;
		uint8_t		sequenceParameterSets	= 0x01;
		uint8_t		pictureParameterSets	= 0x01;
		
		profile = spsProfileIdc;
		level = spsLevelIdc;
		profileCompatibility = spsProfileCompatability;
		
		uint32_t* lenp = (uint32_t*)p;
		p += sizeof(*lenp);
		*(uint32_t*)p = OSSwapHostToBigInt32(atomCode);
		p += sizeof(atomCode);
		*p++ = version;
		*p++ = profile;
		*p++ = profileCompatibility;
		*p++ = level;
		*p++ = (reserved1 << 2) | ((length_size - 1) & 3);
		*p++ = (reserved2 << 5) | (sequenceParameterSets & 0x1f);
		*(uint16_t*)p = OSSwapHostToBigInt16(spsLen);
		p += sizeof(uint16_t);
		memcpy(p, spsData, spsLen);
		p += spsLen;
		*p++ = pictureParameterSets;
		*(uint16_t*)p = OSSwapHostToBigInt16(ppsLen);
		p += sizeof(uint16_t);
		memcpy(p, ppsData, ppsLen);
		p += ppsLen;
		*lenp = OSSwapHostToBigInt32(p - (char*)lenp);
		
		struct uuid_atom_t
		{
			uint32_t length;
			uint32_t code;
			uuid_t   uuid;
			uint32_t bollocks;
		};
		
		// iTunes looks for a specific UUID when deciding whether to accept a
		// movie in iPod High format. It could probably be left out
		// of iPod Standard movies since iTunes doesn't appear to require it,
		// but it's probably safer to leave it in there in case that
		// restriction is tightened in a future iTunes version.
		uuid_t uuid = {0x6b, 0x68, 0x40, 0xf2, 0x5f, 0x24, 0x4f, 0xc5,
			0xba, 0x39, 0xa5, 0x1b, 0xcf, 0x03, 0x23, 0xf3};
		uuid_atom_t* u = (uuid_atom_t*)p;
		u->length = OSSwapHostToBigInt32(sizeof(*u));
		u->code = OSSwapHostToBigInt32('uuid');
		memcpy(u->uuid, uuid, sizeof(u->uuid));
		u->bollocks = OSSwapHostToBigInt32(1);
		p = (char*)(u + 1);
		
		// 'pasp' atom - TODO: Deal with aspect ratios properly!
		atomCode = 'pasp';
		uint32_t h = 1;
		uint32_t v = 1;
		
		lenp = (uint32_t*)p;
		p += sizeof(*lenp);
		*(uint32_t*)p = OSSwapHostToBigInt32(atomCode);
		p += sizeof(atomCode);
		*(uint32_t*)p = OSSwapHostToBigInt32(h);
		p += sizeof(h);
		*(uint32_t*)p = OSSwapHostToBigInt32(v);
		p += sizeof(v);
		*lenp = OSSwapHostToBigInt32(p - (char*)lenp);
		
		// Fill in the length of the image description. (native-endian again)
		imgDescPtr->idSize = p - imgDescData;		
	}
	
	imageDescriptionHandle = (ImageDescriptionHandle)::NewHandle(imgDescPtr->idSize);
	
	memcpy(*imageDescriptionHandle, imgDescData, imgDescPtr->idSize);

	free(imgDescData);
	
	return imageDescriptionHandle;
}

- (void)startDecompressionSession
{
	int err;
	ImageDescriptionHandle imageDescriptionHandle = [self createImageDescriptionHandle];
	
	m_session = nil;
	
	m_tracking.decompressionTrackingCallback = StaticDecompressionComplete;
	m_tracking.decompressionTrackingRefCon = self;
	
	{
		short	width	= (*imageDescriptionHandle)->width;
		short	height	= (*imageDescriptionHandle)->height;
		long	format	= k2vuyPixelFormat;
		
		CFMutableDictionaryRef pixBufAttributes = ::CFDictionaryCreateMutable(NULL, 0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
		
		CFNumberRef number = ::CFNumberCreate(NULL, kCFNumberShortType, &width);
		::CFDictionaryAddValue(pixBufAttributes, kCVPixelBufferWidthKey, number);
		::CFRelease(number);
		
		number = ::CFNumberCreate(NULL, kCFNumberShortType, &height);
		::CFDictionaryAddValue(pixBufAttributes, kCVPixelBufferHeightKey, number);
		::CFRelease(number);
		
		number = ::CFNumberCreate(NULL, kCFNumberSInt32Type, &format);
		::CFDictionaryAddValue(pixBufAttributes, kCVPixelBufferPixelFormatTypeKey, number);
		::CFRelease(number);
		
		err = ::ICMDecompressionSessionCreate(NULL, imageDescriptionHandle, NULL,
											  pixBufAttributes, &m_tracking, &m_session);
		::CFRelease(pixBufAttributes);
		BAIL_IF(err, "ICMDecompressionSessionCreate returned %i, session %p\n", err, m_session);
		
	}
	
bail:
	::DisposeHandle((Handle)imageDescriptionHandle);
}

- (void)setDelegate:(id<DecodeSessionDelegate>)delegate
{
	m_delegate = delegate;
}

- (void)handleH264NAL:(IBMDStreamingH264NALPacket*)nal
{
	int err;
	ICMFrameTimeRecord frameTime = {0};
	
	if (m_needsSPS && m_nalParser->IsNALSequenceParameterSet(nal) == S_OK)
	{
		m_firstSPS = nal;
		m_firstSPS->AddRef();
		m_needsSPS = NO;
	}
	
	if (m_needsPPS && m_nalParser->IsNALPictureParameterSet(nal) == S_OK)
	{
		m_firstPPS = nal;
		m_firstPPS->AddRef();
		m_needsPPS = NO;
	}
	
	if (!m_hasSentNALs)
	{
		// Skip all NALs until we have both SPS and PPS, then we can set up the
		// decompression session.
		if (!m_firstSPS || !m_firstPPS)
			goto bail;

		[self startDecompressionSession];
		m_firstSPS->Release();
		m_firstPPS->Release();
		m_firstSPS = NULL;
		m_firstPPS = NULL;
		m_hasSentNALs = YES;
	}
	
	// Now we actually pass the NAL to QuickTime:
	m_nalDecompressed = NO;
	m_pixBuf = NULL;

	// We handle the frame time ourselves.
	frameTime.recordSize = sizeof(frameTime);
	frameTime.scale = 1;
	frameTime.rate = fixed1;
	frameTime.flags = icmFrameTimeIsNonScheduledDisplayTime;	
	
	// Enqueue the NAL data:
	const uint8_t* nalData;
	nal->GetBytesWithSizePrefix((void**)&nalData);
	err = ::ICMDecompressionSessionDecodeFrame(m_session, nalData, nal->GetPayloadSize()+4, nil, &frameTime, nil);
	BAIL_IF(err, "ICMDecompressionSessionDecodeFrame failed with %ld\n", err);
	
	// Wait for the decompression to be done (serially).  Since we're in
	// a single threaded decode session, we know the callback will be
	// called serially.  If we weren't, we'd have to wait on a pthread_cond_t
	// here.
	err = ::ICMDecompressionSessionSetNonScheduledDisplayTime(m_session, 0, 1, 0);
	BAIL_IF(err, "ICMDecompressionSessionSetNonScheduledDisplayTime failed with %ld\n", err);

	BAIL_IF(!m_nalDecompressed, "Unexpected error: nal decompressed should be true here\n");

	if (m_pixBuf)
	{
		[m_delegate haveVideoFrame:m_pixBuf fromNAL:nal];
		
		::CVPixelBufferRelease(m_pixBuf);
		m_pixBuf = NULL;
	}

bail:
	return;
}

- (void)decompressionCompleteWithResult:(OSStatus)result
								  flags:(ICMDecompressionTrackingFlags)flags
								 pixBuf:(CVPixelBufferRef)pixBuf
								   time:(TimeValue64)time
							   duration:(TimeValue64)duration
							  timeFlags:(ICMValidTimeFlags)timeFlags
{
	if (flags & kICMDecompressionTracking_FrameDecoded)
	{
		::CVPixelBufferRetain(pixBuf);
		m_pixBuf = pixBuf;
	}
	else
	{
		m_pixBuf = NULL;
	}
	
	m_nalDecompressed = YES;
}

@end
