#include "AVIWriterC.h"
#include <windows.h>
#include <vfw.h>
#include "PajaTypes.h"
#include <string>

using namespace PajaTypes;


void
TRACE( const char* szFormat, ...  )
{
	char	szMsg[256];
	va_list	rList;
	va_start( rList, szFormat );
	_vsnprintf( szMsg, 255, szFormat, rList );
	va_end( rList );
	OutputDebugString( szMsg );
}


AVIWriterC::AVIWriterC() :
	m_i32Rate( 0 ),
	m_i32Frame( 0 ),
	m_pAVIFile( 0 ),
	m_pStream( 0 ),
	m_pStreamCompressed( 0 ),
	m_pStreamAudio( 0 ),
	m_i32SampleCountPerFrame( 0 ),
	m_i32SampleSize( 0 ),
	m_i32AudioFrame( 0 ),
	m_pAudioCompressionBuffer( 0 ),
	m_had( 0 ),
	m_hstr( 0 )
{
	memset( &m_rBIH, 0, sizeof( BITMAPINFOHEADER ) );
}

AVIWriterC::~AVIWriterC()
{
	close();
}

bool
AVIWriterC::open( const char* szName, int32 i32Width, int32 i32Height, int32 i32FPS, COMPVARS* pCompVars, int32 i32AudioChannels, int32 i32AudioFreq, HACMDRIVERID hAudioDriverId, WAVEFORMATEX* pCompressedFormat )
{

	m_i32Rate = i32FPS;

	memset( &m_rBIH,0, sizeof( BITMAPINFOHEADER ) );
	m_rBIH.biSize=sizeof( BITMAPINFOHEADER );
	m_rBIH.biWidth = i32Width;
	m_rBIH.biHeight = i32Height;
	m_rBIH.biPlanes = 1;
	m_rBIH.biBitCount = 24;
	m_rBIH.biSizeImage = m_rBIH.biWidth * m_rBIH.biHeight * 3;
	m_rBIH.biCompression = BI_RGB;		//BI_RGB means BRG in reality


	memset( &m_rWaveFormat, 0, sizeof( WAVEFORMATEX ) );
	m_rWaveFormat.wFormatTag = WAVE_FORMAT_PCM;
	m_rWaveFormat.nChannels = i32AudioChannels;
	m_rWaveFormat.nSamplesPerSec = i32AudioFreq;
	m_rWaveFormat.nAvgBytesPerSec = i32AudioFreq * 2 * i32AudioChannels;
	m_rWaveFormat.nBlockAlign = 2 * i32AudioChannels;
	m_rWaveFormat.wBitsPerSample = 16;
	m_rWaveFormat.cbSize = 0;


	AVISTREAMINFO			strHdr; // information for a single stream 
	AVISTREAMINFO			strAudioHdr; // information for a single stream 
	AVICOMPRESSOPTIONS		opts;
	AVICOMPRESSOPTIONS		audioOpts;
	AVICOMPRESSOPTIONS FAR* aopts[2] = { &opts, &audioOpts };
	PAVISTREAM	pStreams[2];
	int32	i32StreamCount = 0;


	std::string	sError;
	HRESULT		hr;

	// Step 0 : Let's make sure we are running on 1.1 
	DWORD wVer = HIWORD( VideoForWindowsVersion() );
	if( wVer < 0x010a ) {
		 // oops, we are too old, blow out of here 
		::MessageBox( NULL, "Version of Video for Windows too old.", "AVI Error", MB_OK | MB_ICONERROR );
		return false;
	}

	// Step 1 : initialize AVI engine
	AVIFileInit();

	// Step 2 : Open the movie file for writing....
	hr = AVIFileOpen( &m_pAVIFile,			// Address to contain the new file interface pointer
		       (LPCSTR)szName,				// Null-terminated string containing the name of the file to open
		       OF_WRITE | OF_CREATE,	    // Access mode to use when opening the file. 
		       NULL );						// use handler determined from file extension.
											// Name your file .avi -> very important

	if( hr != AVIERR_OK ) {
		sError = "Failed to open AVI file for writing.\n\n";
		// Check it succeded.
		switch( hr )
		{
		case AVIERR_BADFORMAT: 
			sError += "The file couldn't be read, indicating a corrupt file or an unrecognized format.";
			break;
		case AVIERR_MEMORY:		
			sError += "The file could not be opened because of insufficient memory."; 
			break;
		case AVIERR_FILEREAD:
			sError += "A disk error occurred while reading the file."; 
			break;
		case AVIERR_FILEOPEN:		
			sError += "A disk error occurred while opening the file.";
			break;
		case REGDB_E_CLASSNOTREG:		
			sError += "According to the registry, the type of file specified in AVIFileOpen does not have a handler to process it";
			break;
		}
		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return false;
	}

	// Fill in the header for the video stream....
	memset( &strHdr, 0, sizeof( strHdr ) );
	strHdr.fccType                = streamtypeVIDEO;	// video stream type
	strHdr.fccHandler             = 0;
	strHdr.dwScale                = 1;					// should be one for video
	strHdr.dwRate                 = i32FPS;		    // fps
	strHdr.dwSuggestedBufferSize  = 0; //m_rBIH.biSizeImage;	// Recommended buffer size, in bytes, for the stream.
	strHdr.dwQuality							= -1;
	SetRect( &strHdr.rcFrame, 0, 0,		    // rectangle for stream
	    (int) m_rBIH.biWidth,
	    (int) m_rBIH.biHeight);
	strcpy( strHdr.szName, "Video" );

	// Step 3 : Create the stream;
	hr = AVIFileCreateStream( m_pAVIFile,		    // file pointer
			         &m_pStream,		    // returned stream pointer
			         &strHdr );	    // stream header

	// Check it succeded.
	if( hr != AVIERR_OK) {
		sError = "AVI Stream creation failed. Check Bitmap info.";
		if( hr == AVIERR_READONLY ) {
			sError += " Read only file.";
		}
		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return false;
	}

	pStreams[i32StreamCount] = m_pStream;
	i32StreamCount++;



/*	memset( &opts, 0, sizeof( opts ) );
	memset( &audioOpts, 0, sizeof( audioOpts ) );
	// Poping codec dialog
	if( !AVISaveOptions( NULL, 0, i32StreamCount, pStreams, (LPAVICOMPRESSOPTIONS FAR *) &aopts) ) {
		AVISaveOptionsFree( i32StreamCount,(LPAVICOMPRESSOPTIONS FAR *) &aopts );
		return false;
	}*/

	//pCompVars

	memset( &opts, 0, sizeof( opts ) );
	opts.fccType = streamtypeVIDEO;
	opts.fccHandler = pCompVars->fccHandler;
	opts.dwKeyFrameEvery = pCompVars->lKey;
	opts.dwQuality = pCompVars->lQ;
	opts.dwBytesPerSecond = pCompVars->lDataRate * 1024;
	opts.dwFlags = 0;
	opts.lpFormat = pCompVars->lpbiOut;
	opts.cbFormat = sizeof( BITMAPINFOHEADER );
	opts.lpParms = 0;
	opts.cbParms = 0;
	opts.dwInterleaveEvery = 0;


	// Step 5:  Create a compressed stream using codec options.
	hr = AVIMakeCompressedStream( &m_pStreamCompressed, 
				m_pStream, 
				&opts, 
				NULL );

	if( hr != AVIERR_OK )
	{
		sError ="AVI Compressed Stream creation failed.";
		
		switch(hr)
		{
		case AVIERR_NOCOMPRESSOR:
			sError+=" A suitable compressor cannot be found.";
				break;
		case AVIERR_MEMORY:
			sError+=" There is not enough memory to complete the operation.";
				break; 
		case AVIERR_UNSUPPORTED:
			sError+="Compression is not supported for this type of data. This error might be returned if you try to compress data that is not audio or video.";
			break;
		}

		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return FALSE;
	}

	// Step 6 : sets the format of a stream at the specified position
	hr = AVIStreamSetFormat( m_pStreamCompressed, 
					0,			// position
					&m_rBIH,	    // stream format
					m_rBIH.biSize +   // format size
					m_rBIH.biClrUsed * sizeof( RGBQUAD ) );

	if( hr != AVIERR_OK ) {
		sError = "AVI Compressed Stream format setting failed.";
		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return FALSE;
	}


	if( hAudioDriverId && pCompressedFormat && i32AudioChannels > 0 )
	{
		// audio part
		memset( &strHdr, 0, sizeof( strHdr ) );
		strHdr.fccType                = streamtypeAUDIO;		// stream type
		strHdr.fccHandler             = 0;
		strHdr.dwScale                = m_rWaveFormat.nBlockAlign;
		strHdr.dwRate                 = i32AudioFreq * 2 * i32AudioChannels;		    // 25 fps
		strHdr.dwSuggestedBufferSize  = 0; //i32AudioFreq * 2 * i32AudioChannels;//format.bmiHeader.biSizeImage;
		strHdr.dwQuality= -1;
		strHdr.wPriority				= 7;
		strHdr.dwInitialFrames = 0.75 * i32AudioFreq;
		strHdr.dwSampleSize = m_rWaveFormat.nBlockAlign;

		// And create the stream;
		hr = AVIFileCreateStream( m_pAVIFile,		    // file pointer
			&m_pStreamAudio,		    // returned stream pointer
								&strHdr );	    // stream header
		if (hr != AVIERR_OK) {
			sError = "Error creating audio output memory";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}


		hr = AVIStreamSetFormat( m_pStreamAudio, 0,
							pCompressedFormat,	    // stream format
							sizeof( WAVEFORMATEX ) + sizeof( pCompressedFormat->cbSize ) );


		if (hr != AVIERR_OK) {
			sError = "Error setting audio output format.";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}


		m_i32SampleSize = 2 * i32AudioChannels;
		m_i32SampleCountPerFrame = (i32AudioFreq + (i32FPS - 1)) / i32FPS;


		// Create audio compressor

		MMRESULT mmr;
		WORD wFormatTag;
		HACMDRIVERID hadid;
		HACMDRIVER had = NULL;
		HACMSTREAM hstr = NULL;

		// open the driver
		mmr = acmDriverOpen( &m_had, hAudioDriverId, 0 );
		if( mmr ) {
			sError = "Error opening audio compressor.";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}

		// open the conversion stream
		// Note the use of the ACM_STREAMOPENF_NONREALTIME flag. Without this
		// some software compressors will report error 512 - not possible
		mmr = acmStreamOpen( &m_hstr,
			m_had,                    // driver handle
			&m_rWaveFormat,                 // source format
			pCompressedFormat,                 // destination format
			NULL,                   // no filter
			0,                      // no callback
			0,                      // instance data (not used)
			ACM_STREAMOPENF_NONREALTIME );   // flags

		if (mmr) {
			sError = "Error opening stream for audio compressor.";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}

		// Alloc compression buffer
		m_pAudioCompressionBuffer = new uint8[m_i32SampleCountPerFrame * m_i32SampleSize];
		if( !m_pAudioCompressionBuffer )
		{
			sError = "Error creating audio compression buffer.";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}
	}

	// releasing memory allocated by AVISaveOptionFree
/*	hr = AVISaveOptionsFree( i32StreamCount,(LPAVICOMPRESSOPTIONS FAR *) &aopts );
	if( hr != AVIERR_OK ) {
		sError = "Error releasing memory";
		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return false;
	}*/

	// Step 6 : Initialize step counter
	m_i32Frame = 0;

	return true;
}

bool
AVIWriterC::write_frame( uint8* pImageData, uint8* pAudioData, int32 i32AudioDataSize )
{
	std::string	sError;
	HRESULT		hr;

	// compress bitmap
	hr = AVIStreamWrite( m_pStreamCompressed,	// stream pointer
		m_i32Frame,						// time of this frame
		1,						// number to write
		pImageData,					// image buffer
		m_rBIH.biSizeImage,		// size of this frame
		AVIIF_KEYFRAME,			// flags....
		NULL,
		NULL );

	if( hr != AVIERR_OK ) {
		sError = "Unable to write data to stream.";
		::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
		return false;
	}

	if( m_pStreamAudio && pAudioData )
	{

    ACMSTREAMHEADER strhdr;
    MMRESULT				mmr;

    //
    // set up the in and out buffers
    //
    memset( &strhdr, 0, sizeof( strhdr ) );
    strhdr.cbStruct = sizeof( strhdr );
    strhdr.pbSrc = pAudioData;						// the source data to convert
    strhdr.cbSrcLength = i32AudioDataSize;    // and its size
    strhdr.pbDst = m_pAudioCompressionBuffer;						// ditto for the dest
    strhdr.cbDstLength = m_i32SampleCountPerFrame * m_i32SampleSize;

    // prep the header
    mmr = acmStreamPrepareHeader( m_hstr, &strhdr, 0 );
    if (mmr) {
			sError = "Error compressing audio (1).";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
    }
    // convert the data
    mmr = acmStreamConvert( m_hstr, &strhdr, 0 );
    if (mmr) {
			sError = "Error compressing audio (2).";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
    }
    // unprepare the header
    mmr = acmStreamUnprepareHeader( m_hstr, &strhdr, 0 );
    if (mmr) {
			sError = "Error compressing audio (3).";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
    }

		hr = AVIStreamWrite( m_pStreamAudio,	// stream pointer
			m_i32AudioFrame,						// time of this frame
			m_i32SampleCountPerFrame,						// number to write
			m_pAudioCompressionBuffer,					// image buffer
			strhdr.cbDstLengthUsed,
			0,			// flags....
			NULL,
			NULL );

		m_i32AudioFrame += m_i32SampleCountPerFrame;
		if( hr != AVIERR_OK ) {
			sError = "Unable to write audio data to stream.";
			::MessageBox( NULL, sError.c_str(), "AVI Error", MB_OK | MB_ICONERROR );
			return false;
		}
	}

	// updating frame counter
	m_i32Frame++;

	return true;
}

void
AVIWriterC::close()
{
	if( m_pStreamAudio ) {
		AVIStreamRelease( m_pStreamAudio );
		m_pStreamAudio = NULL;
	}

	if( m_pStream ) {
		AVIStreamRelease( m_pStream );
		m_pStream = NULL;
	}

	if( m_pStreamCompressed ) {
		AVIStreamRelease( m_pStreamCompressed );
		m_pStreamCompressed = NULL;
	}

	if( m_pAVIFile ) {
		AVIFileRelease( m_pAVIFile );
		m_pAVIFile = NULL;
	}

	m_i32Rate = 0;
	m_i32Frame = 0;
	m_i32SampleCountPerFrame = 0;
	m_i32SampleSize = 0;
	m_i32AudioFrame = 0;

	if( m_hstr )
		acmStreamClose( m_hstr, 0 );
	if( m_had )
		acmDriverClose( m_had, 0 );

	delete [] m_pAudioCompressionBuffer;
	m_pAudioCompressionBuffer = 0;

	// Close engine
	AVIFileExit();
}


int32
AVIWriterC::get_audio_bytes_per_frame() const
{
	return m_i32SampleCountPerFrame * m_i32SampleSize;
}
