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

#include "zipfs.h"

#include <zencore/logging.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <zlib.h>
ZEN_THIRD_PARTY_INCLUDES_END

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 || Cd.CompressionMethod == 8);	 // stored or deflate
		if (Acceptable)
		{
			const uint8_t* Lfh = Cursor + Cd.Offset;
			if (uintptr_t(Lfh - Cursor) < View.GetSize())
			{
				std::string_view FileName(Cd.FileName, Cd.FileNameLength);
				FileItem		 Item;
				Item.View			   = MemoryView{Lfh, size_t(0)};
				Item.CompressionMethod = Cd.CompressionMethod;
				Item.CompressedSize	   = Cd.CompressedSize;
				Item.UncompressedSize  = Cd.OriginalSize;
				m_Files.insert(std::make_pair(FileName, std::move(Item)));
			}
		}

		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.View.GetSize() > 0)
		{
			return IoBuffer(IoBuffer::Wrap, Item.View.GetData(), Item.View.GetSize());
		}
	}

	RwLock::ExclusiveLockScope _(m_FilesLock);

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

	const auto*	   Lfh		= (LocalFileHeader*)(Item.View.GetData());
	const uint8_t* FileData = (const uint8_t*)(Lfh->FileName + Lfh->FileNameLength + Lfh->ExtraFieldLength);

	if (Item.CompressionMethod == 0)
	{
		// Stored — point directly into the buffer
		Item.View = MemoryView(FileData, Item.UncompressedSize);
	}
	else
	{
		// Deflate — decompress using zlib
		Item.DecompressedData = IoBuffer(Item.UncompressedSize);

		z_stream Stream	 = {};
		Stream.next_in	 = const_cast<Bytef*>(FileData);
		Stream.avail_in	 = Item.CompressedSize;
		Stream.next_out	 = (Bytef*)Item.DecompressedData.GetMutableView().GetData();
		Stream.avail_out = Item.UncompressedSize;

		// Use raw inflate (-MAX_WBITS) since zip stores raw deflate streams
		if (inflateInit2(&Stream, -MAX_WBITS) != Z_OK)
		{
			ZEN_WARN("failed to initialize inflate for '{}'", FileName);
			return {};
		}

		int Result = inflate(&Stream, Z_FINISH);
		inflateEnd(&Stream);

		if (Result != Z_STREAM_END)
		{
			ZEN_WARN("failed to decompress '{}' (zlib error {})", FileName, Result);
			return {};
		}

		Item.View = Item.DecompressedData.GetView();
	}

	return IoBuffer(IoBuffer::Wrap, Item.View.GetData(), Item.View.GetSize());
}

}  // namespace zen