// Copyright Epic Games, Inc. All Rights Reserved.

#include "zipfs.h"

namespace zen {

//////////////////////////////////////////////////////////////////////////
namespace {

#if ZEN_COMPILER_MSC
#	pragma warning(push)
#	pragma warning(disable : 4200)
#endif

	using ZipInt16 = uint16_t;

	struct ZipInt32
	{
				 operator uint32_t() const { return *(uint32_t*)Parts; }
		uint16_t Parts[2];
	};

	struct EocdRecord
	{
		enum : uint32_t
		{
			Magic = 0x0605'4b50,
		};
		ZipInt32 Signature;
		ZipInt16 ThisDiskIndex;
		ZipInt16 CdStartDiskIndex;
		ZipInt16 CdRecordThisDiskCount;
		ZipInt16 CdRecordCount;
		ZipInt32 CdSize;
		ZipInt32 CdOffset;
		ZipInt16 CommentSize;
		char	 Comment[];
	};

	struct CentralDirectoryRecord
	{
		enum : uint32_t
		{
			Magic = 0x0201'4b50,
		};

		ZipInt32 Signature;
		ZipInt16 VersionMadeBy;
		ZipInt16 VersionRequired;
		ZipInt16 Flags;
		ZipInt16 CompressionMethod;
		ZipInt16 LastModTime;
		ZipInt16 LastModDate;
		ZipInt32 Crc32;
		ZipInt32 CompressedSize;
		ZipInt32 OriginalSize;
		ZipInt16 FileNameLength;
		ZipInt16 ExtraFieldLength;
		ZipInt16 CommentLength;
		ZipInt16 DiskIndex;
		ZipInt16 InternalFileAttr;
		ZipInt32 ExternalFileAttr;
		ZipInt32 Offset;
		char	 FileName[];
	};

	struct LocalFileHeader
	{
		enum : uint32_t
		{
			Magic = 0x0403'4b50,
		};

		ZipInt32 Signature;
		ZipInt16 VersionRequired;
		ZipInt16 Flags;
		ZipInt16 CompressionMethod;
		ZipInt16 LastModTime;
		ZipInt16 LastModDate;
		ZipInt32 Crc32;
		ZipInt32 CompressedSize;
		ZipInt32 OriginalSize;
		ZipInt16 FileNameLength;
		ZipInt16 ExtraFieldLength;
		char	 FileName[];
	};

#if ZEN_COMPILER_MSC
#	pragma warning(pop)
#endif

}  // namespace

//////////////////////////////////////////////////////////////////////////
ZipFs::ZipFs(IoBuffer&& Buffer)
{
	MemoryView View = Buffer.GetView();

	uint8_t* Cursor = (uint8_t*)(View.GetData()) + View.GetSize();
	if (View.GetSize() < sizeof(EocdRecord))
	{
		return;
	}

	const auto* EocdCursor = (EocdRecord*)(Cursor - sizeof(EocdRecord));

	// It is more correct to search backwards for EocdRecord::Magic as the
	// comment can be of a variable length. But here we're not going to support
	// zip files with comments.
	if (EocdCursor->Signature != EocdRecord::Magic)
	{
		return;
	}

	// Zip64 isn't supported either
	if (EocdCursor->ThisDiskIndex == 0xffff)
	{
		return;
	}

	Cursor = (uint8_t*)EocdCursor - uint32_t(EocdCursor->CdOffset) - uint32_t(EocdCursor->CdSize);

	const auto* CdCursor = (CentralDirectoryRecord*)(Cursor + EocdCursor->CdOffset);
	for (int i = 0, n = EocdCursor->CdRecordCount; i < n; ++i)
	{
		const CentralDirectoryRecord& Cd = *CdCursor;

		bool Acceptable = true;
		Acceptable &= (Cd.OriginalSize > 0);		// has some content
		Acceptable &= (Cd.CompressionMethod == 0);	// is stored uncomrpessed
		if (Acceptable)
		{
			const uint8_t* Lfh = Cursor + Cd.Offset;
			if (uintptr_t(Lfh - Cursor) < View.GetSize())
			{
				std::string_view FileName(Cd.FileName, Cd.FileNameLength);
				m_Files.insert(std::make_pair(FileName, FileItem{Lfh, size_t(0)}));
			}
		}

		uint32_t ExtraBytes = Cd.FileNameLength + Cd.ExtraFieldLength + Cd.CommentLength;
		CdCursor			= (CentralDirectoryRecord*)(Cd.FileName + ExtraBytes);
	}

	m_Buffer = std::move(Buffer);
}

//////////////////////////////////////////////////////////////////////////
IoBuffer
ZipFs::GetFile(const std::string_view& FileName) const
{
	{
		RwLock::SharedLockScope _(m_FilesLock);

		FileMap::const_iterator Iter = m_Files.find(FileName);
		if (Iter == m_Files.end())
		{
			return {};
		}

		const FileItem& Item = Iter->second;
		if (Item.GetSize() > 0)
		{
			return IoBuffer(IoBuffer::Wrap, Item.GetData(), Item.GetSize());
		}
	}

	RwLock::ExclusiveLockScope _(m_FilesLock);

	FileItem& Item = m_Files.find(FileName)->second;
	if (Item.GetSize() > 0)
	{
		return IoBuffer(IoBuffer::Wrap, Item.GetData(), Item.GetSize());
	}

	const auto* Lfh = (LocalFileHeader*)(Item.GetData());
	Item			= MemoryView(Lfh->FileName + Lfh->FileNameLength + Lfh->ExtraFieldLength, Lfh->OriginalSize);
	return IoBuffer(IoBuffer::Wrap, Item.GetData(), Item.GetSize());
}

}  // namespace zen