// Copyright (C) 2009 Gaz Davidson
// for copying permissions see readme.txt

#include "CResLoader.h"

#include "IFileSystem.h"

namespace irr
{
namespace io
{

//! Constructor
CArchiveLoaderRes::CArchiveLoaderRes(io::IFileSystem* fs)
: FileSystem(fs)
{

}

//! returns true if the file maybe is able to be loaded by this class
//! based on the file extension (e.g. ".exe")
bool CArchiveLoaderRes::isALoadableFileFormat(const core::string<c16>& filename) const
{
	return core::hasFileExtension(filename, "exe", "dll", "cpl") ||
	       core::hasFileExtension(filename, "scr", "ocx");
}

//! Creates an archive from the filename
IFileArchive* CArchiveLoaderRes::createArchive(const core::string<c16>& filename, bool ignoreCase, bool ignorePaths) const
{
	IFileArchive *archive = 0;
	io::IReadFile* file = FileSystem->createAndOpenFile(filename);

	if (file)
	{
		archive = createArchive(file, ignoreCase, ignorePaths);
		file->drop();
	}

	return archive;
}

//! Check if the file might be loaded by this class
bool CArchiveLoaderRes::isALoadableFileFormat(io::IReadFile* file) const
{
	file->seek(0);

	// check for "MZ"

	c8 b[2];
	if (file->read(b, 2) != 2 || b[0] != 'M' || b[1] != 'Z')
		return false;
	
	// get PE header pos
	s32 seekPos = 0;
	if (!file->seek(60) || file->read(&seekPos, 4) != 4)
		return false;

	if (seekPos > file->getSize() || seekPos < file->getPos())
		return false;

	// must be "PE\0\0"
	u32 PESig=0;
	if (!file->seek(seekPos) || file->read(&PESig, 4) != 4 || PESig != PEBinarySignature)
		return false;

	// this is a valid PE binary!
	return true;
}

//! creates/loads an archive from the file.
io::IFileArchive* CArchiveLoaderRes::createArchive(io::IReadFile* file, bool ignoreCase, bool ignorePaths) const
{
	IFileArchive *archive = 0;
	if (file)
	{
		file->seek(0);
		archive = new CResReader(FileSystem, file, ignoreCase, ignorePaths);
	}
	return archive;
}


//! constructor
CResReader::CResReader(IFileSystem* fs, IReadFile* file, bool ignoreCase, bool ignorePaths)
 : FileSystem(fs), File(file), IgnoreCase(ignoreCase), IgnorePaths(ignorePaths)
{
	#ifdef _DEBUG
	setDebugName("CResReader");
	#endif

	if (File)
	{
		File->grab();

		Base = File->getFileName();
		Base.replace('\\', '/');

		// fill the file list
		populateFileList();
	}
}	

//! destructor
CResReader::~CResReader()
{
	if (File)
		File->drop();
}

//! opens a file by file name
IReadFile* CResReader::createAndOpenFile(const core::string<c16>& filename)
{
	s32 index = findFile(filename);

	if (index != -1)
		return createAndOpenFile(index);

	return 0;
}

//! opens a file by index
IReadFile* CResReader::createAndOpenFile(u32 index)
{
	if (index >= FileList.size())
		return 0;

	// TODO: Need to reconstruct files from the resource data. 
	// Most data types are missing file headers, without these the file data needs extra processing.

	switch(FileList[index].ResourceType)
	{
	case ERSRCT_BITMAP:
		{
			if (!File->seek(FileList[index].Offset))
				return 0;

			// create memory file
			CResMemoryReadWriteFile* ret = new CResMemoryReadWriteFile(FileList[index].simpleFileName.c_str());

			// allocate memory
			ret->getData().set_used(FileList[index].Length + sizeof(SBMPHeader::SBMPFileHeader));

			// set up header pointer and copy in data
			SBMPHeader *bitmapHeader= (SBMPHeader*)(&(ret->getData()[0]));
			if (File->read(&(ret->getData()[sizeof(SBMPHeader::SBMPFileHeader)]), FileList[index].Length) != FileList[index].Length)
			{
				ret->drop();
				return 0;
			}

			// set header
			bitmapHeader->FileHeader.Id = 0x4d42; // BM
			bitmapHeader->FileHeader.FileSize = FileList[index].Length + sizeof(SBMPHeader::SBMPFileHeader);
			bitmapHeader->FileHeader.Reserved = 0;
			bitmapHeader->FileHeader.BitmapDataOffset = bitmapHeader->FileHeader.FileSize - bitmapHeader->BitmapDataSize;

			return ret;
		}
	case ERSRCT_CURSOR:
	case ERSRCT_ICON:
		{
			// cursors and icons are the same...

			if (!File->seek(FileList[index].Offset))
				return 0;

			// create memory file
			CResMemoryReadWriteFile* ret = new CResMemoryReadWriteFile(FileList[index].simpleFileName.c_str());

			// allocate memory
			ret->getData().set_used(FileList[index].Length + sizeof(SIcoFileHeader) + sizeof(SIconDetails));

			// assign pointers
			SIcoFileHeader *icoHead  = (SIcoFileHeader*)(&(ret->getData()[0]));
			SIconDetails *icoDetails = (SIconDetails*)(&(ret->getData()[sizeof(SIcoFileHeader)]));
			SBMPHeader *bitmapHeader = (SBMPHeader*)(&(ret->getData()[sizeof(SIcoFileHeader) + sizeof(SIconDetails) - sizeof(SBMPHeader::SBMPFileHeader)]));

			// read data
			if (File->read(&(ret->getData()[sizeof(SIcoFileHeader) + sizeof(SIconDetails)]), FileList[index].Length) != FileList[index].Length)
			{
				ret->drop();
				return 0;
			}

			// set up header
			icoHead->itemCount      = 1;
			icoHead->reserved       = 0;
			icoHead->resourceType   = FileList[index].ResourceType == ERSRCT_ICON ? 1 : 2;
			
			icoDetails->width       = bitmapHeader->Width;
			icoDetails->height      = bitmapHeader->Height;
			icoDetails->colourCount = bitmapHeader->Colors;
			icoDetails->reserved    = 0;
			icoDetails->planes      = bitmapHeader->Planes;
			icoDetails->bpp         = bitmapHeader->BPP;
			icoDetails->size        = FileList[index].Length;
			icoDetails->offset      = sizeof(SIcoFileHeader) + sizeof(SIconDetails);

			return ret;
		}
	case ERSRCT_CURSOR_GROUP:
	case ERSRCT_ICON_GROUP:
		{
			if (!File->seek(FileList[index].Offset))
				return 0;

			SIcoFileHeader icoHead;
			if (File->read(&icoHead, sizeof(SIcoFileHeader)) != sizeof(SIcoFileHeader))
				return 0;

			// create memory file
			CResMemoryReadWriteFile* ret = new CResMemoryReadWriteFile(FileList[index].simpleFileName.c_str());

			// reserve space for longer icon directory items
			ret->getData().set_used(FileList[index].Length + icoHead.itemCount*2);

			// copy file header
			memcpy(&ret->getData()[0], &icoHead, sizeof(SIcoFileHeader));

			// calculate real size
			SIconDetails icoDetails;
			SIconDetails *icoDetailsDest = (SIconDetails*)(&(ret->getData()[sizeof(SIcoFileHeader)]));
			u32 totalSize = FileList[index].Length + icoHead.itemCount*2;
			for (u32 i=0; i < icoHead.itemCount; ++i)
			{
				icoDetails.offset = 0;
				// read smaller version of icon details header
				if (File->read(&icoDetails, sizeof(SIconDetails)-2) != sizeof(SIconDetails)-2)
				{
					ret->drop();
					return 0;
				}
				icoDetailsDest[i] = icoDetails;
				totalSize += icoDetails.size;
			}

			// reserve real size
			ret->getData().set_used(totalSize);
			
			// reset details pointer
			icoDetailsDest = (SIconDetails*)(&(ret->getData()[sizeof(SIcoFileHeader)]));

			u32 currentPos = FileList[index].Length + icoHead.itemCount*2;
			for (u32 i=0; i < icoHead.itemCount; ++i)
			{
				// find file
				core::string<c16> fileName = FileList[index].path;
				if (fileName.size())
					fileName.append('/');
				fileName.append( core::string<c16>(icoDetailsDest[i].resourceID) );
				fileName.append( FileList[index].ResourceType == ERSRCT_ICON_GROUP ? ".ico" : ".cur" );
				s32 j = findFile(fileName);

				if (j == -1 || !File->seek(FileList[j].Offset))
				{
					ret->drop();
					return 0;
				}

				// append icon data
				if (File->read(&(ret->getData()[currentPos]), icoDetailsDest[i].size) != icoDetailsDest[i].size)
				{
					ret->drop();
					return 0;
				}

				// set correct offset of data
				icoDetailsDest[i].offset = currentPos;

				// move current position
				currentPos += icoDetailsDest[i].size;
			}
			return ret;
		}
	default:
		return FileSystem->createLimitReadFile(FileList[index].simpleFileName, 
			File, FileList[index].Offset, FileList[index].Length);
	}
}

//! returns count of files in archive
u32 CResReader::getFileCount() const
{
	return FileList.size();
}

//! returns data of file
const IFileArchiveEntry* CResReader::getFileInfo(u32 index)
{
	return &FileList[index];
}

//! returns fileindex
s32 CResReader::findFile(const core::string<c16>& filename)
{
	SResArchiveEntry entry;
	entry.simpleFileName = filename;

	if (IgnoreCase)
		entry.simpleFileName.make_lower();

	if (IgnorePaths)
		core::deletePathFromFilename(entry.simpleFileName);

	s32 res = FileList.binary_search(entry);

	return res;
}

//! return the id of the file Archive
const core::string<c16>& CResReader::getArchiveName()
{
	return Base;
}

u32 CResReader::populateFileList()
{
	FileList.clear();
	File->seek(0);

	// seek to the start of the PE info
	u32 PEStartOffset = 0;
	if (!File->seek(60) || File->read(&PEStartOffset, 4) != 4)
		return 0;

	// signature must be "PE\0\0"
	u32 PESig=0;
	if (!File->seek(PEStartOffset) || File->read(&PESig, 4) != 4 || PESig != PEBinarySignature)
		return 0;

	// read the file header 
	SCOFFFileHeader fileHead;
	if (File->read(&fileHead, sizeof(SCOFFFileHeader)) != sizeof(SCOFFFileHeader))
		return 0;

	// skip the optional header
	if (!File->seek(fileHead.SizeOfOptionalHeader, true))
		return 0;

	// useless empty file
	if (!fileHead.NumberOfSections)
		return 0;

	SSectionHeader sectionHeader;

	// now we parse the sections, looking for the resource section
	for (u16 i=0; i < fileHead.NumberOfSections; ++i)
	{
		if (File->read(&sectionHeader, sizeof(SSectionHeader)) != sizeof(SSectionHeader))
			return 0;

		sectionHeader.VirtualSize = 0; // null terminate string in place, we don't need the data

		if (strcmp(sectionHeader.Name, ".rsrc") == 0)
			ResourceTreeAddress = sectionHeader.PointerToRawData;
	}

	ResourceDataAddress = 0;

	// we need to recursively add the folders starting at the root.
	SResArchiveEntry entry;
	if (!addFolder(0, entry, 0))
	{
		FileList.clear();
	}

	if (!FileList.size())
		return 0;

	// We know the file sizes, the start of the file data block and their 
	// positions relative to each other, so find the general offset from the
	// minimum position
	u32 minPos = 0xFFFFFFFF;
	for (u32 i=0; i < FileList.size(); ++i)
		if (FileList[i].Offset < minPos)
			minPos = FileList[i].Offset;

	u32 currentPos = ResourceDataAddress;
	s32 diff = ResourceDataAddress - minPos;

	// finish up
	for (u32 i=0; i < FileList.size(); ++i)
	{
		// set the correct pointer to the file data
		FileList[i].Offset += diff;

		// append the correct file extension
		addFileExtension(FileList[i]);

		if (IgnoreCase)
		{
			FileList[i].simpleFileName.make_lower();
			FileList[i].path.make_lower();
		}
	}

	return FileList.size();
}

bool CResReader::addFolder(c8 level, SResArchiveEntry entry, u32 offset)
{
	SResourceDirectoryTable table;

	// read the table
	if (!File->seek(ResourceTreeAddress + offset) || 
		File->read(&table, sizeof(SResourceDirectoryTable)) != sizeof(SResourceDirectoryTable))
	{
		return false;
	}

	SResourceDirectoryEntry dirEntry;

	for (u16 i=0; i < table.NameEntryCount + table.IDEntryCount; ++i)
	{
		SResArchiveEntry newEntry = entry;

		if (!File->seek(ResourceTreeAddress + offset + sizeof(SResourceDirectoryTable) + sizeof(SResourceDirectoryEntry)*i))
			return false;

		if (File->read(&dirEntry, sizeof(SResourceDirectoryEntry)) != sizeof(SResourceDirectoryEntry))
			return false;

		// get the name
		if (dirEntry.NameRVAorIntegerID & 0x80000000)
		{
			// this is a string
			core::string<c16> name;
			if (!getName(name, dirEntry.NameRVAorIntegerID ^ 0x80000000))
				return false;

			switch (level)
			{
			case 0:
				// type
				newEntry.ResourceType = ERSRCT_NAMED_RESOURCE_TYPE;
				newEntry.ResourceTypeName = name;
				break;
			case 1:
				// name
				newEntry.simpleFileName = name;
			case 2:
				// language (path)
				if (!IgnorePaths)
					newEntry.path = name;
				break;
			default:
				// eh? these aren't allowed!
				return false;
			}
		}
		else
		{
			// it's an ID
			switch (level)
			{
			case 0:
				// type
				newEntry.ResourceType = E_RESOURCE_TYPE(dirEntry.NameRVAorIntegerID);
				newEntry.ResourceTypeName = "";
				break;
			case 1:
				// name
				newEntry.simpleFileName = core::string<c16>(dirEntry.NameRVAorIntegerID);
			case 2:
				// language (path)
				if (!IgnorePaths)
					newEntry.path = core::string<c16>(dirEntry.NameRVAorIntegerID);
				break;
			default:
				// eh? these aren't allowed!
				return false;
			}
		}

		if (dirEntry.DataEntryRVAorSubdirectoryRVA & 0x80000000)
		{
			// it's a folder
			if (!addFolder(level+1, newEntry, dirEntry.DataEntryRVAorSubdirectoryRVA ^ 0x80000000))
				return false;
		}
		else
		{
			// it's a file- get file data 
			SResourceData fileData;
			if (!File->seek(ResourceTreeAddress + dirEntry.DataEntryRVAorSubdirectoryRVA))
				return false;

			if (File->read(&fileData, sizeof(SResourceData)) != sizeof(SResourceData))
				return false;

			// since we have no idea where this is...
			if (u32(File->getPos()) > ResourceDataAddress)
				ResourceDataAddress = File->getPos();

			newEntry.Offset = fileData.DataRVA; // incomplete
			newEntry.Length = fileData.Size;
			newEntry.path   = "";

			FileList.push_back(newEntry);
		}
	}
	// end of folder
	return true;
}

bool CResReader::getName(core::string<c16>& name, u32 offset)
{
	if (!File->seek(ResourceTreeAddress + offset))
		return false;

	u16 length;
	File->read(&length, 2);
	wchar_t data;
	for (u16 i=0; i < length; ++i)
	{
		if (File->read(&data, 2) != 2)
			return false;
		name.append(data);
	}

	// since we have no idea where this is...
	if (u32(File->getPos()) > ResourceDataAddress)
		ResourceDataAddress = File->getPos();

	return true;
}

// Add the file extension depending on the resource type
void CResReader::addFileExtension(SResArchiveEntry &entry)
{
	// first things first
	switch (entry.ResourceType)
	{
	case ERSRCT_CURSOR:
		entry.simpleFileName.append(".cur");
		return;

	case ERSRCT_CURSOR_GROUP:
		entry.simpleFileName.append(".group.cur");
		return;

	case ERSRCT_BITMAP:
		entry.simpleFileName.append(".bmp");
		return;
	
	case ERSRCT_ICON:
		entry.simpleFileName.append(".ico");
		return;

	case ERSRCT_ICON_GROUP:
		entry.simpleFileName.append(".group.ico");
		return;

	case ERSRCT_MANIFEST:
		entry.simpleFileName.append(".xml");
		return;

	case ERSRCT_NAMED_RESOURCE_TYPE:
		entry.simpleFileName.append('.');
		entry.simpleFileName.append(entry.ResourceTypeName);
		return;

	// no idea what these ones actually contain...
	case ERSRCT_FONT_DIR:
	case ERSRCT_FONT:
	case ERSRCT_MENU:
	case ERSRCT_DIALOG:
	case ERSRCT_STRING:
	case ERSRCT_ACCELERATORS:
	default:
		break;
	}
	
	// ...so just append the data type ID as the extension
	entry.simpleFileName.append('.');
	entry.simpleFileName.append( core::string<c16>(u32(entry.ResourceType)));
}



///////////////////////////////////// Memory file

CResMemoryReadWriteFile::CResMemoryReadWriteFile(const c16* filename) 
:  Data(), FileName(filename), Pos(0)
{
}


s32 CResMemoryReadWriteFile::write(const void* buffer, u32 sizeToWrite)
{
	// no point in writing 0 bytes
	if (sizeToWrite < 1)
		return 0;

	// expand size
	if (Pos + sizeToWrite > Data.size())
		Data.set_used(Pos+sizeToWrite);

	// copy data
	memcpy( (void*) &Data[Pos], buffer, (size_t) sizeToWrite);

	Pos += sizeToWrite;

	return sizeToWrite;
	
}

bool CResMemoryReadWriteFile::seek(long finalPos, bool relativeMovement)
{
	if (relativeMovement)
	{
		if (finalPos + Pos < 0)
			return 0;
		else
			Pos += finalPos;
	}
	else
	{
		Pos = finalPos;
	}

	if (Pos > (s32)Data.size())
		Data.set_used(Pos+1);

	return true;
}

const core::string<c16>& CResMemoryReadWriteFile::getFileName() const
{
	return FileName;
}

long CResMemoryReadWriteFile::getPos() const
{
	return Pos;
}

core::array<c8>& CResMemoryReadWriteFile::getData()
{
	return Data;
}


long CResMemoryReadWriteFile::getSize() const
{
	return Data.size();
}


s32 CResMemoryReadWriteFile::read(void* buffer, u32 sizeToRead)
{
	// cant read past the end
	if (Pos + sizeToRead >= Data.size())
		sizeToRead = Data.size() - Pos;

	// cant read 0 bytes
	if (!sizeToRead)
		return 0;

	// copy data
	memcpy( buffer, (void*) &Data[Pos], (size_t) sizeToRead);
	
	Pos += sizeToRead;

	return sizeToRead;
}





} // namespace io
} // namespace irr