/* 
 *	Copyright (C) Chia-chen Kuo - Jan 2001
 *
 *  This file is part of DVD2AVI, a free MPEG-2 decoder
 *	
 *  DVD2AVI is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2, or (at your option)
 *  any later version.
 *   
 *  DVD2AVI is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *  GNU General Public License for more details.
 *   
 *  You should have received a copy of the GNU General Public License
 *  along with GNU Make; see the file COPYING.  If not, write to
 *  the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA. 
 *
 */

#include "global.h"

#define MAX_AVI_SIZE	2073600000

__forceinline static void Store_RGB24(unsigned char *src[], DWORD frame);
__forceinline static void Store_YUY2(unsigned char *src[], DWORD frame);

static void Luminance_Filter(unsigned char *src, unsigned char *dst);
static void conv420to422(unsigned char *src, unsigned char *dst);
static void conv422to444(unsigned char *src, unsigned char *dst);
static void conv444toRGB24odd(unsigned char *py, unsigned char *pu, unsigned char *pv, unsigned char *dst);
static void conv444toRGB24even(unsigned char *py, unsigned char *pu, unsigned char *pv, unsigned char *dst);
static void Flush_RGB24();
static void conv422toyuy2odd(unsigned char *py, unsigned char *pu, unsigned char *pv, unsigned char *dst);
static void conv422toyuy2even(unsigned char *py, unsigned char *pu, unsigned char *pv, unsigned char *dst);
static void Flush_YUY2();
static void AVIKill();

PAVIFILE pfile;
PAVISTREAM ps, psCompressed;
AVICOMPRESSOPTIONS opts;
LPAVICOMPRESSOPTIONS lpopts = &opts;
AVISTREAMINFO strinfo;
ICCOMPRESSFRAMES iccf;
COMPVARS compvars;
PCOMPVARS pcompvars = &compvars;

static unsigned char *y444;
static int AVI_Init, avi_size, avi_count, frame_count;
static int TFF, RFF, TFB, BFB, playback, Old_Playback, frame_size, frame_type;
static char VideoOut[_MAX_PATH];

static char *FrameType[3] = {
	"", "Interlaced", "Progressive"
};

static const __int64 mmmask_0001 = 0x0001000100010001;
static const __int64 mmmask_0002 = 0x0002000200020002;
static const __int64 mmmask_0003 = 0x0003000300030003;
static const __int64 mmmask_0004 = 0x0004000400040004;
static const __int64 mmmask_0005 = 0x0005000500050005;
static const __int64 mmmask_0007 = 0x0007000700070007;
static const __int64 mmmask_0064 = 0x0040004000400040;
static const __int64 mmmask_0128 = 0x0080008000800080;
static const __int64 mmmask_cbu = 0x0000408D0000408D;
static const __int64 mmmask_cgu_cgv = 0xF377E5FCF377E5FC;
static const __int64 mmmask_crv = 0x0000331300003313;

void Write_Frame(unsigned char *src[], D2VData d2v, DWORD frame)
{
	int i, repeat;

	if (Fault_Flag)
	{
		if (Fault_Flag < CRITICAL_ERROR_LEVEL)
		{
			SetDlgItemText(hDlg, IDC_DEBUG, "V.E.");
			Fault_Flag = 0;		// Fault Tolerance
		}
		else
		{
			if (AVI_Flag)
				AVIKill();
			ThreadKill();
		}
	}

	frame_type = d2v.pf;
	TFF = d2v.trf>>1;
	RFF = d2v.trf & 0x01;

	if (!frame)
	{
		AVI_Init = 1; TFB = BFB = 0; Old_Playback = 0;
		frame_size = 0; avi_size = 0; avi_count = 1; frame_count = 0; playback = 0;

		if (process.locate==LOCATE_RIP && FO_Flag==FO_SWAP)
		{
			if (TFF)
				BFB = 1;
			else
				TFB = 1;

			ZeroMemory(rgb24, Coded_Picture_Width*Coded_Picture_Height*3);
			for (i=0; i<Coded_Picture_Width*Coded_Picture_Height*2; i += 2)
			{
				yuy2[i] = 0;
				yuy2[i+1] = 128;
			}
		}

		if (FO_Flag!=FO_FILM)
		{
			if ((TFF && FO_Flag!=FO_SWAP) || (!TFF && FO_Flag==FO_SWAP))
				SetDlgItemText(hDlg, IDC_DEBUG, "T");
			else
				SetDlgItemText(hDlg, IDC_DEBUG, "B");
		}
	}

	repeat = DetectVideoType(frame, d2v.trf);

	if (FO_Flag!=FO_FILM || repeat)
	{
		if (Store_Flag==STORE_YUY2)
			Store_YUY2(src, frame);
		else
			Store_RGB24(src, frame);
	}

	if (FO_Flag==FO_FILM && repeat==2)
	{
		if (Store_Flag==STORE_YUY2)
			Store_YUY2(src, frame);
		else
			Store_RGB24(src, frame);
	}

	__asm emms;

	if (Statistics_Flag && process.locate==LOCATE_RIP)
	{
		sprintf(szBuffer, "%s", FrameType[frame_type+1]);
		SetDlgItemText(hDlg, IDC_FRAME_TYPE, szBuffer);

		sprintf(szBuffer, "%d", frame+1);
		SetDlgItemText(hDlg, IDC_CODED_NUMBER, szBuffer);

		sprintf(szBuffer, "%d", playback);
		SetDlgItemText(hDlg, IDC_PLAYBACK_NUMBER, szBuffer);

		if (AVI_Flag)
		{
			sprintf(szBuffer, "%d", avi_count-1);
			SetDlgItemText(hDlg, IDC_FILE, szBuffer);
	
			sprintf(szBuffer, "%d MB", avi_size>>20);
			SetDlgItemText(hDlg, IDC_FILE_SIZE, szBuffer);
		}

		if ((frame & 63) == 63)
		{
			process.ed = timeGetTime();
			sprintf(szBuffer, "%.2f FPS", 1000.0*(playback-1)/(process.ed-process.op+1));
			SetDlgItemText(hDlg, IDC_FPS, szBuffer);

			SendDlgItemMessage(hDlg, IDC_PROGRESS, PBM_SETPOS, (int)(Bitrate_Meter*Frame_Rate/(playback-Old_Playback))>>7, 0);

			Bitrate_Meter = 0;
			Old_Playback = playback;
		}
	}

	if (AVI_Flag && avi_size>=MAX_AVI_SIZE)
		AVIKill();
}

static void Store_RGB24(unsigned char *src[], DWORD frame)
{
	// store AVI
	if (AVI_Flag && AVI_Init)
	{
		if (!frame)
		{
			compvars.cbSize = sizeof(compvars);

			if (!ICCompressorChoose(hWnd, ICMF_CHOOSE_DATARATE | ICMF_CHOOSE_KEYFRAME,
				lpbirgb, NULL, pcompvars, NULL))
				ThreadKill();

			// set AVI header
			ZeroMemory(&strinfo, sizeof(AVISTREAMINFO));
			strinfo.fccType					= streamtypeVIDEO;
			strinfo.fccHandler				= compvars.fccHandler;
			strinfo.dwQuality				= -1;
			strinfo.dwSuggestedBufferSize	= birgb.biSizeImage;
			strinfo.dwScale	= 1000;
			strinfo.dwRate = (unsigned int)((FO_Flag==FO_FILM) ? frame_rate*800 : frame_rate*1000);

			// set AVI save options
			opts.cbFormat = sizeof(birgb);
			opts.fccType = streamtypeVIDEO;
			opts.fccHandler = compvars.fccHandler;
			opts.dwKeyFrameEvery = compvars.lKey;
			opts.dwQuality = compvars.lQ;
			opts.dwBytesPerSecond = compvars.lDataRate<<10;
			opts.dwFlags = AVICOMPRESSF_VALID | AVICOMPRESSF_KEYFRAMES | AVICOMPRESSF_DATARATE;
			opts.dwInterleaveEvery = 0;
			opts.lpFormat = lpbirgb;
			opts.cbParms = compvars.cbState;
			opts.lpParms = compvars.lpState;
		}

		AVI_Init = 0; avi_size = 0;
		AVIFileInit();

		sprintf(VideoOut, "%s %02d.avi", szOutput, avi_count++);

		if (AVIFileOpen(&pfile, VideoOut, OF_WRITE | OF_CREATE, NULL) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIFileCreateStream(pfile, &ps, &strinfo) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIMakeCompressedStream(&psCompressed, ps, lpopts, NULL) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIStreamSetFormat(psCompressed, 0, lpbirgb, birgb.biSize) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}
	}

	if (chroma_format==CHROMA420)
	{
		conv420to422(src[1], u422);
		conv420to422(src[2], v422);

		conv422to444(u422, u444);
		conv422to444(v422, v444);
	}
	else
	{
		conv422to444(src[1], u444);
		conv422to444(src[2], v444);
	}

	if (Luminance_Flag)
	{
		Luminance_Filter(src[0], lum);
		y444 = lum;
	}
	else
		y444 = src[0];

	if (BFB)
	{
		conv444toRGB24odd(y444, u444, v444, rgb24);
		Flush_RGB24();
		conv444toRGB24even(y444, u444, v444, rgb24);
		Flush_RGB24();
	}
	else
	{
		conv444toRGB24even(y444, u444, v444, rgb24);
		Flush_RGB24();
		conv444toRGB24odd(y444, u444, v444, rgb24);
		Flush_RGB24();
	}

	if (FO_Flag!=FO_FILM && RFF)
		if (TFF)
		{
			conv444toRGB24odd(y444, u444, v444, rgb24);
			Flush_RGB24();
		}
		else
		{
			conv444toRGB24even(y444, u444, v444, rgb24);
			Flush_RGB24();
		}
}

static void Flush_RGB24()
{
	if (TFB & BFB)
	{
		if (AVI_Flag)
		{
			__asm emms;

			if (AVIStreamWrite(psCompressed, frame_count++, 1, rgb24,
				birgb.biSizeImage, 0, NULL, &frame_size) != AVIERR_OK)
			{
				AVIKill();
				ThreadKill();
			}

			avi_size += frame_size;
		}

		if (Display_Flag)
			RenderRGB24();

		playback++;
		TFB = BFB = 0;
	}
}

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

static void Store_YUY2(unsigned char *src[], DWORD frame)
{
	// Save AVI
	if (AVI_Flag && AVI_Init)
	{
		if (!frame)
		{
			compvars.cbSize = sizeof(compvars);

			if (!ICCompressorChoose(hWnd, ICMF_CHOOSE_DATARATE | ICMF_CHOOSE_KEYFRAME,
				lpbiyuv, NULL, pcompvars, NULL))
				ThreadKill();

			// set AVI header
			ZeroMemory(&strinfo, sizeof(AVISTREAMINFO));
			strinfo.fccType					= streamtypeVIDEO;
			strinfo.fccHandler				= compvars.fccHandler;
			strinfo.dwQuality				= -1;
			strinfo.dwSuggestedBufferSize	= biyuv.biSizeImage;
			strinfo.dwScale	= 1000;
			strinfo.dwRate = (unsigned int)((FO_Flag==FO_FILM) ? frame_rate*800 : frame_rate*1000);

			// set AVI save options
			opts.cbFormat = sizeof(biyuv);
			opts.fccType = streamtypeVIDEO;
			opts.fccHandler = compvars.fccHandler;
			opts.dwKeyFrameEvery = compvars.lKey;
			opts.dwQuality = compvars.lQ;
			opts.dwBytesPerSecond = compvars.lDataRate<<10;
			opts.dwFlags = AVICOMPRESSF_VALID | AVICOMPRESSF_KEYFRAMES | AVICOMPRESSF_DATARATE;
			opts.dwInterleaveEvery = 0;
			opts.lpFormat = lpbiyuv;
			opts.cbParms = compvars.cbState;
			opts.lpParms = compvars.lpState;

			iccf.dwRate = strinfo.dwRate;
			iccf.dwScale = strinfo.dwScale;
			iccf.lQuality = compvars.lQ;
			iccf.lDataRate = compvars.lDataRate<<10;
			iccf.lKeyRate = compvars.lKey;

			ICSendMessage(compvars.hic, ICM_COMPRESS_FRAMES_INFO, (WPARAM)&iccf, (DWORD)sizeof(ICCOMPRESSFRAMES));
		}

		AVI_Init = 0; avi_size = 0;
		AVIFileInit();

		sprintf(VideoOut, "%s %02d.avi", szOutput, avi_count++);

		if (AVIFileOpen(&pfile, VideoOut, OF_WRITE | OF_CREATE, NULL) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIFileCreateStream(pfile, &ps, &strinfo) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIMakeCompressedStream(&psCompressed, ps, lpopts, NULL) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (AVIStreamSetFormat(psCompressed, 0, lpbirgb, birgb.biSize) != AVIERR_OK)
		{
			AVIKill();
			ThreadKill();
		}

		if (!frame && !ICSeqCompressFrameStart(pcompvars, (LPBITMAPINFO)lpbiyuv))
		{
			AVIKill();
			ThreadKill();
		}
	}

	if (chroma_format==CHROMA420)
	{
		conv420to422(src[1], u422);
		conv420to422(src[2], v422);
	}
	else
	{
		u422 = src[1];
		v422 = src[2];
	}

	if (Luminance_Flag)
	{
		Luminance_Filter(src[0], lum);
		y444 = lum;
	}
	else
		y444 = src[0];

	if (BFB)
	{
		conv422toyuy2odd(y444, u422, v422, yuy2);
		Flush_YUY2();
		conv422toyuy2even(y444, u422, v422, yuy2);
		Flush_YUY2();
	}
	else
	{
		conv422toyuy2even(y444, u422, v422, yuy2);
		Flush_YUY2();

		conv422toyuy2odd(y444, u422, v422, yuy2);
		Flush_YUY2();
	}

	if (FO_Flag!=FO_FILM && RFF)
		if (TFF)
		{
			conv422toyuy2odd(y444, u422, v422, yuy2);
			Flush_YUY2();
		}
		else
		{
			conv422toyuy2even(y444, u422, v422, yuy2);
			Flush_YUY2();
		}
}

static void Flush_YUY2()
{
	void *yuy2c; int key_flag, yuy2c_size;

	if (TFB & BFB)
	{
		if (AVI_Flag)
		{
			__asm emms;

			if ((yuy2c = ICSeqCompressFrame(pcompvars, 0, yuy2, &key_flag, &yuy2c_size)) == NULL)
			{
				AVIKill();
				ThreadKill();
			}

			if (AVIStreamWrite(ps, frame_count++, 1, yuy2c, yuy2c_size,
				key_flag ? AVIIF_KEYFRAME : 0, NULL, &frame_size) != AVIERR_OK)
			{
				AVIKill();
				ThreadKill();
			}

			avi_size += frame_size;
		}

		if (DDOverlay_Flag && Display_Flag)
			RenderYUY2();

		playback++;
		TFB = BFB = 0;
	}
}

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

static void Luminance_Filter(unsigned char *src, unsigned char *dst)
{
	int AREA = Coded_Picture_Width * Coded_Picture_Height;

	__asm
	{
		mov			eax, [src]
		mov			ebx, [dst]
		mov			esi, 0x00
		mov			edi, [AREA]
		pxor		mm0, mm0
		movq		mm5, [LumOffsetMask]
		movq		mm6, [LumGainMask]
		movq		mm7, [mmmask_0064]

lumconv:
		movq		mm1, [eax+esi]
		movq		mm2, mm1

		punpcklbw	mm1, mm0
		punpckhbw	mm2, mm0

		pmullw		mm1, mm6
		pmullw		mm2, mm6

		paddw		mm1, mm7
		paddw		mm2, mm7

		psrlw		mm1, 7
		psrlw		mm2, 7

		paddw		mm1, mm5
		paddw		mm2, mm5

		packuswb	mm1, mm0
		packuswb	mm2, mm0

		add			esi, 0x08
		cmp			esi, edi
		movq		[ebx+esi-8], mm1
		movq		[ebx+esi-4], mm2
		jl			lumconv
	}
}

static void conv422to444(unsigned char *src, unsigned char *dst)
{
	int RIGHT_BOUND = (Coded_Picture_Width>>1) - 8;
	int SRC_STRIDE = Coded_Picture_Width>>1;

	__asm
	{
		mov			eax, [src]
		mov			ebx, [dst]
		mov			edi, [Coded_Picture_Height]

		movq		mm1, [mmmask_0001]
		pxor		mm0, mm0

convyuv444init:
		movq		mm7, [eax]
		mov			esi, 0x00

convyuv444:
		movq		mm2, mm7
		movq		mm7, [eax+esi+8]
		movq		mm3, mm2
		movq		mm4, mm7

		psrlq		mm3, 8
		psllq		mm4, 56
		por			mm3, mm4

		movq		mm4, mm2
		movq		mm5, mm3

		punpcklbw	mm4, mm0
		punpcklbw	mm5, mm0

		movq		mm6, mm4
		paddusw		mm4, mm1
		paddusw		mm4, mm5
		psrlw		mm4, 1
		psllq		mm4, 8
		por			mm4, mm6

		punpckhbw	mm2, mm0
		punpckhbw	mm3, mm0

		movq		mm6, mm2
		paddusw		mm2, mm1
		paddusw		mm2, mm3

		movq		[ebx+esi*2], mm4

		psrlw		mm2, 1
		psllq		mm2, 8
		por			mm2, mm6

		add			esi, 0x08
		cmp			esi, [RIGHT_BOUND]
		movq		[ebx+esi*2-8], mm2
		jl			convyuv444

		movq		mm2, mm7
		punpcklbw	mm2, mm0
		movq		mm3, mm2

		psllq		mm2, 8
		por			mm2, mm3

		movq		[ebx+esi*2], mm2

		punpckhbw	mm7, mm0
		movq		mm6, mm7

		psllq		mm6, 8
		por			mm6, mm7

		movq		[ebx+esi*2+8], mm6

		add			eax, [SRC_STRIDE]		
		add			ebx, [Coded_Picture_Width]
		sub			edi, 0x01
		cmp			edi, 0x00
		jg			convyuv444init
	}
}

static void conv420to422(unsigned char *src, unsigned char *dst)
{
	int PROGRESSIVE_HEIGHT = (Coded_Picture_Height>>1) - 2;
	int INTERLACED_HEIGHT = (Coded_Picture_Height>>2) - 2;
	int HALF_WIDTH = Coded_Picture_Width>>1;
	int DOUBLE_WIDTH = Coded_Picture_Width<<1;

	if (frame_type)
	{
		__asm
		{
			mov			eax, [src]
			mov			ebx, [dst]
			mov			ecx, ebx
			add			ecx, [HALF_WIDTH]
			mov			esi, 0x00
			movq		mm3, [mmmask_0003]
			pxor		mm0, mm0
			movq		mm4, [mmmask_0002]

			mov			edx, eax
			add			edx, [HALF_WIDTH]
convyuv422topp:
			movd		mm1, [eax+esi]
			movd		mm2, [edx+esi]
			movd		[ebx+esi], mm1
			punpcklbw	mm1, mm0
			pmullw		mm1, mm3
			paddusw		mm1, mm4
			punpcklbw	mm2, mm0
			paddusw		mm2, mm1
			psrlw		mm2, 0x02
			packuswb	mm2, mm0

			add			esi, 0x04
			cmp			esi, [HALF_WIDTH]
			movd		[ecx+esi-4], mm2
			jl			convyuv422topp

			add			eax, [HALF_WIDTH]
			add			ebx, [Coded_Picture_Width]
			add			ecx, [Coded_Picture_Width]
			mov			esi, 0x00

			mov			edi, [PROGRESSIVE_HEIGHT]
convyuv422p:
			movd		mm1, [eax+esi]

			punpcklbw	mm1, mm0
			mov			edx, eax

			pmullw		mm1, mm3
			sub			edx, [HALF_WIDTH]

			movd		mm5, [edx+esi]
			movd		mm2, [edx+esi]

			punpcklbw	mm5, mm0
			punpcklbw	mm2, mm0
			paddusw		mm5, mm1
			paddusw		mm2, mm1
			paddusw		mm5, mm4
			paddusw		mm2, mm4
			psrlw		mm5, 0x02
			psrlw		mm2, 0x02
			packuswb	mm5, mm0
			packuswb	mm2, mm0

			mov			edx, eax
			add			edx, [HALF_WIDTH]
			add			esi, 0x04
			cmp			esi, [HALF_WIDTH]
			movd		[ebx+esi-4], mm5
			movd		[ecx+esi-4], mm2

			jl			convyuv422p

			add			eax, [HALF_WIDTH]
			add			ebx, [Coded_Picture_Width]
			add			ecx, [Coded_Picture_Width]
			mov			esi, 0x00
			sub			edi, 0x01
			cmp			edi, 0x00
			jg			convyuv422p

			mov			edx, eax
			sub			edx, [HALF_WIDTH]
convyuv422bottomp:
			movd		mm1, [eax+esi]
			movd		mm5, [edx+esi]
			punpcklbw	mm5, mm0
			movd		[ecx+esi], mm1

			punpcklbw	mm1, mm0
			pmullw		mm1, mm3
			paddusw		mm5, mm1
			paddusw		mm5, mm4
			psrlw		mm5, 0x02
			packuswb	mm5, mm0