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

#include <zenremotestore/builds/buildcontent.h>

#include <zencore/compactbinaryutil.h>
#include <zencore/fmtutils.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_set.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

using namespace std::literals;

void
WriteBuildContentToCompactBinary(CbObjectWriter&						PartManifestWriter,
								 const SourcePlatform					Platform,
								 std::span<const std::filesystem::path> Paths,
								 std::span<const IoHash>				RawHashes,
								 std::span<const uint64_t>				RawSizes,
								 std::span<const uint32_t>				Attributes,
								 std::span<const IoHash>				SequenceRawHashes,
								 std::span<const uint32_t>				ChunkCounts,
								 std::span<const IoHash>				LocalChunkHashes,
								 std::span<const uint64_t>				LocalChunkRawSizes,
								 const std::vector<uint32_t>&			AbsoluteChunkOrders,
								 const std::span<const uint32_t>		LooseLocalChunkIndexes,
								 const std::span<IoHash>				BlockHashes)
{
	ZEN_ASSERT(Platform != SourcePlatform::_Count);
	PartManifestWriter.AddString("platform"sv, ToString(Platform));

	uint64_t TotalSize = 0;
	for (const uint64_t Size : RawSizes)
	{
		TotalSize += Size;
	}
	PartManifestWriter.AddInteger("totalSize", TotalSize);

	PartManifestWriter.BeginObject("files"sv);
	{
		compactbinary_helpers::WriteArray(Paths, "paths"sv, PartManifestWriter);
		compactbinary_helpers::WriteArray(RawHashes, "rawhashes"sv, PartManifestWriter);
		compactbinary_helpers::WriteArray(RawSizes, "rawsizes"sv, PartManifestWriter);
		if (Platform == SourcePlatform::Windows)
		{
			compactbinary_helpers::WriteArray(Attributes, "attributes"sv, PartManifestWriter);
		}
		if (Platform == SourcePlatform::Linux || Platform == SourcePlatform::MacOS)
		{
			compactbinary_helpers::WriteArray(Attributes, "mode"sv, PartManifestWriter);
		}
	}
	PartManifestWriter.EndObject();	 // files

	PartManifestWriter.BeginObject("chunkedContent");
	{
		compactbinary_helpers::WriteArray(SequenceRawHashes, "sequenceRawHashes"sv, PartManifestWriter);
		compactbinary_helpers::WriteArray(ChunkCounts, "chunkcounts"sv, PartManifestWriter);
		compactbinary_helpers::WriteArray(AbsoluteChunkOrders, "chunkorders"sv, PartManifestWriter);
	}
	PartManifestWriter.EndObject();	 // chunkedContent

	size_t LooseChunkCount = LooseLocalChunkIndexes.size();
	if (LooseChunkCount > 0)
	{
		PartManifestWriter.BeginObject("chunkAttachments");
		{
			PartManifestWriter.BeginArray("rawHashes"sv);
			for (uint32_t ChunkIndex : LooseLocalChunkIndexes)
			{
				PartManifestWriter.AddBinaryAttachment(LocalChunkHashes[ChunkIndex]);
			}
			PartManifestWriter.EndArray();	// rawHashes

			PartManifestWriter.BeginArray("chunkRawSizes"sv);
			for (uint32_t ChunkIndex : LooseLocalChunkIndexes)
			{
				PartManifestWriter.AddInteger(LocalChunkRawSizes[ChunkIndex]);
			}
			PartManifestWriter.EndArray();	// chunkSizes
		}
		PartManifestWriter.EndObject();	 //
	}

	if (BlockHashes.size() > 0)
	{
		PartManifestWriter.BeginObject("blockAttachments");
		{
			compactbinary_helpers::WriteBinaryAttachmentArray(BlockHashes, "rawHashes"sv, PartManifestWriter);
		}
		PartManifestWriter.EndObject();	 // blocks
	}
}

void
ReadBuildContentFromCompactBinary(CbObjectView						  BuildPartManifest,
								  SourcePlatform&					  OutPlatform,
								  std::vector<std::filesystem::path>& OutPaths,
								  std::vector<IoHash>&				  OutRawHashes,
								  std::vector<uint64_t>&			  OutRawSizes,
								  std::vector<uint32_t>&			  OutAttributes,
								  std::vector<IoHash>&				  OutSequenceRawHashes,
								  std::vector<uint32_t>&			  OutChunkCounts,
								  std::vector<uint32_t>&			  OutAbsoluteChunkOrders,
								  std::vector<IoHash>&				  OutLooseChunkHashes,
								  std::vector<uint64_t>&			  OutLooseChunkRawSizes,
								  std::vector<IoHash>&				  OutBlockRawHashes)
{
	OutPlatform = FromString(BuildPartManifest["platform"sv].AsString(), SourcePlatform::_Count);

	CbObjectView FilesObject = BuildPartManifest["files"sv].AsObjectView();

	OutPaths	 = compactbinary_helpers::ReadArray<std::filesystem::path>("paths"sv, FilesObject);
	OutRawHashes = compactbinary_helpers::ReadArray<IoHash>("rawhashes"sv, FilesObject);
	OutRawSizes	 = compactbinary_helpers::ReadArray<uint64_t>("rawsizes"sv, FilesObject);

	uint64_t PathCount = OutPaths.size();
	if (OutRawHashes.size() != PathCount)
	{
		throw std::runtime_error(fmt::format("Number of raw hashes entries does not match number of paths"));
	}
	if (OutRawSizes.size() != PathCount)
	{
		throw std::runtime_error(fmt::format("Number of raw sizes entries does not match number of paths"));
	}

	std::vector<uint32_t> ModeArray = compactbinary_helpers::ReadArray<uint32_t>("mode"sv, FilesObject);
	if (ModeArray.size() != PathCount && ModeArray.size() != 0)
	{
		throw std::runtime_error(fmt::format("Number of attribute entries does not match number of paths"));
	}

	std::vector<uint32_t> AttributeArray = compactbinary_helpers::ReadArray<uint32_t>("attributes"sv, FilesObject);
	if (AttributeArray.size() != PathCount && AttributeArray.size() != 0)
	{
		throw std::runtime_error(fmt::format("Number of attribute entries does not match number of paths"));
	}

	if (ModeArray.size() > 0)
	{
		if (OutPlatform == SourcePlatform::_Count)
		{
			OutPlatform = SourcePlatform::Linux;  // Best guess - under dev format
		}
		OutAttributes = std::move(ModeArray);
	}
	else if (AttributeArray.size() > 0)
	{
		if (OutPlatform == SourcePlatform::_Count)
		{
			OutPlatform = SourcePlatform::Windows;
		}
		OutAttributes = std::move(AttributeArray);
	}
	else
	{
		if (OutPlatform == SourcePlatform::_Count)
		{
			OutPlatform = GetSourceCurrentPlatform();
		}
	}

	if (CbObjectView ChunkContentView = BuildPartManifest["chunkedContent"sv].AsObjectView(); ChunkContentView)
	{
		OutSequenceRawHashes = compactbinary_helpers::ReadArray<IoHash>("sequenceRawHashes"sv, ChunkContentView);
		OutChunkCounts		 = compactbinary_helpers::ReadArray<uint32_t>("chunkcounts"sv, ChunkContentView);
		if (OutChunkCounts.size() != OutSequenceRawHashes.size())
		{
			throw std::runtime_error(fmt::format("Number of chunk count entries does not match number of paths"));
		}
		OutAbsoluteChunkOrders = compactbinary_helpers::ReadArray<uint32_t>("chunkorders"sv, ChunkContentView);
	}
	else if (FilesObject["chunkcounts"sv])
	{
		// Legacy zen style

		std::vector<uint32_t> LegacyChunkCounts = compactbinary_helpers::ReadArray<uint32_t>("chunkcounts"sv, FilesObject);
		if (LegacyChunkCounts.size() != PathCount)
		{
			throw std::runtime_error(fmt::format("Number of chunk count entries does not match number of paths"));
		}
		std::vector<uint32_t> LegacyAbsoluteChunkOrders = compactbinary_helpers::ReadArray<uint32_t>("chunkorders"sv, FilesObject);

		CbArrayView	   ChunkOrdersArray = BuildPartManifest["chunkorders"sv].AsArrayView();
		const uint64_t ChunkOrdersCount = ChunkOrdersArray.Num();

		tsl::robin_set<IoHash, IoHash::Hasher> FoundRawHashes;
		FoundRawHashes.reserve(PathCount);

		OutChunkCounts.reserve(PathCount);
		OutAbsoluteChunkOrders.reserve(ChunkOrdersCount);

		uint32_t OrderIndexOffset = 0;
		for (uint32_t PathIndex = 0; PathIndex < OutPaths.size(); PathIndex++)
		{
			const IoHash& PathRawHash	   = OutRawHashes[PathIndex];
			uint32_t	  LegacyChunkCount = LegacyChunkCounts[PathIndex];

			if (FoundRawHashes.insert(PathRawHash).second)
			{
				OutSequenceRawHashes.push_back(PathRawHash);
				OutChunkCounts.push_back(LegacyChunkCount);
				std::span<uint32_t> AbsoluteChunkOrder =
					std::span<uint32_t>(LegacyAbsoluteChunkOrders).subspan(OrderIndexOffset, LegacyChunkCount);
				OutAbsoluteChunkOrders.insert(OutAbsoluteChunkOrders.end(), AbsoluteChunkOrder.begin(), AbsoluteChunkOrder.end());
			}
			OrderIndexOffset += LegacyChunkCounts[PathIndex];
		}
	}
	else
	{
		// Legacy C# style

		tsl::robin_set<IoHash, IoHash::Hasher> FoundRawHashes;
		FoundRawHashes.reserve(PathCount);
		uint32_t OrderIndexOffset = 0;
		for (uint32_t PathIndex = 0; PathIndex < OutPaths.size(); PathIndex++)
		{
			if (OutRawSizes[PathIndex] > 0)
			{
				const IoHash& PathRawHash = OutRawHashes[PathIndex];
				if (FoundRawHashes.insert(PathRawHash).second)
				{
					OutSequenceRawHashes.push_back(PathRawHash);
					OutChunkCounts.push_back(1);
					OutAbsoluteChunkOrders.push_back(OrderIndexOffset);
					OutLooseChunkHashes.push_back(PathRawHash);
					OutLooseChunkRawSizes.push_back(OutRawSizes[PathIndex]);
					OrderIndexOffset += 1;
				}
			}
		}
	}

	CbObjectView ChunkAttachmentsView = BuildPartManifest["chunkAttachments"sv].AsObjectView();
	{
		OutLooseChunkHashes	  = compactbinary_helpers::ReadBinaryAttachmentArray("rawHashes"sv, ChunkAttachmentsView);
		OutLooseChunkRawSizes = compactbinary_helpers::ReadArray<uint64_t>("chunkRawSizes"sv, ChunkAttachmentsView);
		if (OutLooseChunkHashes.size() != OutLooseChunkRawSizes.size())
		{
			throw std::runtime_error(fmt::format("Number of attachment chunk hashes does not match number of attachemnt chunk raw sizes"));
		}
	}

	CbObjectView BlocksView = BuildPartManifest["blockAttachments"sv].AsObjectView();
	{
		OutBlockRawHashes = compactbinary_helpers::ReadBinaryAttachmentArray("rawHashes"sv, BlocksView);
	}
}

}  // namespace zen