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

#include <zenutil/chunkblock.h>

#include <zencore/compactbinarybuilder.h>
#include <zencore/logging.h>

#include <vector>

namespace zen {

using namespace std::literals;

ChunkBlockDescription
ParseChunkBlockDescription(const CbObjectView& BlockObject)
{
	ChunkBlockDescription Result;
	Result.BlockHash = BlockObject["rawHash"sv].AsHash();
	if (Result.BlockHash != IoHash::Zero)
	{
		CbArrayView ChunksArray = BlockObject["rawHashes"sv].AsArrayView();
		Result.ChunkHashes.reserve(ChunksArray.Num());
		for (CbFieldView ChunkView : ChunksArray)
		{
			Result.ChunkHashes.push_back(ChunkView.AsHash());
		}

		CbArrayView			  ChunkRawLengthsArray = BlockObject["chunkRawLengths"sv].AsArrayView();
		std::vector<uint32_t> ChunkLengths;
		Result.ChunkRawLengths.reserve(ChunkRawLengthsArray.Num());
		for (CbFieldView ChunkView : ChunkRawLengthsArray)
		{
			Result.ChunkRawLengths.push_back(ChunkView.AsUInt32());
		}
	}
	return Result;
}

std::vector<ChunkBlockDescription>
ParseChunkBlockDescriptionList(const CbObjectView& BlocksObject)
{
	if (!BlocksObject)
	{
		return {};
	}
	std::vector<ChunkBlockDescription> Result;
	CbArrayView						   Blocks = BlocksObject["blocks"].AsArrayView();
	Result.reserve(Blocks.Num());
	for (CbFieldView BlockView : Blocks)
	{
		CbObjectView BlockObject = BlockView.AsObjectView();
		Result.emplace_back(ParseChunkBlockDescription(BlockObject));
	}
	return Result;
}

CbObject
BuildChunkBlockDescription(const ChunkBlockDescription& Block, CbObjectView MetaData)
{
	ZEN_ASSERT(Block.ChunkRawLengths.size() == Block.ChunkHashes.size());

	CbObjectWriter Writer;
	Writer.AddHash("rawHash"sv, Block.BlockHash);
	Writer.BeginArray("rawHashes"sv);
	{
		for (const IoHash& ChunkHash : Block.ChunkHashes)
		{
			Writer.AddHash(ChunkHash);
		}
	}
	Writer.EndArray();
	Writer.BeginArray("chunkRawLengths");
	{
		for (uint32_t ChunkSize : Block.ChunkRawLengths)
		{
			Writer.AddInteger(ChunkSize);
		}
	}
	Writer.EndArray();

	Writer.AddObject("metadata", MetaData);

	return Writer.Save();
}

CompressedBuffer
GenerateChunkBlock(std::vector<std::pair<IoHash, FetchChunkFunc>>&& FetchChunks, ChunkBlockDescription& OutBlock)
{
	const size_t ChunkCount = FetchChunks.size();

	std::vector<SharedBuffer> ChunkSegments;
	ChunkSegments.resize(1);
	ChunkSegments.reserve(1 + ChunkCount);
	OutBlock.ChunkHashes.reserve(ChunkCount);
	OutBlock.ChunkRawLengths.reserve(ChunkCount);
	{
		IoBuffer		  TempBuffer(ChunkCount * 9);
		MutableMemoryView View			 = TempBuffer.GetMutableView();
		uint8_t*		  BufferStartPtr = reinterpret_cast<uint8_t*>(View.GetData());
		uint8_t*		  BufferEndPtr	 = BufferStartPtr;
		BufferEndPtr += WriteVarUInt(gsl::narrow<uint64_t>(ChunkCount), BufferEndPtr);
		for (const auto& It : FetchChunks)
		{
			std::pair<uint64_t, CompressedBuffer> Chunk		= It.second(It.first);
			uint64_t							  ChunkSize = 0;
			std::span<const SharedBuffer>		  Segments	= Chunk.second.GetCompressed().GetSegments();
			for (const SharedBuffer& Segment : Segments)
			{
				ChunkSize += Segment.GetSize();
				ChunkSegments.push_back(Segment);
			}
			BufferEndPtr += WriteVarUInt(ChunkSize, BufferEndPtr);
			OutBlock.ChunkHashes.push_back(It.first);
			OutBlock.ChunkRawLengths.push_back(gsl::narrow<uint32_t>(Chunk.first));
		}
		ZEN_ASSERT(BufferEndPtr <= View.GetDataEnd());
		ptrdiff_t TempBufferLength = std::distance(BufferStartPtr, BufferEndPtr);
		ChunkSegments[0]		   = SharedBuffer(IoBuffer(TempBuffer, 0, gsl::narrow<size_t>(TempBufferLength)));
	}
	CompressedBuffer CompressedBlock =
		CompressedBuffer::Compress(CompositeBuffer(std::move(ChunkSegments)), OodleCompressor::Mermaid, OodleCompressionLevel::None);
	OutBlock.BlockHash = CompressedBlock.DecodeRawHash();
	return CompressedBlock;
}

bool
IterateChunkBlock(const SharedBuffer& BlockPayload, std::function<void(CompressedBuffer&& Chunk, const IoHash& AttachmentHash)> Visitor)
{
	ZEN_ASSERT(BlockPayload);
	if (BlockPayload.GetSize() < 1)
	{
		return false;
	}

	MemoryView	   BlockView = BlockPayload.GetView();
	const uint8_t* ReadPtr	 = reinterpret_cast<const uint8_t*>(BlockView.GetData());
	uint32_t	   NumberSize;
	uint64_t	   ChunkCount = ReadVarUInt(ReadPtr, NumberSize);
	ReadPtr += NumberSize;
	std::vector<uint64_t> ChunkSizes;
	ChunkSizes.reserve(ChunkCount);
	while (ChunkCount--)
	{
		ChunkSizes.push_back(ReadVarUInt(ReadPtr, NumberSize));
		ReadPtr += NumberSize;
	}
	for (uint64_t ChunkSize : ChunkSizes)
	{
		IoBuffer		 Chunk(IoBuffer::Wrap, ReadPtr, ChunkSize);
		IoHash			 AttachmentRawHash;
		uint64_t		 AttachmentRawSize;
		CompressedBuffer CompressedChunk = CompressedBuffer::FromCompressed(SharedBuffer(Chunk), AttachmentRawHash, AttachmentRawSize);

		if (!CompressedChunk)
		{
			ZEN_ERROR("Invalid chunk in block");
			return false;
		}
		Visitor(std::move(CompressedChunk), AttachmentRawHash);
		ReadPtr += ChunkSize;
		ZEN_ASSERT(ReadPtr <= BlockView.GetDataEnd());
	}
	return true;
};

}  // namespace zen