/******************************************************************************
 * SAGE - Scalable Adaptive Graphics Environment
 *
 * Copyright (C) 2004 Electronic Visualization Laboratory,
 * University of Illinois at Chicago
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above
 *    copyright notice, this list of conditions and the following disclaimer
 *    in the documentation and/or other materials provided with the distribution.
 *  * Neither the name of the University of Illinois at Chicago nor
 *    the names of its contributors may be used to endorse or promote
 *    products derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * Direct questions, comments etc about SAGE to http://www.evl.uic.edu/cavern/forum/
 *
 *****************************************************************************/

#include "StdAfx.h"
#include "avifile.h"

/// FPS defines the Number of Frames per second;
/// Here we set FPS to 1 meaning one Frame per Second.
/// Increase it it you need more Frames Per Second.
#define FPS			1	// Frames Per Second

/// VIDEOCODEC defines the Codec that need to be used for the Movie.
/// Change the definition to use a different codec. For example,
/// Change it to mmioFOURCC('D','I','V','X') to use DIVX codec; or
/// Change it to mmioFOURCC('I','V','5','0') to use IV50 codec etc..
/// Also, you can change it to 0 to avoid the codecs altogether. In that case, Frames would be saved as they are;
/// However, make sure you have the codec installed on the machine before using its Fourcc here.
#define VIDEOCODEC	mmioFOURCC('M','P','G','4')

//#define VIDEOCODEC	0	/// Frames will be saved as they are - movie size would be HUGE !!

CAviFile::CAviFile(LPCTSTR lpszFileName)
{

	AVIFileInit();

	_tcscpy(m_szFileName,lpszFileName);

	m_hHeap=NULL;
	m_hAviDC=NULL;
	m_lpBits=NULL;
	m_lSample=NULL;
	m_pAviFile=NULL;
	m_pAviStream=NULL;
	m_pAviCompressedStream=NULL;

	pAppendFrame[0]= &CAviFile::AppendDummy;			// VC8 requires & for Function Pointer; Remove it if your compiler complains;
	pAppendFrame[1]= &CAviFile::AppendFrameFirstTime;
	pAppendFrame[2]= &CAviFile::AppendFrameUsual;

	pAppendFrameBits[0]=&CAviFile::AppendDummy;
	pAppendFrameBits[1]=&CAviFile::AppendFrameFirstTime;
	pAppendFrameBits[2]=&CAviFile::AppendFrameUsual;

	nAppendFuncSelector=1;		//0=Dummy	1=FirstTime	2=Usual
}

CAviFile::~CAviFile(void)
{
	ReleaseMemory();

	AVIFileExit();
}

void CAviFile::ReleaseMemory()
{
	nAppendFuncSelector=0;		//Point to DummyFunction

	if(m_hAviDC)
	{
		DeleteDC(m_hAviDC);
		m_hAviDC=NULL;
	}
	if(m_pAviCompressedStream)
	{
		AVIStreamRelease(m_pAviCompressedStream);
		m_pAviCompressedStream=NULL;
	}
	if(m_pAviStream)
	{
		AVIStreamRelease(m_pAviStream);
		m_pAviStream=NULL;
	}
	if(m_pAviFile)
	{
		AVIFileRelease(m_pAviFile);
		m_pAviFile=NULL;
	}
	if(m_lpBits)
	{
		HeapFree(m_hHeap,HEAP_NO_SERIALIZE,m_lpBits);
		m_lpBits=NULL;
	}
	if(m_hHeap)
	{
		HeapDestroy(m_hHeap);
		m_hHeap=NULL;
	}
}

HRESULT	CAviFile::AppendFrameFirstTime(HBITMAP	hBitmap)
{
	int	nMaxWidth=GetSystemMetrics(SM_CXSCREEN),nMaxHeight=GetSystemMetrics(SM_CYSCREEN);

	BITMAPINFO	bmpInfo;

	m_hAviDC=CreateCompatibleDC(NULL);
	if(m_hAviDC==NULL)
	{
		//MessageBox(NULL,"Unable to Create Compatible DC","Error",MB_OK|MB_ICONERROR);
		goto TerminateInit;
	}

	ZeroMemory(&bmpInfo,sizeof(BITMAPINFO));
	bmpInfo.bmiHeader.biSize=sizeof(BITMAPINFOHEADER);

	GetDIBits(m_hAviDC,hBitmap,0,0,NULL,&bmpInfo,DIB_RGB_COLORS);

	bmpInfo.bmiHeader.biCompression=BI_RGB;

	if(bmpInfo.bmiHeader.biHeight>nMaxHeight)	nMaxHeight=bmpInfo.bmiHeader.biHeight;
	if(bmpInfo.bmiHeader.biWidth>nMaxWidth)	nMaxWidth=bmpInfo.bmiHeader.biWidth;

	m_hHeap=HeapCreate(HEAP_NO_SERIALIZE,nMaxWidth*nMaxHeight*4,0);
	if(m_hHeap==NULL)
	{
		//MessageBox(NULL,"Unable to Allocate Memory","Error",MB_OK);
		goto TerminateInit;
	}

	m_lpBits=HeapAlloc(m_hHeap,HEAP_ZERO_MEMORY|HEAP_NO_SERIALIZE,nMaxWidth*nMaxHeight*4);
	if(m_lpBits==NULL)
	{
		//MessageBox(NULL,"Unable to Allocate Memory","Error",MB_OK);
		goto TerminateInit;
	}

	if(FAILED(AVIFileOpen(&m_pAviFile,m_szFileName,OF_CREATE|OF_WRITE,NULL)))
	{
		//MessageBox(NULL,"Unable to Create the Movie File","Error",MB_OK|MB_ICONERROR);
		goto TerminateInit;
	}

	ZeroMemory(&m_AviStreamInfo,sizeof(AVISTREAMINFO));
	m_AviStreamInfo.fccType=streamtypeVIDEO;
	m_AviStreamInfo.fccHandler=VIDEOCODEC;
	m_AviStreamInfo.dwScale=1;
	m_AviStreamInfo.dwRate=FPS;		//Frames Per Second;
	m_AviStreamInfo.dwQuality=-1;	//Default Quality
	m_AviStreamInfo.dwSuggestedBufferSize=nMaxWidth*nMaxHeight*4;
    SetRect(&m_AviStreamInfo.rcFrame,0,0,bmpInfo.bmiHeader.biWidth,bmpInfo.bmiHeader.biHeight);
	strcpy(m_AviStreamInfo.szName,"Video Stream");

	if(FAILED(AVIFileCreateStream(m_pAviFile,&m_pAviStream,&m_AviStreamInfo)))
	{
		//MessageBox(NULL,"Unable to Create Stream","Error",MB_OK|MB_ICONERROR);
		goto TerminateInit;

	}

	ZeroMemory(&m_AviCompressOptions,sizeof(AVICOMPRESSOPTIONS));
	m_AviCompressOptions.fccType=streamtypeVIDEO;
	m_AviCompressOptions.fccHandler=m_AviStreamInfo.fccHandler;
	m_AviCompressOptions.dwFlags=AVICOMPRESSF_KEYFRAMES|AVICOMPRESSF_VALID;//|AVICOMPRESSF_DATARATE;
	m_AviCompressOptions.dwKeyFrameEvery=15;
	//m_AviCompressOptions.dwBytesPerSecond=1000/8;
	//m_AviCompressOptions.dwQuality=100;

	if(FAILED(AVIMakeCompressedStream(&m_pAviCompressedStream,m_pAviStream,&m_AviCompressOptions,NULL)))
	{
		//MessageBox(NULL,"Unable to Create Compressed Stream","Error",MB_OK);
		goto TerminateInit;
	}

	if(FAILED(AVIStreamSetFormat(m_pAviCompressedStream,0,(LPVOID)&bmpInfo,bmpInfo.bmiHeader.biSize)))
	{
		//MessageBox(NULL,"Unable to Set Format","Error",MB_OK);
		goto TerminateInit;
	}

	nAppendFuncSelector=2;		//Point to UsualAppend Function

	return AppendFrameUsual(hBitmap);

TerminateInit:

	ReleaseMemory();	MessageBox(NULL,"Error Occured While Rendering the Movie","Error",MB_OK|MB_ICONERROR);

	return E_FAIL;
}

HRESULT CAviFile::AppendFrameUsual(HBITMAP hBitmap)
{
	BITMAPINFO	bmpInfo;

	bmpInfo.bmiHeader.biBitCount=0;
	bmpInfo.bmiHeader.biSize=sizeof(BITMAPINFOHEADER);

	GetDIBits(m_hAviDC,hBitmap,0,0,NULL,&bmpInfo,DIB_RGB_COLORS);

	bmpInfo.bmiHeader.biCompression=BI_RGB;

	GetDIBits(m_hAviDC,hBitmap,0,bmpInfo.bmiHeader.biHeight,m_lpBits,&bmpInfo,DIB_RGB_COLORS);

	if(FAILED(AVIStreamWrite(m_pAviCompressedStream,m_lSample++,1,m_lpBits,bmpInfo.bmiHeader.biSizeImage,0,NULL,NULL)))
		return E_FAIL;

    return S_OK;
}

HRESULT CAviFile::AppendDummy(HBITMAP)
{
	return E_FAIL;
}

HRESULT CAviFile::AppendNewFrame(HBITMAP hBitmap)
{
	return (this->*pAppendFrame[nAppendFuncSelector])((HBITMAP)hBitmap);
}

HRESULT	CAviFile::AppendNewFrame(int nWidth, int nHeight, LPVOID pBits,int nBitsPerPixel)
{
	return (this->*pAppendFrameBits[nAppendFuncSelector])(nWidth,nHeight,pBits,nBitsPerPixel);
}

HRESULT	CAviFile::AppendFrameFirstTime(int nWidth, int nHeight, LPVOID pBits,int nBitsPerPixel)
{
	int	nMaxWidth=GetSystemMetrics(SM_CXSCREEN),nMaxHeight=GetSystemMetrics(SM_CYSCREEN);

	BITMAPINFO	bmpInfo;

	m_hAviDC=CreateCompatibleDC(NULL);
	if(m_hAviDC==NULL)
	{
		//MessageBox(NULL,"Unable to Create Compatible DC","Error",MB_OK|MB_ICONERROR);
		goto TerminateInitBits;
	}

	ZeroMemory(&bmpInfo,sizeof(BITMAPINFO));
	bmpInfo.bmiHeader.biPlanes=1;
	bmpInfo.bmiHeader.biWidth=nWidth;
	bmpInfo.bmiHeader.biHeight=nHeight;
	bmpInfo.bmiHeader.biCompression=BI_RGB;
	bmpInfo.bmiHeader.biBitCount=nBitsPerPixel;
	bmpInfo.bmiHeader.biSize=sizeof(BITMAPINFOHEADER);
	bmpInfo.bmiHeader.biSizeImage=bmpInfo.bmiHeader.biWidth*bmpInfo.bmiHeader.biHeight*bmpInfo.bmiHeader.biBitCount/8;

	if(bmpInfo.bmiHeader.biHeight>nMaxHeight)	nMaxHeight=bmpInfo.bmiHeader.biHeight;
	if(bmpInfo.bmiHeader.biWidth>nMaxWidth)	nMaxWidth=bmpInfo.bmiHeader.biWidth;

	m_hHeap=HeapCreate(HEAP_NO_SERIALIZE,nMaxWidth*nMaxHeight*4,0);
	if(m_hHeap==NULL)
	{
		//MessageBox(NULL,"Unable to Allocate Memory","Error",MB_OK);
		goto TerminateInitBits;
	}

	m_lpBits=HeapAlloc(m_hHeap,HEAP_ZERO_MEMORY|HEAP_NO_SERIALIZE,nMaxWidth*nMaxHeight*4);
	if(m_lpBits==NULL)
	{
		//MessageBox(NULL,"Unable to Allocate Memory","Error",MB_OK);
		goto TerminateInitBits;
	}

	if(FAILED(AVIFileOpen(&m_pAviFile,m_szFileName,OF_CREATE|OF_WRITE,NULL)))
	{
		//MessageBox(NULL,"Unable to Create the Movie File","Error",MB_OK|MB_ICONERROR);
		goto TerminateInitBits;
	}

	ZeroMemory(&m_AviStreamInfo,sizeof(AVISTREAMINFO));
	m_AviStreamInfo.fccType=streamtypeVIDEO;
	m_AviStreamInfo.fccHandler=VIDEOCODEC;
	m_AviStreamInfo.dwScale=1;
	m_AviStreamInfo.dwRate= FPS;	//Frames Per Second;
	m_AviStreamInfo.dwQuality=-1;	//Default Quality
	m_AviStreamInfo.dwSuggestedBufferSize=nMaxWidth*nMaxHeight*4;
    SetRect(&m_AviStreamInfo.rcFrame,0,0,bmpInfo.bmiHeader.biWidth,bmpInfo.bmiHeader.biHeight);
	strcpy(m_AviStreamInfo.szName,"Video Stream");

	if(FAILED(AVIFileCreateStream(m_pAviFile,&m_pAviStream,&m_AviStreamInfo)))
	{
		//MessageBox(NULL,"Unable to Create Stream","Error",MB_OK|MB_ICONERROR);
		goto TerminateInitBits;

	}

	ZeroMemory(&m_AviCompressOptions,sizeof(AVICOMPRESSOPTIONS));
	m_AviCompressOptions.fccType=streamtypeVIDEO;
	m_AviCompressOptions.fccHandler=m_AviStreamInfo.fccHandler;
	m_AviCompressOptions.dwFlags=AVICOMPRESSF_KEYFRAMES|AVICOMPRESSF_VALID;//|AVICOMPRESSF_DATARATE;
	m_AviCompressOptions.dwKeyFrameEvery=15;
	//m_AviCompressOptions.dwBytesPerSecond=1000/8;
	//m_AviCompressOptions.dwQuality=100;

	if(FAILED(AVIMakeCompressedStream(&m_pAviCompressedStream,m_pAviStream,&m_AviCompressOptions,NULL)))
	{
		//MessageBox(NULL,"Unable to Create Compressed Stream","Error",MB_OK);
		goto TerminateInitBits;
	}

	if(FAILED(AVIStreamSetFormat(m_pAviCompressedStream,0,(LPVOID)&bmpInfo,bmpInfo.bmiHeader.biSize)))
	{
		//MessageBox(NULL,"Unable to Set Format","Error",MB_OK);
		goto TerminateInitBits;
	}

	nAppendFuncSelector=2;		//Point to UsualAppend Function

	return AppendFrameUsual(nWidth,nHeight,pBits,nBitsPerPixel);

TerminateInitBits:

	ReleaseMemory();	MessageBox(NULL,"Error Occured While Rendering the Movie","Error",MB_OK|MB_ICONERROR);

	return E_FAIL;

}

HRESULT	CAviFile::AppendFrameUsual(int nWidth, int nHeight, LPVOID pBits,int nBitsPerPixel)
{
	DWORD	dwSize=nWidth*nHeight*nBitsPerPixel/8;

	if(FAILED(AVIStreamWrite(m_pAviCompressedStream,m_lSample++,1,pBits,dwSize,0,NULL,NULL)))
		return E_FAIL;

    return S_OK;
}

HRESULT	CAviFile::AppendDummy(int nWidth, int nHeight, LPVOID pBits,int nBitsPerPixel)
{
	return E_FAIL;
}