//
// SaveSurfaceToFile.cpp
//
// Copyright (c) 2001 David Galeano
//
// Permission to use, copy, modify and distribute this software
// is hereby granted, provided that both the copyright notice and 
// this permission notice appear in all copies of the software, 
// derivative works or modified versions.
//
// THE AUTHOR ALLOWS FREE USE OF THIS SOFTWARE IN ITS "AS IS"
// CONDITION AND DISCLAIMS ANY LIABILITY OF ANY KIND FOR ANY DAMAGES 
// WHATSOEVER RESULTING FROM THE USE OF THIS SOFTWARE.
//

#define DIRECT3D_VERSION	0x0800

#include <d3d8.h>

#include <stdio.h>

//-----------------------------------------------------------------------------
//
// SaveSurfaceToTGAFile
//
//-----------------------------------------------------------------------------
HRESULT SaveSurfaceToTGAFile(LPDIRECT3DDEVICE8 pD3DDevice, const char *szFileName, LPDIRECT3DSURFACE8 pSurface)
{
	HRESULT hr=S_OK;

	D3DSURFACE_DESC d3dsd;
	LPDIRECT3DSURFACE8 pSurfaceCopy;

	pSurface->GetDesc(&d3dsd);

	hr=pD3DDevice->CreateImageSurface(d3dsd.Width, d3dsd.Height, d3dsd.Format, &pSurfaceCopy);

	if(SUCCEEDED(hr))
	{
		hr=pD3DDevice->CopyRects(pSurface, NULL, 0, pSurfaceCopy, NULL);

		// save the copy to the file
		D3DLOCKED_RECT d3dLR;
		if(SUCCEEDED(hr = pSurfaceCopy->LockRect(&d3dLR, NULL, D3DLOCK_NO_DIRTY_UPDATE|D3DLOCK_NOSYSLOCK|D3DLOCK_READONLY)))
		{
			// open file
			FILE *pFile=fopen(szFileName, "wb");
			if(pFile!=NULL)
			{
				setvbuf(pFile, NULL, _IOFBF, 64*1024);

				// write header
				char targaheader[18];

				// Set unused fields of header to 0
				memset(targaheader, 0, sizeof(targaheader));

				targaheader[2] = 2;	/* image type = uncompressed RGB */
				targaheader[12] = (char) (d3dsd.Width & 0xFF);
				targaheader[13] = (char) (d3dsd.Width >> 8);
				targaheader[14] = (char) (d3dsd.Height & 0xFF);
				targaheader[15] = (char) (d3dsd.Height >> 8);
				targaheader[16] = (d3dsd.Format==D3DFMT_A8R8G8B8 || d3dsd.Format==D3DFMT_A1R5G5B5 ? 32 : 24);
				targaheader[17] = 0x20;	/* Top-down, non-interlaced */

				fwrite(targaheader, 18, 1, pFile);

				// write file
				const BYTE *pImage=(const BYTE *)d3dLR.pBits;
				for(DWORD y=0; y<d3dsd.Height; y++)
				{
					if(d3dsd.Format==D3DFMT_X8R8G8B8)
					{
						const BYTE *pImageAux=pImage;
						for(DWORD x=0; x<d3dsd.Width; x++)
						{
							putc(pImageAux[0], pFile);
							putc(pImageAux[1], pFile);
							putc(pImageAux[2], pFile);
							pImageAux+=4;
						}
					}
					else if(d3dsd.Format==D3DFMT_A8R8G8B8)
					{
						fwrite(pImage, d3dsd.Width, 4, pFile);
					}
					else if(d3dsd.Format==D3DFMT_R5G6B5)
					{
						const BYTE *pImageAux=pImage;
						for(DWORD x=0; x<d3dsd.Width; x++)
						{
							const DWORD dwColor=*((WORD *)pImageAux);

							BYTE color;

							color=BYTE((dwColor&0x001f)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x7e0)>>5)<<2);
							putc(color, pFile);

							color=BYTE(((dwColor&0xf800)>>11)<<3);
							putc(color, pFile);

							pImageAux+=2;
						}
					}
					else if(d3dsd.Format==D3DFMT_X1R5G5B5)
					{
						const BYTE *pImageAux=pImage;
						for(DWORD x=0; x<d3dsd.Width; x++)
						{
							const DWORD dwColor=*((WORD *)pImageAux);

							BYTE color;

							color=BYTE((dwColor&0x001f)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x3e0)>>5)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x7c00)>>10)<<3);
							putc(color, pFile);

							pImageAux+=2;
						}
					}
					else if(d3dsd.Format==D3DFMT_A1R5G5B5)
					{
						const BYTE *pImageAux=pImage;
						for(DWORD x=0; x<d3dsd.Width; x++)
						{
							const DWORD dwColor=*((WORD *)pImageAux);

							BYTE color;

							color=BYTE((dwColor&0x001f)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x3e0)>>5)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x7c00)>>10)<<3);
							putc(color, pFile);

							color=BYTE(((dwColor&0x8000)>>15)*255);
							putc(color, pFile);

							pImageAux+=2;
						}
					}

					pImage+=d3dLR.Pitch;
				}

				fclose(pFile);
			}

		    pSurfaceCopy->UnlockRect();
		}

		pSurfaceCopy->Release();
	}

	return hr;
}

//-----------------------------------------------------------------------------
//
// SaveSurfaceToGrayscaleTGAFile
//
//-----------------------------------------------------------------------------
HRESULT SaveSurfaceToGrayscaleTGAFile(LPDIRECT3DDEVICE8 pD3DDevice, const char *szFileName, LPDIRECT3DSURFACE8 pSurface)
{
	HRESULT hr=S_OK;

	LPDIRECT3DSURFACE8 pSurfaceCopy=NULL;

	D3DSURFACE_DESC d3dsd;

	pSurface->GetDesc(&d3dsd);

	if(d3dsd.Pool!=D3DPOOL_SYSTEMMEM)
	{
		hr=pD3DDevice->CreateImageSurface(d3dsd.Width, d3dsd.Height, d3dsd.Format, &pSurfaceCopy);

		if(SUCCEEDED(hr))
		{
			hr=pD3DDevice->CopyRects(pSurface, NULL, 0, pSurfaceCopy, NULL);
		}
	}
	else
	{
		pSurfaceCopy=pSurface;
	}

	if(SUCCEEDED(hr))
	{
		// save the copy to the file
		D3DLOCKED_RECT d3dLR;
		if(SUCCEEDED(hr = pSurfaceCopy->LockRect(&d3dLR, NULL, D3DLOCK_NO_DIRTY_UPDATE|D3DLOCK_NOSYSLOCK|D3DLOCK_READONLY)))
		{
			// open file
			FILE *pFile=fopen(szFileName, "wb");
			if(pFile!=NULL)
			{
				setvbuf(pFile, NULL, _IOFBF, 64*1024);

				// write header
				char targaheader[18];

				// Set unused fields of header to 0
				memset(targaheader, 0, sizeof(targaheader));

				targaheader[2] = 3;		/* image type = uncompressed gray-scale */
				targaheader[12] = (char) (d3dsd.Width & 0xFF);
				targaheader[13] = (char) (d3dsd.Width >> 8);
				targaheader[14] = (char) (d3dsd.Height & 0xFF);
				targaheader[15] = (char) (d3dsd.Height >> 8);
				targaheader[16] = 8;
				targaheader[17] = 0x20;	/* Top-down, non-interlaced */

				fwrite(targaheader, 18, 1, pFile);

				// write file
				const BYTE *pImage=(const BYTE *)d3dLR.pBits;
				for(DWORD y=0; y<d3dsd.Height; y++)
				{
					const BYTE *pImageAux=pImage;
					for(DWORD x=0; x<d3dsd.Width; x++)
					{
						if(d3dsd.Format==D3DFMT_X8R8G8B8 || d3dsd.Format==D3DFMT_A8R8G8B8)
						{
							// save average greyscale
							putc((pImageAux[0]+pImageAux[1]+pImageAux[2])/3, pFile);
							pImageAux+=4;
						}
						else if(d3dsd.Format==D3DFMT_R5G6B5)
						{
							const DWORD dwColor=*((WORD *)pImageAux);

							DWORD colorAverage;

							colorAverage=((((dwColor&0xf800)>>11)<<3|0x07));
							colorAverage+=((((dwColor&0x7e0)>>5)<<2)|0x03);
							colorAverage+=(((dwColor&0x001f)<<3)|0x07);
							putc(colorAverage/3, pFile);

							pImageAux+=2;
						}
						else if(d3dsd.Format==D3DFMT_X1R5G5B5)
						{
							// save only green chanel
							const DWORD dwColor=*((WORD *)pImageAux);

							DWORD colorAverage;

							colorAverage=((((dwColor&0x7c00)>>10)<<3)|0x07);
							colorAverage+=((((dwColor&0x3e0)>>5)<<3)|0x07);
							colorAverage+=(((dwColor&0x001f)<<3)|0x07);
							putc(colorAverage/3, pFile);

							pImageAux+=2;
						}					
					}
					pImage+=d3dLR.Pitch;
				}

				fclose(pFile);
			}

		    pSurfaceCopy->UnlockRect();
		}
	}

	if(d3dsd.Pool!=D3DPOOL_SYSTEMMEM)
	{
		if(pSurfaceCopy!=NULL)
		{
			pSurfaceCopy->Release();
		}
	}

	return hr;
}

//-----------------------------------------------------------------------------
//
// CaptureScreenToFile
//
//-----------------------------------------------------------------------------
HRESULT CaptureScreenToFile(LPDIRECT3DDEVICE8 pD3DDevice)
{
	HRESULT hr=S_OK;

	LPDIRECT3DSURFACE8 pBackBuffer;

	if(SUCCEEDED(hr = pD3DDevice->GetRenderTarget(&pBackBuffer)))
	{
		char szName[MAX_PATH];

		// find a free filename
		DWORD nIndFile=0;
		do
		{
			sprintf(szName, "screen%02u.tga", nIndFile);
			nIndFile++;
		}
		while(GetFileAttributes(szName)!=DWORD(-1));

		SaveSurfaceToTGAFile(pD3DDevice, szName, pBackBuffer);

		pBackBuffer->Release();
	}

	return hr;
}