/*
 * HLLib
 * Copyright (C) 2006 Ryan Gregg

 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later
 * version.
 */

#include "HLLib.h"
#include "VBSPFile.h"
#include "Streams.h"
#include "Checksum.h"

using namespace HLLib;

#define HL_VBSP_LUMP_ENTITIES		0
#define HL_VBSP_LUMP_PAKFILE		40

#define HL_VBSP_ZIP_LOCAL_FILE_HEADER_SIGNATURE					0x04034b50
#define HL_VBSP_ZIP_FILE_HEADER_SIGNATURE						0x02014b50
#define HL_VBSP_ZIP_END_OF_CENTRAL_DIRECTORY_RECORD_SIGNATURE	0x06054b50

#define HL_VBSP_ZIP_CHECKSUM_LENGTH								0x00008000

const char *CVBSPFile::lpAttributeNames[] = { "Version", "Map Revision" };
const char *CVBSPFile::lpItemAttributeNames[] = { "Version", "Four CC", "Zip Disk", "Zip Comment", "Create Version", "Extract Version", "Flags", "Compression Method", "CRC", "Disk", "Comment" };

CVBSPFile::CVBSPFile() : CPackage(), pHeaderView(0), pFileHeaderView(0), pEndOfCentralDirectoryRecordView(0), pHeader(0), pEndOfCentralDirectoryRecord(0)
{

}

CVBSPFile::~CVBSPFile()
{
	this->Close();
}

HLPackageType CVBSPFile::GetType() const
{
	return HL_PACKAGE_VBSP;
}

const hlChar *CVBSPFile::GetExtension() const
{
	return "bsp";
}

const hlChar *CVBSPFile::GetDescription() const
{
	return "Half-Life 2 Level";
}

hlBool CVBSPFile::MapDataStructures()
{
	if(sizeof(VBSPHeader) > this->pMapping->GetMappingSize())
	{
		LastError.SetErrorMessage("Invalid file: the file map is too small for it's header.");
		return hlFalse;
	}

	if(!this->pMapping->Map(this->pHeaderView, 0, sizeof(VBSPHeader)))
	{
		return hlFalse;
	}
	this->pHeader = static_cast<const VBSPHeader *>(this->pHeaderView->GetView());

	if(memcmp(this->pHeader->lpSignature, "VBSP", 4) != 0)
	{
		LastError.SetErrorMessage("Invalid file: the file's signature does not match.");
		return hlFalse;
	}

	// Versions:
	//  19-20:			Source
	//  0x00040014:		Dark Messiah - Looks like the 32 bit version has been split into two 16 bit fields.
	if(this->pHeader->iVersion != 19 && this->pHeader->iVersion != 20 && this->pHeader->iVersion != 0x00040014)
	{
		LastError.SetErrorMessageFormated("Invalid VBSP version (v%i): you have a version of a VBSP file that HLLib does not know how to read. Check for product updates.", this->pHeader->iVersion);
		return hlFalse;
	}

	if(sizeof(ZIPEndOfCentralDirectoryRecord) <= this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiLength)
	{
		hlUInt uiTest, uiOffset = this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiOffset;
		while(uiOffset < this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiOffset + this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiLength - sizeof(uiTest))
		{
			Mapping::CView *pTestView = 0;

			if(!this->pMapping->Map(pTestView, uiOffset, sizeof(hlUInt)))
			{
				return hlFalse;
			}

			uiTest = *(hlUInt *)pTestView->GetView();

			this->pMapping->Unmap(pTestView);

			switch(uiTest)
			{
				case HL_VBSP_ZIP_END_OF_CENTRAL_DIRECTORY_RECORD_SIGNATURE:
				{
					if(!this->pMapping->Map(pTestView, uiOffset, sizeof(ZIPEndOfCentralDirectoryRecord)))
					{
						return hlFalse;
					}

					const ZIPEndOfCentralDirectoryRecord EndOfCentralDirRecord = *static_cast<const ZIPEndOfCentralDirectoryRecord *>(pTestView->GetView());

					this->pMapping->Unmap(pTestView);

					if(!this->pMapping->Map(this->pEndOfCentralDirectoryRecordView, uiOffset, sizeof(ZIPEndOfCentralDirectoryRecord) + EndOfCentralDirRecord.uiCommentLength))
					{
						return hlFalse;
					}

					this->pEndOfCentralDirectoryRecord = static_cast<const ZIPEndOfCentralDirectoryRecord *>(this->pEndOfCentralDirectoryRecordView->GetView());

					if(!this->pMapping->Map(this->pFileHeaderView, this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiOffset + this->pEndOfCentralDirectoryRecord->uiStartOfCentralDirOffset, this->pEndOfCentralDirectoryRecord->uiCentralDirectorySize))
					{
						return hlFalse;
					}

					return hlTrue;
				}
				case HL_VBSP_ZIP_FILE_HEADER_SIGNATURE:
				{
					if(!this->pMapping->Map(pTestView, uiOffset, sizeof(ZIPFileHeader)))
					{
						return hlFalse;
					}

					const ZIPFileHeader FileHeader = *static_cast<const ZIPFileHeader *>(pTestView->GetView());

					this->pMapping->Unmap(pTestView);

					uiOffset += sizeof(ZIPFileHeader) + FileHeader.uiFileNameLength + FileHeader.uiExtraFieldLength + FileHeader.uiFileCommentLength;
					break;
				}
				case HL_VBSP_ZIP_LOCAL_FILE_HEADER_SIGNATURE:
				{
					if(!this->pMapping->Map(pTestView, uiOffset, sizeof(ZIPLocalFileHeader)))
					{
						return hlFalse;
					}

					const ZIPLocalFileHeader LocalFileHeader = *static_cast<const ZIPLocalFileHeader *>(pTestView->GetView());

					this->pMapping->Unmap(pTestView);

					uiOffset += sizeof(ZIPLocalFileHeader) + LocalFileHeader.uiFileNameLength + LocalFileHeader.uiExtraFieldLength + LocalFileHeader.uiCompressedSize;
					break;
				}
				default:
				{
					LastError.SetErrorMessageFormated("Invalid file: unknown ZIP section signature %#.8x.", uiTest);
					return hlFalse;
				}
			}
		}

		LastError.SetErrorMessage("Invalid file: unexpected end of file while scanning for end of ZIP central directory record.");
		return hlFalse;
	}

	return hlTrue;
}

hlVoid CVBSPFile::UnmapDataStructures()
{
	this->pMapping->Unmap(this->pFileHeaderView);

	this->pEndOfCentralDirectoryRecord = 0;
	this->pMapping->Unmap(this->pEndOfCentralDirectoryRecordView);

	this->pHeader = 0;
	this->pMapping->Unmap(this->pHeaderView);
}

CDirectoryFolder *CVBSPFile::CreateRoot()
{
	CDirectoryFolder *pRoot = new CDirectoryFolder(this);

	hlChar lpFileName[256];

	if(this->pHeader->lpLumps[HL_VBSP_LUMP_ENTITIES].uiLength != 0)
	{
		this->GetFileName(lpFileName, sizeof(lpFileName) - 4);
		if(*lpFileName == '\0')
		{
			pRoot->AddFile("entities.ent", HL_VBSP_LUMP_ENTITIES);
		}
		else
		{
			strcat(lpFileName, ".ent");
			pRoot->AddFile(lpFileName, HL_VBSP_LUMP_ENTITIES);
		}
	}

	if(this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiLength != 0)
	{
		this->GetFileName(lpFileName, sizeof(lpFileName) - 4);
		if(*lpFileName == '\0')
		{
			pRoot->AddFile("pakfile.zip", HL_VBSP_LUMP_PAKFILE);
		}
		else
		{
			strcat(lpFileName, ".zip");
			pRoot->AddFile(lpFileName, HL_VBSP_LUMP_PAKFILE);
		}
	}

	hlUInt uiTest, uiOffset = 0;
	while(uiOffset < this->pEndOfCentralDirectoryRecord->uiCentralDirectorySize - sizeof(uiTest))
	{
		uiTest = *(hlUInt *)((hlByte *)this->pFileHeaderView->GetView() + uiOffset);

		switch(uiTest)
		{
			case HL_VBSP_ZIP_FILE_HEADER_SIGNATURE:
			{
				ZIPFileHeader *pFileHeader = static_cast<ZIPFileHeader *>((hlVoid *)((hlByte *)this->pFileHeaderView->GetView() + uiOffset));

				hlChar *lpFileName = new hlChar[pFileHeader->uiFileNameLength + 1];
				memcpy(lpFileName, (hlByte *)pFileHeader + sizeof(ZIPFileHeader), pFileHeader->uiFileNameLength);
				lpFileName[pFileHeader->uiFileNameLength] = '\0';

				// Check if we have just a file, or if the file has directories we need to create.
				if(strchr(lpFileName, '/') == 0 && strchr(lpFileName, '\\') == 0)
				{
					pRoot->AddFile(lpFileName, HL_ID_INVALID, pFileHeader);
				}
				else
				{
					// Tokenize the file path and create the directories.
					CDirectoryFolder *pInsertFolder = pRoot;

					hlChar lpTemp[256] = "";
					hlChar *lpToken = strtok(lpFileName, "/\\");
					while(lpToken != 0)
					{
						strcpy(lpTemp, lpToken);

						lpToken = strtok(0, "/\\");

						if(lpToken != 0)
						{
							// Check if the directory exists.
							CDirectoryItem *pItem = pInsertFolder->GetItem(lpTemp);
							if(pItem == 0 || pItem->GetType() == HL_ITEM_FILE)
							{
								// It doesn't, create it.
								pInsertFolder = pInsertFolder->AddFolder(lpTemp);
							}
							else
							{
								// It does, use it.
								pInsertFolder = static_cast<CDirectoryFolder *>(pItem);
							}
						}
					}

					// The file name is the last token, add it.
					pInsertFolder->AddFile(lpTemp, HL_ID_INVALID, pFileHeader);
				}

				delete []lpFileName;

				uiOffset += sizeof(ZIPFileHeader) + pFileHeader->uiFileNameLength + pFileHeader->uiExtraFieldLength + pFileHeader->uiFileCommentLength;
				break;
			}
			default:
			{
				uiOffset = this->pEndOfCentralDirectoryRecord->uiCentralDirectorySize;
				break;
			}
		}
	}

	return pRoot;
}

hlUInt CVBSPFile::GetAttributeCountInternal() const
{
	return HL_VBSP_PACKAGE_COUNT;
}

const hlChar *CVBSPFile::GetAttributeNameInternal(HLPackageAttribute eAttribute) const
{
	if(eAttribute < HL_VBSP_PACKAGE_COUNT)
	{
		return this->lpAttributeNames[eAttribute];
	}

	return 0;
}

hlBool CVBSPFile::GetAttributeInternal(HLPackageAttribute eAttribute, HLAttribute &Attribute) const
{
	switch(eAttribute)
	{
	case HL_VBSP_PACKAGE_VERSION:
		hlAttributeSetInteger(&Attribute, this->lpAttributeNames[eAttribute], this->pHeader->iVersion);
		return hlTrue;
	case HL_VBSP_PACKAGE_MAP_REVISION:
		hlAttributeSetInteger(&Attribute, this->lpAttributeNames[eAttribute], this->pHeader->iMapRevision);
		return hlTrue;
	default:
		return hlFalse;
	}
}

hlUInt CVBSPFile::GetItemAttributeCountInternal() const
{
	return HL_VBSP_ITEM_COUNT;
}

const hlChar *CVBSPFile::GetItemAttributeNameInternal(HLPackageAttribute eAttribute) const
{
	if(eAttribute < HL_VBSP_ITEM_COUNT)
	{
		return this->lpItemAttributeNames[eAttribute];
	}

	return 0;
}


hlBool CVBSPFile::GetItemAttributeInternal(const CDirectoryItem *pItem, HLPackageAttribute eAttribute, HLAttribute &Attribute) const
{
	switch(pItem->GetType())
	{
		case HL_ITEM_FILE:
		{
			const CDirectoryFile *pFile = static_cast<const CDirectoryFile *>(pItem);
			if(pFile->GetData())
			{
				const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());
				switch(eAttribute)
				{
					case HL_VBSP_ZIP_ITEM_CREATE_VERSION:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiVersionMadeBy, hlFalse);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_EXTRACT_VERSION:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiVersionNeededToExtract, hlFalse);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_FLAGS:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiFlags, hlTrue);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_COMPRESSION_METHOD:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiCompressionMethod, hlTrue);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_CRC:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiCRC32, hlTrue);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_DISK:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], pDirectoryItem->uiDiskNumberStart, hlFalse);
						return hlTrue;
					}
					case HL_VBSP_ZIP_ITEM_COMMENT:
					{
						hlChar *lpComment = new hlChar[pDirectoryItem->uiFileCommentLength + 1];
						memcpy(lpComment, (hlByte *)pDirectoryItem + sizeof(ZIPFileHeader) + pDirectoryItem->uiFileNameLength + pDirectoryItem->uiExtraFieldLength, pDirectoryItem->uiFileCommentLength);
						lpComment[pDirectoryItem->uiFileCommentLength] = '\0';

						hlAttributeSetString(&Attribute, this->lpItemAttributeNames[eAttribute], lpComment);

						delete []lpComment;
						return hlTrue;
					}
				}
			}
			else
			{
				switch(eAttribute)
				{
					case HL_VBSP_ITEM_VERSION:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], this->pHeader->lpLumps[pFile->GetID()].uiVersion, hlFalse);
						return hlTrue;
					}
					case HL_VBSP_ITEM_FOUR_CC:
					{
						hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], *(hlUInt *)this->pHeader->lpLumps[pFile->GetID()].lpFourCC, hlTrue);
						return hlTrue;
					}
				}

				if(pItem->GetID() == HL_VBSP_LUMP_PAKFILE)
				{
					switch(eAttribute)
					{
						case HL_VBSP_ZIP_PACKAGE_DISK:
						{
							hlAttributeSetUnsignedInteger(&Attribute, this->lpItemAttributeNames[eAttribute], this->pEndOfCentralDirectoryRecord->uiNumberOfThisDisk, hlFalse);
							return hlTrue;
						}
						case HL_VBSP_ZIP_PACKAGE_COMMENT:
						{
							hlChar *lpComment = new hlChar[this->pEndOfCentralDirectoryRecord->uiCommentLength + 1];
							memcpy(lpComment, (hlByte *)this->pEndOfCentralDirectoryRecord + sizeof(ZIPEndOfCentralDirectoryRecord), this->pEndOfCentralDirectoryRecord->uiCommentLength);
							lpComment[this->pEndOfCentralDirectoryRecord->uiCommentLength] = '\0';

							hlAttributeSetString(&Attribute, this->lpItemAttributeNames[eAttribute], lpComment);

							delete []lpComment;
							return hlTrue;
						}
					}
				}
			}
			break;
		}
	}

	return hlFalse;
}

hlBool CVBSPFile::GetFileExtractableInternal(const CDirectoryFile *pFile, hlBool &bExtractable) const
{
	if(pFile->GetData())
	{
		const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());

		bExtractable = pDirectoryItem->uiCompressionMethod == 0 && pDirectoryItem->uiDiskNumberStart == this->pEndOfCentralDirectoryRecord->uiNumberOfThisDisk;
	}
	else
	{
		bExtractable = hlTrue;
	}

	return hlTrue;
}

hlBool CVBSPFile::GetFileValidationInternal(const CDirectoryFile *pFile, HLValidation &eValidation) const
{
	if(pFile->GetData())
	{
		const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());

		if(pDirectoryItem->uiCompressionMethod != 0 || pDirectoryItem->uiDiskNumberStart != this->pEndOfCentralDirectoryRecord->uiNumberOfThisDisk)
		{
			eValidation = HL_VALIDATES_ASSUMED_OK;
			return hlTrue;
		}

		hlULong uiChecksum = 0;
		Streams::IStream *pStream = 0;
		if(const_cast<CVBSPFile *>(this)->CreateStreamInternal(pFile, pStream))
		{
			if(pStream->Open(HL_MODE_READ))
			{
				hlUInt uiTotalBytes = 0, uiFileBytes = pStream->GetStreamSize(), uiBufferSize;
				hlByte lpBuffer[HL_VBSP_ZIP_CHECKSUM_LENGTH];

				hlBool bCancel = hlFalse;
				if(pValidateFileProgressProc != 0)
				{
					pValidateFileProgressProc(const_cast<CDirectoryFile *>(pFile), uiTotalBytes, uiFileBytes, &bCancel);
				}

				while((uiBufferSize = pStream->Read(lpBuffer, sizeof(lpBuffer))) != 0)
				{
					if(bCancel)
					{
						eValidation = HL_VALIDATES_CANCELED;
						break;
					}

					uiChecksum = CRC32(lpBuffer, uiBufferSize, uiChecksum);

					if(pValidateFileProgressProc != 0)
					{
						pValidateFileProgressProc(const_cast<CDirectoryFile *>(pFile), uiTotalBytes, uiFileBytes, &bCancel);
					}
				}

				pStream->Close();
			}

			const_cast<CVBSPFile *>(this)->ReleaseStreamInternal(*pStream);
			delete pStream;
		}

		eValidation = (hlULong)pDirectoryItem->uiCRC32 == uiChecksum ? HL_VALIDATES_OK : HL_VALIDATES_CORRUPT;
	}
	else
	{
		eValidation = HL_VALIDATES_ASSUMED_OK;
	}

	return hlTrue;
}

hlBool CVBSPFile::GetFileSizeInternal(const CDirectoryFile *pFile, hlUInt &uiSize) const
{
	if(pFile->GetData())
	{
		const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());

		uiSize = pDirectoryItem->uiUncompressedSize;
	}
	else
	{
		uiSize = this->pHeader->lpLumps[pFile->GetID()].uiLength;
	}

	return hlTrue;
}

hlBool CVBSPFile::GetFileSizeOnDiskInternal(const CDirectoryFile *pFile, hlUInt &uiSize) const
{
	if(pFile->GetData())
	{
		const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());

		uiSize = pDirectoryItem->uiCompressedSize;
	}
	else
	{
		uiSize = this->pHeader->lpLumps[pFile->GetID()].uiLength;
	}

	return hlTrue;
}

hlBool CVBSPFile::CreateStreamInternal(const CDirectoryFile *pFile, Streams::IStream *&pStream)
{
	if(pFile->GetData())
	{
		const ZIPFileHeader *pDirectoryItem = static_cast<const ZIPFileHeader *>(pFile->GetData());

		if(pDirectoryItem->uiCompressionMethod != 0)
		{
			LastError.SetErrorMessageFormated("Compression format %#.2x not supported.", pDirectoryItem->uiCompressionMethod);
			return hlFalse;
		}

		if(pDirectoryItem->uiDiskNumberStart != this->pEndOfCentralDirectoryRecord->uiNumberOfThisDisk)
		{
			LastError.SetErrorMessageFormated("File resides on disk %u.", pDirectoryItem->uiDiskNumberStart);
			return hlFalse;
		}

		Mapping::CView *pDirectoryEnrtyView = 0;

		if(!this->pMapping->Map(pDirectoryEnrtyView, this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiOffset + pDirectoryItem->uiRelativeOffsetOfLocalHeader, sizeof(ZIPLocalFileHeader)))
		{
			return hlFalse;
		}

		const ZIPLocalFileHeader DirectoryEntry = *static_cast<const ZIPLocalFileHeader *>(pDirectoryEnrtyView->GetView());

		this->pMapping->Unmap(pDirectoryEnrtyView);

		if(DirectoryEntry.uiSignature != HL_VBSP_ZIP_LOCAL_FILE_HEADER_SIGNATURE)
		{
			LastError.SetErrorMessageFormated("Invalid file data offset.", pDirectoryItem->uiDiskNumberStart);
			return hlFalse;
		}

		pStream = new Streams::CMappingStream(*this->pMapping, this->pHeader->lpLumps[HL_VBSP_LUMP_PAKFILE].uiOffset + pDirectoryItem->uiRelativeOffsetOfLocalHeader + sizeof(ZIPLocalFileHeader) + DirectoryEntry.uiFileNameLength + DirectoryEntry.uiExtraFieldLength, DirectoryEntry.uiUncompressedSize);
	}
	else
	{
		pStream = new Streams::CMappingStream(*this->pMapping, this->pHeader->lpLumps[pFile->GetID()].uiOffset, this->pHeader->lpLumps[pFile->GetID()].uiLength);
	}

	return hlTrue;
}

hlVoid CVBSPFile::GetFileName(hlChar *lpBuffer, hlUInt uiBufferSize)
{
	if(lpBuffer == 0 || uiBufferSize == 0)
	{
		return;
	}

	const hlChar *lpMappingName = this->pMapping->GetFileName();

	if(lpMappingName == 0 && *lpMappingName == '\0')
	{
		*lpBuffer = '\0';
		return;
	}

	const hlChar *lpForward = strrchr(lpMappingName, '\\');
	const hlChar *lpBackward = strrchr(lpMappingName, '/');
	const hlChar *lpStart = lpForward > lpBackward ? lpForward : lpBackward;

	if(lpStart == 0)
	{
		lpStart = lpMappingName;
	}
	else
	{
		lpStart++;
	}

	const hlChar *lpEnd = strrchr(lpStart, '.');

	if(lpEnd == 0)
	{
		lpEnd = lpStart + strlen(lpStart);
	}

	if((hlUInt)(lpEnd - lpStart) + 1 < uiBufferSize)
	{
		uiBufferSize = (hlUInt)(lpEnd - lpStart) + 1;
	}

	strncpy(lpBuffer, lpStart, uiBufferSize);
	lpBuffer[uiBufferSize - 1] = '\0';
}