//
// Tools_Zlib.cpp
//
// Created by Richard Rice - 21/08/2007
// 
// Copyright (C) 2007 Rebellion, All Rights Reserved.
//

#include "Core/Asura.h"
#include "Load&Save/Tools_Zlib.h"

#ifdef ASURA_USE_ZLIB

//asura
#include "ChunkSystem/Asura_Chunk_Loaders.h"
#include "ChunkSystem/Asura_Zlib.h"
#include "Core/Asura_Endian.h"
#include "Core/Asura_Maths.h"


//middleware
#include "zlib/zlib.h"

//tools
#include "PC/FileSystem/Tools_FileSystem.h"

#ifdef ASURA_MAIN
#include "Load&Save/LevelExport/Tools_Level_Export.h"
#endif



static const u_int g_uInputBufferSize = 16384;
static const u_int g_uOutputBufferSize = 16384;

class Tools_Zlib_Compressor
{
public :
	Tools_Zlib_Compressor();
	~Tools_Zlib_Compressor();

	bool SaveCompressedFile(const void* pBuffer, const u_int uSize, const char* szFilename, bool bBigEndian);
	bool SaveCompressedFile(const char* szSourceFile, const char* szDestFile, bool bBigEndian, bool bAllowAobrt);

private :
	bool SaveCompressedFile(const char* szFilename, bool bBigEndian, bool bAllowAbort);
	void Cleanup();

	//if compressing from memory
	const void* m_pSourceBuffer;
	u_int m_uBufferSize;

	//if compressing from a file
	HANDLE m_pSourceFile;


	char* m_pcInputBuffer;
	char* m_pcOutputBuffer;

	Asura_String m_xFilename;
	HANDLE m_pDestFile;

	z_stream m_xZlibStream;
	bool m_bStreamOpened;

};


Tools_Zlib_Compressor::Tools_Zlib_Compressor()
{
	m_pSourceBuffer = 0;
	m_uBufferSize = 0;
	m_pSourceFile = INVALID_HANDLE_VALUE;
	m_pDestFile = INVALID_HANDLE_VALUE;
	m_bStreamOpened = false;
	m_pcInputBuffer = 0;
	m_pcOutputBuffer = new char[g_uOutputBufferSize];
}

Tools_Zlib_Compressor::~Tools_Zlib_Compressor()
{
	delete [] m_pcOutputBuffer;
}

void Tools_Zlib_Compressor::Cleanup()
{
	if(m_pSourceFile != INVALID_HANDLE_VALUE)
	{
		CloseHandle(m_pSourceFile);
		m_pSourceFile = INVALID_HANDLE_VALUE;
	}

	if(m_pDestFile != INVALID_HANDLE_VALUE)
	{
		CloseHandle(m_pDestFile);
		m_pDestFile = INVALID_HANDLE_VALUE;
		//delete the partially saved file
		DeleteFile(m_xFilename.GetString());
	}

	if(m_bStreamOpened)
	{
		deflateEnd(&m_xZlibStream);
		m_bStreamOpened = false;
	}
	delete [] m_pcInputBuffer;
	m_pcInputBuffer = 0;

	m_pSourceBuffer = 0;
	m_uBufferSize = 0;
}

#define WRITE_FILE(pBuffer,iSize)\
WriteFile(m_pDestFile,(long*)pBuffer,(iSize),&uBytesWritten,0);\
if(uBytesWritten!=(unsigned)(iSize))\
	return false;

bool Tools_Zlib_Compressor::SaveCompressedFile(const void* pBuffer, const u_int uSize, const char* szFilename,bool bBigEndian)
{
	m_pSourceBuffer = pBuffer;
	m_uBufferSize = uSize;
	
	bool bOkay = SaveCompressedFile(szFilename, bBigEndian, false);
	Cleanup();

	return bOkay;
}

bool Tools_Zlib_Compressor::SaveCompressedFile(const char* szSourceFile, const char* szDestFile, bool bBigEndian, bool bAllowAbort)
{
	m_pSourceFile = CreateFile(szSourceFile,GENERIC_READ,FILE_SHARE_READ,0,OPEN_EXISTING,0,0);

	if(m_pSourceFile == INVALID_HANDLE_VALUE) return false;

	bool bOkay = SaveCompressedFile(szDestFile, bBigEndian, bAllowAbort);
	Cleanup();

	return bOkay;
}



bool Tools_Zlib_Compressor::SaveCompressedFile(const char* szFilename,bool bBigEndian, bool bAllowAbort)
{
	DWORD uBytesWritten;

	m_xFilename = szFilename;
	
	m_pDestFile = CreateFile(m_xFilename.GetString(),GENERIC_WRITE,0,0,CREATE_ALWAYS,0,0);
	if(m_pDestFile == INVALID_HANDLE_VALUE)
	{
		return false;
	}

	//write the file ID
	WRITE_FILE(ASURA_ZLIB_COMPRESSED_FILE_ID, ASURA_ZLIB_COMPRESSED_FILE_ID_ACTUAL_CHARS);

	//leave space for our header
	Asura_Zlib_Header xHeader;
	WRITE_FILE(&xHeader, sizeof(Asura_Zlib_Header))

	//initialise the compression stream
    m_xZlibStream.zalloc = Z_NULL;
    m_xZlibStream.zfree = Z_NULL;
    m_xZlibStream.opaque = Z_NULL;

	int iRet = deflateInit2(&m_xZlibStream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, ASURA_ZLIB_DEFAULT_WINDOW_BITS, 9, Z_DEFAULT_STRATEGY);
	if(iRet != Z_OK)
	{
		return false;
	}
	m_bStreamOpened = true;

	u_int uSourceBytesRemaining = 0;
	u_int uSize = 0;
	if(m_pSourceFile != INVALID_HANDLE_VALUE)
	{
		m_pcInputBuffer = new char[g_uInputBufferSize];
		uSize = uSourceBytesRemaining = GetFileSize(m_pSourceFile, 0);
	}
	else
	{
		uSize = m_uBufferSize;
	}
     

	int iFlushFlag = Z_NO_FLUSH;
	bool bAborted = false;

	do
	{
		if(m_pSourceBuffer)
		{
			//Can feed in the whole input buffer at once
			//Need to gratuitously forget the constness of the source buffer though
			m_xZlibStream.next_in = (Bytef*) const_cast<void*>(m_pSourceBuffer);
			m_xZlibStream.avail_in = m_uBufferSize;
			iFlushFlag = Z_FINISH;
		}
		else
		{
			//Get the next chunk of file
			u_int uBytesToRead = Asura_Maths::Min(uSourceBytesRemaining,g_uInputBufferSize);
			DWORD uBytesRead = 0;

			ReadFile(m_pSourceFile, m_pcInputBuffer, uBytesToRead, &uBytesRead, 0);

			if(uBytesRead != uBytesToRead)
			{
				Asura_Assert(0,"Error reading file to compress");
				return false;
			}

			uSourceBytesRemaining -= uBytesToRead;
			if(!uSourceBytesRemaining)
			{
				iFlushFlag = Z_FINISH;
			}

			m_xZlibStream.next_in = (Bytef*) m_pcInputBuffer;
			m_xZlibStream.avail_in = uBytesToRead;
		}
	
		
		/* run deflate() on input until output buffer not full*/
		do 
		{
		    m_xZlibStream.avail_out = g_uOutputBufferSize;
		    m_xZlibStream.next_out = (Bytef*) m_pcOutputBuffer;

		    iRet = deflate(&m_xZlibStream, iFlushFlag);
		    Asura_Assert(iRet != Z_STREAM_ERROR,"Error compressing file");

		    u_int uBytesToWrite = g_uOutputBufferSize - m_xZlibStream.avail_out;
			WRITE_FILE(m_pcOutputBuffer, uBytesToWrite);

			#ifdef ASURA_MAIN
			if(bAllowAbort)
			{
				Tools_Level_Export::UpdateStatus(Tools_Level_Export::EXPORT_STAGE_COMPRESSING_FILE, "Compressing File", -1);
				if(Tools_Level_Export::IsAbortFlagSet())
				{
					bAborted = true;
					break;
				}
			}
			#endif

		} while (m_xZlibStream.avail_out == 0);
		
		//I think all the input should be used by this point
		Asura_Assert(bAborted || m_xZlibStream.avail_in == 0,"Error compressing file");

	}while(iFlushFlag != Z_FINISH && !bAborted);



	if(!bAborted)
	{
    	Asura_Assert(iRet == Z_STREAM_END,"Error compressing file");
	
		//go back and fill in the header
		xHeader.m_uBigEndian = bBigEndian ? 1 : 0;
		xHeader.m_uCompressedSize = GetFileSize(m_pDestFile,0)-(8+sizeof(xHeader));
		xHeader.m_uUncompressedSize = uSize;

		if(bBigEndian)
		{
			//I think its just our header that will need endian swapping , and the compressed data should be okay.
			xHeader.m_uCompressedSize = Asura_Endian::SwapBytes(xHeader.m_uCompressedSize);
			xHeader.m_uUncompressedSize = Asura_Endian::SwapBytes(xHeader.m_uUncompressedSize);
		}

		SetFilePointer(m_pDestFile, 8, 0, FILE_BEGIN);
		WRITE_FILE(&xHeader, sizeof(Asura_Zlib_Header))
	}


    //clean up and return
    deflateEnd(&m_xZlibStream);
	m_bStreamOpened = false;
	CloseHandle(m_pDestFile);
	m_pDestFile = INVALID_HANDLE_VALUE;
	if(m_pSourceFile != INVALID_HANDLE_VALUE)
	{
		CloseHandle(m_pSourceFile);
		m_pSourceFile = INVALID_HANDLE_VALUE;

	}

	if(bAborted)
	{
		DeleteFile(m_xFilename.GetString());
	}

	return !bAborted;
}


bool Tools_Zlib::SaveCompressedFile(const void* pBuffer, const u_int uSize, const char* szFilename, bool bBigEndian /*=false*/)
{
	Tools_Zlib_Compressor xCompressor;
	return xCompressor.SaveCompressedFile(pBuffer,uSize,szFilename,bBigEndian);

}

bool Tools_Zlib::CompressFile(const char* szFilename, bool bBigEndian /*= false*/, bool bAllowAbort /*=false*/)
{
	if(Tools_FileSystem::IsReadOnly(szFilename))
	{
		return false;
	}

	//get a temporary name for the compressed file
	Asura_String xFullFilename = szFilename;
	Asura_String xPath,xName;

	xFullFilename.SplitFilename(&xPath, &xName, 0);
	if(xPath.IsBlank()) xPath = ".";

	char acTempName[MAX_PATH];
	if(!GetTempFileName(xPath.GetString(),xName.GetString(),0,acTempName))
	{
		Asura_Assert(0,"Error getting temporary filename");
		return false;
	}

	//copy and compress to the new file
	Tools_Zlib_Compressor xCompressor;
	if(!xCompressor.SaveCompressedFile(szFilename,acTempName,bBigEndian, bAllowAbort))
	{
		return false;
	}

	//delete the original and rename the temporary file
	if(DeleteFile(szFilename))
	{
		if(!MoveFile(acTempName,szFilename))
		{
			Asura_Assert(0,"Error renaming compressed file back to original name.");
			return false;
		}
	}
	else
	{
		Asura_Assert(0,"Failed to delete old copy of file being compressed");
		DeleteFile(acTempName);
		return false;
	}

	return true;
}

bool Tools_Zlib::CompressFile(const char* szSourceFile, const char* szDestFile, bool bBigEndian /*= false*/)
{
	if(Tools_FileSystem::IsReadOnly(szDestFile))
	{
		return false;
	}
	if(!stricmp(szSourceFile, szDestFile))
	{
		//source and destination the same
		//not really intended via this function , and the test won't catch all cases anyway
		Asura_Assert(0,"Not expecting the source and destination file to be the same for this function");
		return CompressFile(szSourceFile, bBigEndian);
	}


	//copy and compress to the new file
	Tools_Zlib_Compressor xCompressor;
	return xCompressor.SaveCompressedFile(szSourceFile,szDestFile,bBigEndian, false);
}




#endif //ASURA_USE_ZLIB
