//
// Tools_AVI_Writer.cpp
//
// Created by Richard Rice - 08/05/2002
// 
// Copyright (C) 2002 Rebellion, All Rights Reserved.
//

#include "Core/Asura.h"
#include "Rendering/Asura_Viewport.h"
#include "Rendering/Asura_Direct3D.h"
#include "Core/Asura_Windows.h"
#include "Debugging/Asura_Assert.h"
#include "Core/Asura_Timers.h"

#ifdef ASURA_SUPPORT_COMMAND_CONSOLE
#include "Interface/Asura_CommandConsole.h"
#include "Interface/Asura_CommandConsole_IO.h"
#endif


#include "Tools_System.h"
#include "Misc/Tools_AVI_Writer.h"
#include "Misc/Tools_Timed_Messages.h"
#include "PC/OS/Tools_Os.h"
#include "PC/FileSystem/Tools_FileSystem.h"

Tools_AVI_Writer TheAviWriter;


Tools_AVI_Writer::Tools_AVI_Writer()
{
	AviFile = NULL;
	AviStream = NULL;
	AviStreamCompressed = NULL;
	AviOptionsArray[0] = &AviCompressOptions;
	AviBuffer = 0;
	Frame = 0;
	m_iBitDepth = 24;
	m_bWriteBmps = false;
	m_pxTempSurface = NULL;
	m_eSurfaceFormat = D3DFMT_UNKNOWN;

	memset(&AviCompressOptions, 0, sizeof(AviCompressOptions));

}

Tools_AVI_Writer::~Tools_AVI_Writer()
{
	EndAVI();
}

bool Tools_AVI_Writer::StartAVI(const char* filename,float fFrameRate)
{
	return StartAVI(filename,fFrameRate,Asura_Viewport::GetWidth(),Asura_Viewport::GetHeight(),0);
}

bool Tools_AVI_Writer::StartAVI(const char* filename,float fFrameRate,int iWidth,int iHeight,u_int uFlags)
{
	EndAVI();//stop any previous recording

	m_bHalfSize = (uFlags & TOOLS_AVIFLAG_HALFSIZE)!= 0;
	bool bUsePreviousSettings = (uFlags & TOOLS_AVIFLAG_USEPREVIOUSCOMPRESSIONSETTINGS) !=0;
	m_bWriteBmps = (uFlags & TOOLS_AVIFLAG_WRITEBMPS)!=0;
	int iScale = m_bHalfSize ? 2 : 1;
	
	if(iWidth>Asura_Viewport::GetWidth()/iScale || iHeight > Asura_Viewport::GetHeight()/iScale)
	{
		Asura_Assert(iWidth<=Asura_Viewport::GetWidth()/iScale && iHeight <= Asura_Viewport::GetHeight()/iScale,"");
		return false;
	}

	if(!CreateTempSurface()) return false;

	
	m_fFrameRate = fFrameRate;
	Asura_Timers::SetFixedGameTime(1.0f/m_fFrameRate);

	Tools_FileSystem::EnsureDirectoryExists(filename);
	
	Asura_Assert(strlen(filename)<MAX_PATH-10,"");//some leeway to allow for extra stuff being tacked on the end
	strcpy(m_szFilename,filename);
	m_pszPartPosition = 0;
	m_iPart = 0;

	m_iPitch = ((iWidth*m_iBitDepth/8)+3)&~3;
	

	BitmapInfo.biSize = sizeof(BITMAPINFOHEADER) ;
	BitmapInfo.biWidth = iWidth ;
	BitmapInfo.biHeight = iHeight ;
	BitmapInfo.biPlanes = 1 ;
	BitmapInfo.biBitCount = (unsigned short)m_iBitDepth ;
	BitmapInfo.biCompression = BI_RGB ;
	BitmapInfo.biSizeImage = BitmapInfo.biHeight*m_iPitch;
	BitmapInfo.biXPelsPerMeter = 0 ;
	BitmapInfo.biYPelsPerMeter = 0 ;
	BitmapInfo.biClrUsed = 0;
	BitmapInfo.biClrImportant = 0 ;

	AviBuffer = new char[BitmapInfo.biSizeImage];

	if(m_bWriteBmps)
	{
		//remove extension from filename
		char* pcDotPos = strrchr(m_szFilename,'.');
		if(pcDotPos) *pcDotPos = 0;

		//if the filename ends with a frame sequence number , remove that as well
		char* pcUnderscore = strrchr(m_szFilename,'_');
		if(pcUnderscore)
		{
			//check to see if all a remaining characters are numbers
			char* pcPos = pcUnderscore + 1;
			int iCount = 0;
			while(*pcPos)
			{
				if(*pcPos<'0' || *pcPos > '9')
				{
					//not a number
					break;
				}
				pcPos++;
				iCount++;
			}

			if(*pcPos == 0 && iCount>=4)
			{
				//ends with four or more digits , so call it a frame number , and remove it
				*pcUnderscore = 0;
			}
		}


		OutputMessage(0xffffffff,"Writing %s_XXXX.bmp",m_szFilename);
	}
	else
	{
		//make sure extension is correct
		char* pcDotPos = strrchr(m_szFilename,'.');
		if(pcDotPos) *pcDotPos = 0;
		strcat(m_szFilename,".avi");

		
		if(!OpenFile(bUsePreviousSettings))
		{
			OutputMessage(0xffffffff,"Failed to create %s",m_szFilename);
			EndAVI();
			return false;
		}

				
		OutputMessage(0xffffffff,"Writing %s",m_szFilename);
	}
	return true;
}

bool Tools_AVI_Writer::OpenFile(bool bUsePreviousSettings)
{
	if(m_bWriteBmps) return true;//nothing to do here if writing bmps

	HRESULT hr;
	
	AVIFileInit();

	hr = AVIFileOpen(&AviFile,		    // returned file pointer
		       m_szFilename,		            // file name
		       OF_WRITE | OF_CREATE,	    // mode to open file with
		       NULL);			    // use handler determined

	if (hr != AVIERR_OK)
	{
		return false;
	}

	//fill in the stream information
	memset(&AviStreamInfo, 0, sizeof(AviStreamInfo));
	AviStreamInfo.fccType                = streamtypeVIDEO;// stream type
	AviStreamInfo.fccHandler             = 0;
	AviStreamInfo.dwScale                = 1000;
	AviStreamInfo.dwRate                 = Asura_Maths::FloatToInt(m_fFrameRate*1000.0f);
	AviStreamInfo.dwSuggestedBufferSize  = BitmapInfo.biSizeImage;
	SetRect(&AviStreamInfo.rcFrame, 0, 0,		    // rectangle for stream
	    (int) BitmapInfo.biWidth,
	    (int) BitmapInfo.biHeight);

	// And create the stream;
	hr = AVIFileCreateStream(AviFile,		    // file pointer
			         &AviStream,		    // returned stream pointer
			         &AviStreamInfo);	    // stream header
	if (hr != AVIERR_OK) 
	{
		return false;
	}

	bool bOkay = true;

	if(!bUsePreviousSettings)
	{
		//wait for enter to be released , so that it desn't result in the options dialog
		//being closed immediately.
		while(Asura_UserInput::KeyPressed(ASURA_KEYID_ENTER))
		{
			Asura_Windows::CheckForWindowsMessages();
			Asura_UserInput::Update();
		}
		
		bOkay = AVISaveOptions(Asura_Windows::GetWindowHandle(), 0, 1, &AviStream, (LPAVICOMPRESSOPTIONS FAR *) &AviOptionsArray) != 0;
	}
	if(!bOkay)
	{
		return false;
	}

	hr = AVIMakeCompressedStream(&AviStreamCompressed, AviStream, &AviCompressOptions, NULL);
	if (hr != AVIERR_OK) 
	{
		return false;
	}
	

	
	hr = AVIStreamSetFormat(AviStreamCompressed, 0,
			       &BitmapInfo,	    // stream format
			       BitmapInfo.biSize +   // format size
			       BitmapInfo.biClrUsed * sizeof(RGBQUAD));
	if (hr != AVIERR_OK) 
	{
		
		return false;
	}

	Frame = 0;
	m_iFileSize = 0;
	m_iMaxFrameSize = 0;

	return true;
}



void Tools_AVI_Writer::WriteFrame()
{
	
	bool success = false;

	if(!AviBuffer) return;
	if(!m_pxTempSurface) return;

	//make sure screen resolution hasn't changed
	int iScale = m_bHalfSize ? 2 : 1;
	if(Asura_Viewport::GetWidth()/iScale<BitmapInfo.biWidth || Asura_Viewport::GetHeight()/iScale<BitmapInfo.biHeight)
	{
		EndAVI();
		return;
	}

	//deal with files that get too large
	if(!StartNewFileIfRequired())
	{
		EndAVI();
		return;
	}

	
	Direct3D_Surface* backbuffer = 0;
	HRESULT hres;
	hres = Asura_Direct3D::D3DDevice->GetBackBuffer(0,0,D3DBACKBUFFER_TYPE_MONO,&backbuffer);
	if(FAILED(hres))
	{
		Asura_Assert(0,"Failed to get backbuffer for avi");
		goto Cleanup;
	}

	hres = Asura_Direct3D::D3DDevice->StretchRect(backbuffer, NULL, m_pxTempSurface, NULL, D3DTEXF_POINT);
	ASURA_D3D_OK( hres);
	
	D3DLOCKED_RECT rect;
	hres = m_pxTempSurface->LockRect(&rect,NULL,D3DLOCK_READONLY);
	if(FAILED(hres))
	{
		Asura_Assert(0,"Failed to lock buffer for avi");
		goto Cleanup;
	}
	
	//determine the pixel format
	int red_shift,red_scale,red_mask;
	int green_shift,green_scale,green_mask;
	int blue_shift,blue_scale,blue_mask;
	int bytesPerPixel;

	switch(m_eSurfaceFormat)
	{
		case D3DFMT_X1R5G5B5 :
			red_mask = 31<<10;
			red_shift = 10;
			red_scale = 255/31;
			green_mask = 31<<5;
			green_shift = 5;
			green_scale = 255/31;
			blue_mask = 31;
			blue_shift = 0;
			blue_scale = 255/31;
			bytesPerPixel = 2;
			break;
		case D3DFMT_R5G6B5 :
			red_mask = 31<<11;
			red_shift = 11;
			red_scale = 255/31;
			green_mask = 63<<5;
			green_shift = 5;
			green_scale = 255/63;
			blue_mask = 31;
			blue_shift = 0;
			blue_scale = 255/31;
			bytesPerPixel = 2;
			break;
		case D3DFMT_X8R8G8B8 :
		case D3DFMT_A8R8G8B8 :
			red_mask = 255<<16;
			red_shift = 16;
			red_scale = 1;
			green_mask = 255<<8;
			green_shift = 8;
			green_scale = 1;
			blue_mask = 255;
			blue_shift = 0;
			blue_scale = 1;
			bytesPerPixel = 4;
			break;

		default :
			Asura_Assert(0,"Unexpected pixel format for screenshot");
			goto Cleanup;

	}

	//if the avi being written is smaller than the current resolution , then we need to take the top left corner indented a bit
	int iXOffset,iYOffset;
	iXOffset = (Asura_Viewport::GetWidth()-BitmapInfo.biWidth*iScale)/2;
	iYOffset = (Asura_Viewport::GetHeight()-BitmapInfo.biHeight*iScale)/2;
	
	// write 24 bit image
	unsigned char* BufferPtr,*BufferPtr2;
	unsigned char* pucDestLineStart;
	BufferPtr = (unsigned char *)rect.pBits+rect.Pitch*(iYOffset+iScale*BitmapInfo.biHeight-1);
	BufferPtr += iXOffset*bytesPerPixel;
	BufferPtr2 = BufferPtr-rect.Pitch;
	pucDestLineStart = (unsigned char*)AviBuffer;

	for (int i=BitmapInfo.biHeight-1; i>=0; --i)
	{
		int j;
		unsigned char* pucDestination = pucDestLineStart;
	
		{
			for (j=0; j<BitmapInfo.biWidth; ++j)
			{
				int r=0;
				int g=0;
				int b=0;

				if(m_bHalfSize)
				{
					//average out the 4 pixels to get the colour to use
					unsigned int auColour[4];
					if(bytesPerPixel == 4)
					{
						auColour[0]=*(unsigned int*)&BufferPtr[(j*2)*4];
						auColour[1]=*(unsigned int*)&BufferPtr[(j*2+1)*4];
						auColour[2]=*(unsigned int*)&BufferPtr2[(j*2)*4];
						auColour[3]=*(unsigned int*)&BufferPtr2[(j*2+1)*4];
					}
					else
					{
						auColour[0]=*(unsigned short*)&BufferPtr[(j*2)*2];
						auColour[1]=*(unsigned short*)&BufferPtr[(j*2+1)*2];
						auColour[2]=*(unsigned short*)&BufferPtr2[(j*2)*2];
						auColour[3]=*(unsigned short*)&BufferPtr2[(j*2+1)*2];
					}

					for(int k = 0;k<4;k++)
					{
						b += (((auColour[k]&blue_mask)>>blue_shift)*blue_scale);
						g += (BYTE)(((auColour[k]&green_mask)>>green_shift)*green_scale);
						r += (BYTE)(((auColour[k]&red_mask)>>red_shift)*red_scale);
					}
					r/=4;
					g/=4;
					b/=4;
				}
				else
				{
					unsigned int colour;
					if(bytesPerPixel == 4)
					{
						colour=*(unsigned int*)&BufferPtr[j*4];
					}
					else
					{
						colour=*(unsigned short*)&BufferPtr[j*2];
					}
					b = (((colour&blue_mask)>>blue_shift)*blue_scale);
					g = (BYTE)(((colour&green_mask)>>green_shift)*green_scale);
					r = (BYTE)(((colour&red_mask)>>red_shift)*red_scale);
				}

				switch(m_iBitDepth)
				{
					case 16 :
						r>>=3;
						g>>=3;
						b>>=3;
						*(unsigned short*)pucDestination = (unsigned short)((r<<10)|(g<<5)|(b));
						pucDestination+=2;
						break;
					case 24 :
						*pucDestination++ = (unsigned char)b;
						*pucDestination++ = (unsigned char)g;
						*pucDestination++ = (unsigned char)r;
						break;
					default :
						Asura_Assert(0,"Unexpected bit depth");
						break;
				}
			}
		}
	
		if(m_bHalfSize)
		{
			BufferPtr-=2*rect.Pitch;
			BufferPtr2-=2*rect.Pitch;
		}
		else
		{
			BufferPtr-=rect.Pitch;
		}
		pucDestLineStart += m_iPitch;
	}
	
	
	if(m_bWriteBmps)
	{
		success = WriteBmpFrame();
	}
	else
	{
		LONG iBytesWritten;
		
		//write out the frame		
		success = !AVIStreamWrite(	AviStreamCompressed,	// stream pointer
									Frame,				// time of this frame
									1,				// number to write
									(LPBYTE) AviBuffer,		// pointer to data
									BitmapInfo.biSizeImage,	// size of this frame
									AVIIF_KEYFRAME,			 // flags....
									NULL,
									&iBytesWritten);
	
		//keep track of file size
		if(iBytesWritten > m_iMaxFrameSize) m_iMaxFrameSize = iBytesWritten;
		m_iFileSize += iBytesWritten;
	}
	Frame++;

	

	if(!success)
	{
		const char* szFileType = m_bWriteBmps ? "BMP" : "AVI";
		OutputMessage(0xff0000,"Error writing %s file :",szFileType);
		//Not entirely clear whether AVIStreamWrite's errors are accessed using GetLastError()
		//I'll give it a go and see...
		OutputMessage(0xff0000,"%s",Tools_OS::GetLastSystemError());
	}

	Cleanup :
	m_pxTempSurface->UnlockRect();
	if(backbuffer)
	{
		backbuffer->Release();
	}

	if(!success)
	{
		EndAVI();
	}
}


bool Tools_AVI_Writer::StartNewFileIfRequired()
{
	if(m_bWriteBmps) return true;//nothing to do here if writing bmps

	const int iMaxFileSize = 0x7fffffff;

	//allow frame to be written if there is enough space for another frame , plus some extra to allow
	//for any header information
	if(m_iFileSize < iMaxFileSize - m_iMaxFrameSize - 10000) return true;

	//close existing file
	CloseFile();

	//if this is the first part ,then rename it
	if(m_iPart==0)
	{
		char m_szOriginalFilename[MAX_PATH];
		strcpy(m_szOriginalFilename,m_szFilename);

		//remove extension
		char* pcDot = strrchr(m_szFilename,'.');
		if(pcDot) *pcDot = 0;

		//remember where the file postfix is going to be put
		m_pszPartPosition = &m_szFilename[strlen(m_szFilename)];

		sprintf(m_pszPartPosition,"_Part%02d.avi",m_iPart);

		//rename the existing file.
		if(Tools_FileSystem::FileExists(m_szFilename))
		{
			if(!DeleteFile(m_szFilename))
			{
				OutputMessage(0xff0000,"Failed to delete %s",m_szFilename);
				return false;
			}
		}
		MoveFile(m_szOriginalFilename,m_szFilename);
	}

	//generate name for next part
	m_iPart++;
	sprintf(m_pszPartPosition,"_Part%02d.avi",m_iPart);

	//and open it
	bool bOkay = OpenFile(true);
	if(bOkay)
	{
		OutputMessage(0xffffffff,"Writing %s",m_szFilename);
	}
	else
	{
		OutputMessage(0xffffffff,"Failed to create %s",m_szFilename);
	}

	return bOkay;


}


void Tools_AVI_Writer::EndAVI()
{
	CloseFile();

	if(m_pxTempSurface)
	{
		m_pxTempSurface->Release();
		m_pxTempSurface = NULL;
	}

	delete [] AviBuffer;
	AviBuffer = 0;

	Asura_Timers::SetFixedGameTime(0);
}


void Tools_AVI_Writer::CloseFile()
{
	if (AviStreamCompressed && AviStreamCompressed!=AviStream)
	{
		AVIStreamClose(AviStreamCompressed);
	}
	AviStreamCompressed = 0;

	if (AviStream)
	{
		AVIStreamClose(AviStream);
	}
	AviStream = 0;


	if (AviFile)
	{
		AVIFileClose(AviFile);
	}
	AviFile = 0;

	AVIFileExit();

	Frame = 0;
}



void Tools_AVI_Writer::OutputMessage(int iColour,const char* pszString,...)
{
	char szBuffer[255];

	// format string and place in buffer
	va_list ap;
	va_start(ap, pszString);
    Asura_PrintfToBfr( szBuffer, pszString, ap, 255, true );
	va_end(ap);

	if(Tools_System::IsActive())
	{
		Tools_Timed_Message_System::AddMessage(iColour,"%s",szBuffer);
	}
	#ifdef ASURA_SUPPORT_COMMAND_CONSOLE
	else
	{
		//not in tools.. so send to command console instead
		Asura_CommandConsole::Printf(L"%S",szBuffer);
	}
	#endif

}

bool Tools_AVI_Writer::WriteBmpFrame()
{
	//get filename for current frame ,and attempt to open the file
	char szBmpFilename[MAX_PATH];
	Asura_PrintfToBfr(szBmpFilename,"%s_%04d.bmp",MAX_PATH,true,m_szFilename,Frame);

	HANDLE hFile = CreateFile(szBmpFilename,GENERIC_WRITE,0,0,CREATE_ALWAYS,FILE_ATTRIBUTE_NORMAL,0);
	if( hFile == INVALID_HANDLE_VALUE) return false;

    #define WRITE_AND_CHECK( src_ptr, num_bytes ) \
    { \
    	DWORD uNumBytesWritten; \
        BOOL bOK = WriteFile( hFile, (src_ptr), (num_bytes), &uNumBytesWritten, NULL ); \
		if(!bOK) \
		{ \
			CloseHandle(hFile); \
			return false; \
		} \
    }
	
	//write the header
	u_short usType = 'B'+'M'*256;
	WRITE_AND_CHECK(&usType,sizeof(u_short));

	u_int uFilesize = 14 + sizeof(BitmapInfo) + BitmapInfo.biSizeImage;
	WRITE_AND_CHECK(&uFilesize,sizeof(u_int));

	u_int uReserved = 0;
	WRITE_AND_CHECK(&uReserved,sizeof(u_int));

	u_int uImageOffset = 14 + sizeof(BitmapInfo);
	WRITE_AND_CHECK(&uImageOffset,sizeof(u_int));

	WRITE_AND_CHECK(&BitmapInfo,sizeof(BitmapInfo));

	//and the image
	WRITE_AND_CHECK(AviBuffer,BitmapInfo.biSizeImage);

	CloseHandle(hFile);

	#undef WRITE_AND_CHECK

	return true;
}

void Tools_AVI_Writer::PreD3DReset()
{
	if(m_pxTempSurface)
	{
		m_pxTempSurface->Release();
		m_pxTempSurface = 0;
	}
}

void Tools_AVI_Writer::PostD3DReset()
{
	if(IsWritingFile())
	{
		if(!CreateTempSurface())
		{
			EndAVI();
		}
	}
}

bool Tools_AVI_Writer::CreateTempSurface()
{
	Asura_Assert(m_pxTempSurface == 0,"");

	Direct3D_Surface* backbuffer = 0;
	HRESULT hres;
	hres = Asura_Direct3D::D3DDevice->GetBackBuffer(0,0,D3DBACKBUFFER_TYPE_MONO,&backbuffer);
	if(SUCCEEDED(hres))
	{
		D3DSURFACE_DESC desc;
		hres = backbuffer->GetDesc(&desc);
		if(SUCCEEDED(hres))
		{
			m_eSurfaceFormat = desc.Format;

			// Create a lockable RT surface to copy to
			hres = Asura_Direct3D::D3DDevice->CreateRenderTarget(desc.Width, desc.Height, desc.Format, D3DMULTISAMPLE_NONE, 0, TRUE, &m_pxTempSurface, NULL);
			Asura_Assert(SUCCEEDED(hres), "Failed to create temporary AVI surface");
		}
		else
		{
			Asura_Assert(0, "Failed to get backbuffer desc");
		}
		backbuffer->Release();
	}
	else
	{
		Asura_Assert(0, "Failed to get backbuffer");
	}
	if(!m_pxTempSurface) return false;

	return true;
}
