#include <Carbon/Carbon.h>
#include <QuickTime/QuickTime.h>
#include "VobSubCodec.h"
#include "Codecprintf.h"
#include <zlib.h>
#include "avcodec.h"
#include "intreadwrite.h"

// Constants
const UInt8 kNumPixelFormatsSupportedVobSub = 1;

// Data structures
typedef struct	{
	ComponentInstance       self;
	ComponentInstance       delegateComponent;
	ComponentInstance       target;
	OSType**                wantedDestinationPixelTypeH;
	ImageCodecMPDrawBandUPP drawBandUPP;
	
	UInt32                  palette[16];
	
	int                     compressed;
	uint8_t                 *codecData;
	unsigned int            bufferSize;
	
	AVCodec                 *avCodec;
	AVCodecContext          *avContext;
	AVSubtitle              subtitle;
} VobSubCodecGlobalsRecord, *VobSubCodecGlobals;

typedef struct {
	long		width;
	long		height;
	long		bufferSize;
	char		decoded;
} VobSubDecompressRecord;

typedef struct {
	// color format is 32-bit ARGB
	UInt32  pixelColor[16];
	UInt32  duration;
} PacketControlData;

// Setup required for ComponentDispatchHelper.c
#define IMAGECODEC_BASENAME() 		VobSubCodec
#define IMAGECODEC_GLOBALS() 		VobSubCodecGlobals storage

#define CALLCOMPONENT_BASENAME()	IMAGECODEC_BASENAME()
#define	CALLCOMPONENT_GLOBALS()		IMAGECODEC_GLOBALS()

#define COMPONENT_UPP_PREFIX()		uppImageCodec
#define COMPONENT_DISPATCH_FILE		"VobSubCodecDispatch.h"
#define COMPONENT_SELECT_PREFIX()  	kImageCodec

#define	GET_DELEGATE_COMPONENT()	(storage->delegateComponent)

#include <CoreServices/Components.k.h>
#include <QuickTime/ImageCodec.k.h>
#include <QuickTime/ComponentDispatchHelper.c>


// dest must be at least as large as src
void ExtractVobSubPacket(UInt8 *dest, UInt8 *framedSrc, int srcSize);
static ComponentResult ReadPacketControls(UInt8 *packet, UInt32 palette[16], PacketControlData *controlDataOut);
extern void initLib();

ComponentResult VobSubCodecOpen(VobSubCodecGlobals glob, ComponentInstance self)
{
	ComponentResult err;
	// Allocate memory for our globals, set them up and inform the component manager that we've done so
	glob = (VobSubCodecGlobals)NewPtrClear(sizeof(VobSubCodecGlobalsRecord));
	if (err = MemError()) goto bail;

	SetComponentInstanceStorage(self, (Handle)glob);

	glob->self = self;
	glob->target = self;
	glob->wantedDestinationPixelTypeH = (OSType **)NewHandle(sizeof(OSType) * 2);
	if (err = MemError()) goto bail;
	glob->drawBandUPP = NULL;
	
	// Open and target an instance of the base decompressor as we delegate
	// most of our calls to the base decompressor instance
	err = OpenADefaultComponent(decompressorComponentType, kBaseCodecType, &glob->delegateComponent);
	if (err) goto bail;

	ComponentSetTarget(glob->delegateComponent, self);

bail:
	return err;
}

ComponentResult VobSubCodecClose(VobSubCodecGlobals glob, ComponentInstance self)
{
	int i;
	
	// Make sure to close the base component and dealocate our storage
	if (glob) {
		if (glob->delegateComponent) {
			CloseComponent(glob->delegateComponent);
		}
		if (glob->wantedDestinationPixelTypeH) {
			DisposeHandle((Handle)glob->wantedDestinationPixelTypeH);
		}
		if (glob->drawBandUPP) {
			DisposeImageCodecMPDrawBandUPP(glob->drawBandUPP);
		}
		if (glob->codecData) {
			av_free(glob->codecData);
		}
		if (glob->avCodec) {
			avcodec_close(glob->avContext);
		}
		if (glob->avContext) {
			av_free(glob->avContext);
		}
		if (glob->subtitle.rects) {
			for (i = 0; i < glob->subtitle.num_rects; i++) {
				av_free(glob->subtitle.rects[i].rgba_palette);
				av_free(glob->subtitle.rects[i].bitmap);
			}
			av_free(glob->subtitle.rects);
		}

		DisposePtr((Ptr)glob);
	}

	return noErr;
}

ComponentResult VobSubCodecVersion(VobSubCodecGlobals glob)
{	
	return kVobSubCodecVersion;
}

ComponentResult VobSubCodecTarget(VobSubCodecGlobals glob, ComponentInstance target)
{
	glob->target = target;
	return noErr;
}

ComponentResult VobSubCodecGetMPWorkFunction(VobSubCodecGlobals glob, ComponentMPWorkFunctionUPP *workFunction, void **refCon)
{
	if (glob->drawBandUPP == NULL)
		glob->drawBandUPP = NewImageCodecMPDrawBandUPP((ImageCodecMPDrawBandProcPtr)VobSubCodecDrawBand);
	
	return ImageCodecGetBaseMPWorkFunction(glob->delegateComponent, workFunction, refCon, glob->drawBandUPP, glob);
}

ComponentResult VobSubCodecInitialize(VobSubCodecGlobals glob, ImageSubCodecDecompressCapabilities *cap)
{
	cap->decompressRecordSize = sizeof(VobSubDecompressRecord);
	cap->canAsync = true;
	cap->baseCodecShouldCallDecodeBandForAllFrames = true;
	
	if(cap->recordSize > offsetof(ImageSubCodecDecompressCapabilities, baseCodecShouldCallDecodeBandForAllFrames)) 
		cap->subCodecIsMultiBufferAware = true;

	return noErr;
}

static ComponentResult SetupColorPalette(VobSubCodecGlobals glob, ImageDescriptionHandle imageDescription) {
	OSErr err = noErr;
	Handle descExtension = NewHandle(0);
	
	err = GetImageDescriptionExtension(imageDescription, &descExtension, kVobSubIdxExtension, 1);
	if (err) goto bail;
	
	char *string = (char *) *descExtension;
	
	char *palette = strnstr(string, "palette:", GetHandleSize(descExtension));
	
	if (palette != NULL) {
		sscanf(palette, "palette: %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx, %lx", 
			   &glob->palette[ 0], &glob->palette[ 1], &glob->palette[ 2], &glob->palette[ 3], 
			   &glob->palette[ 4], &glob->palette[ 5], &glob->palette[ 6], &glob->palette[ 7], 
			   &glob->palette[ 8], &glob->palette[ 9], &glob->palette[10], &glob->palette[11], 
			   &glob->palette[12], &glob->palette[13], &glob->palette[14], &glob->palette[15]);
	}
	
bail:
	DisposeHandle(descExtension);
	return err;
}

ComponentResult VobSubCodecPreflight(VobSubCodecGlobals glob, CodecDecompressParams *p)
{
	CodecCapabilities *capabilities = p->capabilities;
	OSTypePtr         formats = *glob->wantedDestinationPixelTypeH;

	// Specify the minimum image band height supported by the component
	// bandInc specifies a common factor of supported image band heights
	capabilities->bandMin = (**p->imageDescription).height;
	capabilities->bandInc = capabilities->bandMin;

	// Indicate the wanted destination using the wantedDestinationPixelTypeH previously set up
	capabilities->wantedPixelSize  = 0; 	
	
	// we want 4:4:4:4 ARGB
	*formats++ = k32ARGBPixelFormat;
	*formats++ = 0;
	
	p->wantedDestinationPixelTypes = glob->wantedDestinationPixelTypeH;

	// Specify the number of pixels the image must be extended in width and height if
	// the component cannot accommodate the image at its given width and height
	capabilities->extendWidth = 0;
	capabilities->extendHeight = 0;
	
	// get the color palette info from the image description
	SetupColorPalette(glob, p->imageDescription);

	if (isImageDescriptionExtensionPresent(p->imageDescription, kMKVCompressionExtension))
		glob->compressed = 1;
	
	if (!glob->avCodec) {
		initLib();
		
		glob->avCodec = avcodec_find_decoder(CODEC_ID_DVD_SUBTITLE);
		glob->avContext = avcodec_alloc_context();
		
		if (avcodec_open(glob->avContext, glob->avCodec)) {
			Codecprintf(NULL, "Error opening DVD subtitle decoder\n");
			return codecErr;
		}
	}
	
	return noErr;
}

ComponentResult VobSubCodecBeginBand(VobSubCodecGlobals glob, CodecDecompressParams *p, ImageSubCodecDecompressRecord *drp, long flags)
{
	VobSubDecompressRecord *myDrp = (VobSubDecompressRecord *)drp->userDecompressRecord;

	// Let base codec know that all our frames are keyframes
	drp->frameType = kCodecFrameTypeKey;

	myDrp->width = (**p->imageDescription).width;
	myDrp->height = (**p->imageDescription).height;
	myDrp->bufferSize = p->bufferSize;
	myDrp->decoded = p->frameTime ? (0 != (p->frameTime->flags & icmFrameAlreadyDecoded)) : false;
	
	return noErr;
}

void DecompressZlib(VobSubCodecGlobals glob, uint8_t *data, long *bufferSize)
{
	ComponentResult err = noErr;
	z_stream strm;
	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = 0;
	strm.next_in = Z_NULL;
	err = inflateInit(&strm);
	if (err != Z_OK) return;
	
	strm.avail_in = *bufferSize;
	strm.next_in = data;
	
	// first, get the size of the decompressed data
	strm.avail_out = 2;
	strm.next_out = glob->codecData;
	
	err = inflate(&strm, Z_SYNC_FLUSH);
	if (err < Z_OK) goto bail;
	if (strm.avail_out != 0) goto bail;
	
	// reallocate our buffer to be big enough to store the decompressed packet
	*bufferSize = AV_RB16(glob->codecData);
	glob->codecData = av_fast_realloc(glob->codecData, &glob->bufferSize, *bufferSize);
	
	// then decompress the rest of it
	strm.avail_out = glob->bufferSize - 2;
	strm.next_out = glob->codecData + 2;
	
	inflate(&strm, Z_SYNC_FLUSH);
bail:
	inflateEnd(&strm);
}

ComponentResult VobSubCodecDecodeBand(VobSubCodecGlobals glob, ImageSubCodecDecompressRecord *drp, unsigned long flags)
{
	VobSubDecompressRecord *myDrp = (VobSubDecompressRecord *)drp->userDecompressRecord;
	UInt8 *data = (UInt8 *) drp->codecData;
	int ret, got_sub;
	
	if (glob->codecData == NULL) {
		glob->codecData = malloc(myDrp->bufferSize + 2);
		glob->bufferSize = myDrp->bufferSize + 2;
	}
	
	// make sure we have enough space to store the packet
	glob->codecData = av_fast_realloc(glob->codecData, &glob->bufferSize, myDrp->bufferSize + 2);
	
	if (glob->compressed) {
		DecompressZlib(glob, data, &myDrp->bufferSize);
		
	// the header of a spu PS packet starts 0x000001bd
	// if it's raw spu data, the 1st 2 bytes are the length of the data
	} else if (data[0] + data[1] == 0) {
		// remove the MPEG framing
		ExtractVobSubPacket(glob->codecData, data, myDrp->bufferSize);
	} else {
		memcpy(glob->codecData, drp->codecData, myDrp->bufferSize);
	}
	
	ret = avcodec_decode_subtitle(glob->avContext, &glob->subtitle, &got_sub, glob->codecData, myDrp->bufferSize);
	
	if (ret < 0 || !got_sub) {
		Codecprintf(NULL, "Error decoding DVD subtitle %d / %ld\n", ret, myDrp->bufferSize);
		return codecErr;
	}
	
	myDrp->decoded = true;
	
	return noErr;
}

ComponentResult VobSubCodecDrawBand(VobSubCodecGlobals glob, ImageSubCodecDecompressRecord *drp)
{
	OSErr err = noErr;
	VobSubDecompressRecord *myDrp = (VobSubDecompressRecord *)drp->userDecompressRecord;
	PacketControlData controlData;
	uint8_t *data = glob->codecData;
	unsigned int i, j, x, y;
	int usePalette = 0;
	
	if(!myDrp->decoded)
		err = VobSubCodecDecodeBand(glob, drp, 0);
	if (err) return err;
	
	// clear the buffer to pure transparent
	memset(drp->baseAddr, 0, myDrp->height * drp->rowBytes);
	
	err = ReadPacketControls(data, glob->palette, &controlData);
	if (err == noErr)
		usePalette = true;
	
	for (i = 0; i < glob->subtitle.num_rects; i++) {
		AVSubtitleRect *rect = &glob->subtitle.rects[i];
		uint8_t *line = (uint8_t *)drp->baseAddr + drp->rowBytes * rect->y + rect->x*4;
		uint8_t *sub = rect->bitmap;
		unsigned int w = FFMIN(rect->w, myDrp->width  - rect->x);
		unsigned int h = FFMIN(rect->h, myDrp->height - rect->y);
		
		if (usePalette) {
			for (j = 0; j < 4; j++)
				rect->rgba_palette[j] = controlData.pixelColor[j];
		}
		
		for (y = 0; y < h; y++) {
			uint32_t *pixel = (uint32_t *) line;
			
			for (x = 0; x < w; x++)
				pixel[x] = rect->rgba_palette[sub[x]];
			
			line += drp->rowBytes;
			sub += rect->linesize;
		}
	}
	
	return err;
}

ComponentResult VobSubCodecEndBand(VobSubCodecGlobals glob, ImageSubCodecDecompressRecord *drp, OSErr result, long flags)
{
	return noErr;
}

ComponentResult VobSubCodecQueueStarting(VobSubCodecGlobals glob)
{
	return noErr;
}

ComponentResult VobSubCodecQueueStopping(VobSubCodecGlobals glob)
{
	return noErr;
}

ComponentResult VobSubCodecGetCompressedImageSize(VobSubCodecGlobals glob, ImageDescriptionHandle desc, Ptr data, long dataSize, ICMDataProcRecordPtr dataProc, long *size)
{
	if (size == NULL) 
		return paramErr;
	
	return unimpErr;
}

ComponentResult VobSubCodecGetCodecInfo(VobSubCodecGlobals glob, CodecInfo *info)
{
	OSErr err = noErr;
	if (info == NULL) {
		err = paramErr;
	} else {
		CodecInfo **tempCodecInfo;

		err = GetComponentResource((Component)glob->self, codecInfoResourceType, kVobSubCodecResource, (Handle *)&tempCodecInfo);
		if (err == noErr) {
			*info = **tempCodecInfo;
			DisposeHandle((Handle)tempCodecInfo);
		}
	}

	return err;
}

void ExtractVobSubPacket(UInt8 *dest, UInt8 *framedSrc, int srcSize) {
	int copiedBytes = 0;
	UInt8 *currentPacket = framedSrc;
	
	while (currentPacket - framedSrc < srcSize) {
		// 3-byte start code: 0x00 00 01
		if (currentPacket[0] + currentPacket[1] != 0 || currentPacket[2] != 1) {
			Codecprintf(NULL, "VobSub Codec: !! Unknown header: %02x %02x %02x\n", currentPacket[0], currentPacket[1], currentPacket[2]);
			return;
		}
		
		int packet_length;
		
		switch (currentPacket[3]) {
			case 0xba:
				// discard PS packets; nothing in them we're interested in
				// here, packet_length is the additional stuffing
				packet_length = currentPacket[13] & 0x7;
				
				currentPacket += 14 + packet_length;
				break;
				
			case 0xbe:
			case 0xbf:
				// skip padding and navigation data 
				// (navigation shouldn't be present anyway)
				packet_length = currentPacket[4];
				packet_length <<= 8;
				packet_length += currentPacket[5];
				
				currentPacket += 6 + packet_length;
				break;
				
			case 0xbd:
				// a private stream packet, contains subtitle data
				packet_length = currentPacket[4];
				packet_length <<= 8;
				packet_length += currentPacket[5];
				
				int header_data_length = currentPacket[8];
				
				memcpy(&dest[copiedBytes], 
					   // header's 9 bytes + extension, we don't want 1st byte of packet
					   &currentPacket[9 + header_data_length + 1], 
					   // we don't want the 1-byte stream ID, or the header
					   packet_length - 1 - (header_data_length + 3));
				
				copiedBytes += packet_length - 1 - (header_data_length + 3);
				currentPacket += packet_length + 6;
				break;
				
			default:
				// unknown packet, probably video, return for now
				Codecprintf(NULL, "VobSubCodec - Unknown packet type %x, aborting\n", (int)currentPacket[3]);
				return;
		} // switch (currentPacket[3])
	} // while (currentPacket - framedSrc < srcSize)
}

ComponentResult ReadPacketControls(UInt8 *packet, UInt32 palette[16], PacketControlData *controlDataOut) {
	// to set whether the key sequences 0x03 - 0x06 have been seen
	UInt16 controlSeqSeen = 0;
	int i = 0;
	Boolean loop = TRUE;
	int controlOffset = (packet[2] << 8) + packet[3] + 4;
	uint8_t *controlSeq = packet + controlOffset;
	
	memset(controlDataOut, 0, sizeof(PacketControlData));
	
	while (loop) {
		switch (controlSeq[i]) {
			case 0x00:
				// subpicture identifier, we don't care
				i++;
				break;
				
			case 0x01:
				// start displaying, we don't care
				i++;
				break;
				
			case 0x03:
				// palette info
				controlDataOut->pixelColor[3] += palette[controlSeq[i+1] >> 4 ];
				controlDataOut->pixelColor[2] += palette[controlSeq[i+1] & 0xf];
				controlDataOut->pixelColor[1] += palette[controlSeq[i+2] >> 4 ];
				controlDataOut->pixelColor[0] += palette[controlSeq[i+2] & 0xf];
				
				i += 3;
				controlSeqSeen |= 0x0f;
				break;
				
			case 0x04:
				// alpha info
				controlDataOut->pixelColor[3] += (controlSeq[i + 1] & 0xf0) << 20;
				controlDataOut->pixelColor[2] += (controlSeq[i + 1] & 0x0f) << 24;
				controlDataOut->pixelColor[1] += (controlSeq[i + 2] & 0xf0) << 20;
				controlDataOut->pixelColor[0] += (controlSeq[i + 2] & 0x0f) << 24;
				
				// double the nibble
				controlDataOut->pixelColor[3] += (controlSeq[i + 1] & 0xf0) << 24;
				controlDataOut->pixelColor[2] += (controlSeq[i + 1] & 0x0f) << 28;
				controlDataOut->pixelColor[1] += (controlSeq[i + 2] & 0xf0) << 24;
				controlDataOut->pixelColor[0] += (controlSeq[i + 2] & 0x0f) << 28;
				
				i += 3;
				controlSeqSeen |= 0xf0;
				break;
				
			case 0x05:
				// coordinates of image, ffmpeg takes care of this
				i += 7;
				break;
				
			case 0x06:
				// offset of the first graphic line, and second, ffmpeg takes care of this
				i += 5;
				break;
				
			case 0xff:
				// end of control sequence
				loop = FALSE;
				break;
				
			default:
				Codecprintf(NULL, " !! Unknown control sequence 0x%02x  aborting (offset %x)\n", controlSeq[i], i);
				loop = FALSE;
				break;
		}
	}
	
	// force fully transparent to transparent black; needed? for graphicsModePreBlackAlpha
	for (i = 0; i < 4; i++) {
		if ((controlDataOut->pixelColor[i] & 0xff000000) == 0)
			controlDataOut->pixelColor[i] = 0;
	}
	
	if (controlSeqSeen != 0xff)
		return -1;
	return noErr;
}