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

#include "zipfs.h"

#include <zencore/iobuffer.h>

#if ZEN_WITH_TESTS

ZEN_THIRD_PARTY_INCLUDES_START
#	include <doctest/doctest.h>
#	include <zlib.h>
ZEN_THIRD_PARTY_INCLUDES_END

#	include <cstring>
#	include <vector>

TEST_SUITE_BEGIN("server.zipfs");

namespace {

// Helpers to build a minimal zip file in memory
struct ZipBuilder
{
	std::vector<uint8_t> Data;

	struct Entry
	{
		std::string Name;
		uint32_t	LocalHeaderOffset;
		uint16_t	CompressionMethod;
		uint32_t	CompressedSize;
		uint32_t	UncompressedSize;
	};

	std::vector<Entry> Entries;

	void Append(const void* Src, size_t Size)
	{
		const uint8_t* Bytes = (const uint8_t*)Src;
		Data.insert(Data.end(), Bytes, Bytes + Size);
	}

	void AppendU16(uint16_t V) { Append(&V, 2); }
	void AppendU32(uint32_t V) { Append(&V, 4); }

	void AddFile(const std::string& Name, const void* Content, size_t ContentSize, bool Deflate)
	{
		std::vector<uint8_t> FileData;
		uint16_t			 Method = 0;

		if (Deflate)
		{
			// Compress with raw deflate (no zlib/gzip header)
			uLongf				 BoundSize = compressBound((uLong)ContentSize);
			std::vector<uint8_t> TempBuf(BoundSize);

			z_stream Stream	 = {};
			Stream.next_in	 = (Bytef*)Content;
			Stream.avail_in	 = (uInt)ContentSize;
			Stream.next_out	 = TempBuf.data();
			Stream.avail_out = (uInt)TempBuf.size();

			deflateInit2(&Stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -MAX_WBITS, 8, Z_DEFAULT_STRATEGY);
			deflate(&Stream, Z_FINISH);
			deflateEnd(&Stream);

			TempBuf.resize(Stream.total_out);
			FileData = std::move(TempBuf);
			Method	 = 8;
		}
		else
		{
			FileData.assign((const uint8_t*)Content, (const uint8_t*)Content + ContentSize);
		}

		Entry E;
		E.Name				= Name;
		E.LocalHeaderOffset = (uint32_t)Data.size();
		E.CompressionMethod = Method;
		E.CompressedSize	= (uint32_t)FileData.size();
		E.UncompressedSize	= (uint32_t)ContentSize;
		Entries.push_back(E);

		// Local file header
		AppendU32(0x04034b50);			   // signature
		AppendU16(20);					   // version needed
		AppendU16(0);					   // flags
		AppendU16(Method);				   // compression method
		AppendU16(0);					   // last mod time
		AppendU16(0);					   // last mod date
		AppendU32(0);					   // crc32 (not validated by ZipFs)
		AppendU32(E.CompressedSize);	   // compressed size
		AppendU32(E.UncompressedSize);	   // uncompressed size
		AppendU16((uint16_t)Name.size());  // file name length
		AppendU16(0);					   // extra field length
		Append(Name.data(), Name.size());  // file name
		Append(FileData.data(), FileData.size());
	}

	zen::IoBuffer Build()
	{
		uint32_t CdOffset = (uint32_t)Data.size();

		for (const Entry& E : Entries)
		{
			// Central directory record
			AppendU32(0x02014b50);				 // signature
			AppendU16(20);						 // version made by
			AppendU16(20);						 // version needed
			AppendU16(0);						 // flags
			AppendU16(E.CompressionMethod);		 // compression method
			AppendU16(0);						 // last mod time
			AppendU16(0);						 // last mod date
			AppendU32(0);						 // crc32
			AppendU32(E.CompressedSize);		 // compressed size
			AppendU32(E.UncompressedSize);		 // uncompressed size
			AppendU16((uint16_t)E.Name.size());	 // file name length
			AppendU16(0);						 // extra field length
			AppendU16(0);						 // comment length
			AppendU16(0);						 // disk index
			AppendU16(0);						 // internal file attr
			AppendU32(0);						 // external file attr
			AppendU32(E.LocalHeaderOffset);		 // offset
			Append(E.Name.data(), E.Name.size());
		}

		uint32_t CdSize = (uint32_t)Data.size() - CdOffset;

		// End of central directory record
		AppendU32(0x06054b50);				  // signature
		AppendU16(0);						  // this disk
		AppendU16(0);						  // cd start disk
		AppendU16((uint16_t)Entries.size());  // cd records this disk
		AppendU16((uint16_t)Entries.size());  // cd records total
		AppendU32(CdSize);					  // cd size
		AppendU32(CdOffset);				  // cd offset
		AppendU16(0);						  // comment length

		zen::IoBuffer Buffer(Data.size());
		std::memcpy(Buffer.GetMutableView().GetData(), Data.data(), Data.size());
		return Buffer;
	}
};

}  // namespace

TEST_CASE("zipfs.stored")
{
	const char* Content = "Hello, World!";

	ZipBuilder Zip;
	Zip.AddFile("test.txt", Content, std::strlen(Content), false);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("test.txt");
	REQUIRE(Result);
	CHECK(Result.GetView().GetSize() == std::strlen(Content));
	CHECK(std::memcmp(Result.GetView().GetData(), Content, std::strlen(Content)) == 0);
}

TEST_CASE("zipfs.deflate")
{
	const char* Content = "This is some content that will be deflate compressed in the zip file.";

	ZipBuilder Zip;
	Zip.AddFile("compressed.txt", Content, std::strlen(Content), true);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("compressed.txt");
	REQUIRE(Result);
	CHECK(Result.GetView().GetSize() == std::strlen(Content));
	CHECK(std::memcmp(Result.GetView().GetData(), Content, std::strlen(Content)) == 0);
}

TEST_CASE("zipfs.mixed")
{
	const char* StoredContent  = "stored content";
	const char* DeflateContent = "deflate content that is compressed";

	ZipBuilder Zip;
	Zip.AddFile("stored.txt", StoredContent, std::strlen(StoredContent), false);
	Zip.AddFile("deflated.txt", DeflateContent, std::strlen(DeflateContent), true);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Stored = Fs.GetFile("stored.txt");
	REQUIRE(Stored);
	CHECK(Stored.GetView().GetSize() == std::strlen(StoredContent));
	CHECK(std::memcmp(Stored.GetView().GetData(), StoredContent, std::strlen(StoredContent)) == 0);

	zen::IoBuffer Deflated = Fs.GetFile("deflated.txt");
	REQUIRE(Deflated);
	CHECK(Deflated.GetView().GetSize() == std::strlen(DeflateContent));
	CHECK(std::memcmp(Deflated.GetView().GetData(), DeflateContent, std::strlen(DeflateContent)) == 0);
}

TEST_CASE("zipfs.not_found")
{
	const char* Content = "data";

	ZipBuilder Zip;
	Zip.AddFile("exists.txt", Content, std::strlen(Content), false);

	zen::ZipFs Fs(Zip.Build());

	zen::IoBuffer Result = Fs.GetFile("missing.txt");
	CHECK(!Result);
}

TEST_SUITE_END();

#endif	// ZEN_WITH_TESTS