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

#include <algorithm>

#include <zenutil/chunkrequests.h>

#include <zencore/blake3.h>
#include <zencore/iobuffer.h>
#include <zencore/sharedbuffer.h>
#include <zencore/stream.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {
namespace {
	struct RequestHeader
	{
		enum
		{
			kMagic = 0xAAAA'77AC
		};
		uint32_t Magic;
		uint32_t ChunkCount;
		uint32_t Reserved1;
		uint32_t Reserved2;
	};

	struct ResponseHeader
	{
		uint32_t Magic = 0xbada'b00f;
		uint32_t ChunkCount;
		uint32_t Reserved1 = 0;
		uint32_t Reserved2 = 0;
	};

	struct ResponseChunkEntry
	{
		uint32_t CorrelationId;
		uint32_t Flags = 0;
		uint64_t ChunkSize;
	};
}  // namespace

IoBuffer
BuildChunkBatchRequest(const std::vector<RequestChunkEntry>& Entries)
{
	RequestHeader RequestHdr;
	RequestHdr.Magic			  = (uint32_t)RequestHeader::kMagic;
	RequestHdr.ChunkCount		  = gsl::narrow<uint32_t>(Entries.size());
	UniqueBuffer	  Buffer	  = UniqueBuffer::Alloc(sizeof(RequestHeader) + sizeof(RequestChunkEntry) * RequestHdr.ChunkCount);
	MutableMemoryView WriteBuffer = Buffer.GetMutableView();
	WriteBuffer					  = WriteBuffer.CopyFrom(MemoryView(&RequestHdr, sizeof(RequestHeader)));
	WriteBuffer.CopyFrom(MemoryView(Entries.data(), sizeof(RequestChunkEntry) * RequestHdr.ChunkCount));
	return Buffer.MoveToShared().AsIoBuffer();
}

std::optional<std::vector<RequestChunkEntry>>
ParseChunkBatchRequest(const IoBuffer& Payload)
{
	if (Payload.Size() <= sizeof(RequestHeader))
	{
		return {};
	}

	BinaryReader Reader(Payload);

	RequestHeader RequestHdr;
	Reader.Read(&RequestHdr, sizeof RequestHdr);

	if (RequestHdr.Magic != RequestHeader::kMagic)
	{
		return {};
	}

	std::vector<RequestChunkEntry> RequestedChunks;
	RequestedChunks.resize(RequestHdr.ChunkCount);
	Reader.Read(RequestedChunks.data(), sizeof(RequestChunkEntry) * RequestHdr.ChunkCount);
	return RequestedChunks;
}

std::vector<IoBuffer>
BuildChunkBatchResponse(const std::vector<RequestChunkEntry>& Requests, std::span<IoBuffer> Chunks)
{
	ZEN_ASSERT(Requests.size() == Chunks.size());
	size_t ChunkCount = Requests.size();

	std::vector<IoBuffer> OutBlobs;
	OutBlobs.reserve(1 + ChunkCount);
	OutBlobs.emplace_back(sizeof(ResponseHeader) + ChunkCount * sizeof(ResponseChunkEntry));

	uint8_t*	   ResponsePtr = reinterpret_cast<uint8_t*>(OutBlobs[0].MutableData());
	ResponseHeader ResponseHdr;
	ResponseHdr.ChunkCount = gsl::narrow<uint32_t>(Requests.size());
	memcpy(ResponsePtr, &ResponseHdr, sizeof(ResponseHdr));
	ResponsePtr += sizeof(ResponseHdr);
	for (uint32_t ChunkIndex = 0; ChunkIndex < ChunkCount; ++ChunkIndex)
	{
		const IoBuffer&	   FoundChunk(Chunks[ChunkIndex]);
		ResponseChunkEntry ResponseChunk;
		ResponseChunk.CorrelationId = Requests[ChunkIndex].CorrelationId;
		if (FoundChunk)
		{
			ResponseChunk.ChunkSize = FoundChunk.Size();
		}
		else
		{
			ResponseChunk.ChunkSize = uint64_t(-1);
		}
		memcpy(ResponsePtr, &ResponseChunk, sizeof(ResponseChunk));
		ResponsePtr += sizeof(ResponseChunk);
	}
	OutBlobs.insert(OutBlobs.end(), Chunks.begin(), Chunks.end());
	auto It = std::remove_if(OutBlobs.begin() + 1, OutBlobs.end(), [](const IoBuffer& B) { return B.GetSize() == 0; });
	OutBlobs.erase(It, OutBlobs.end());
	return OutBlobs;
}

std::vector<IoBuffer>
ParseChunkBatchResponse(const IoBuffer& Buffer)
{
	MemoryView			  View	 = Buffer.GetView();
	const ResponseHeader* Header = (const ResponseHeader*)View.GetData();
	if (Header->Magic != 0xbada'b00f)
	{
		return {};
	}
	View.MidInline(sizeof(ResponseHeader));
	const ResponseChunkEntry* Entries = (const ResponseChunkEntry*)View.GetData();
	View.MidInline(sizeof(ResponseChunkEntry) * Header->ChunkCount);
	std::vector<IoBuffer> Result(Header->ChunkCount);
	for (uint32_t Index = 0; Index < Header->ChunkCount; Index++)
	{
		const ResponseChunkEntry& Entry = Entries[Index];
		if (Result.size() < Entry.CorrelationId + 1)
		{
			Result.resize(Entry.CorrelationId + 1);
		}
		if (Entry.ChunkSize != uint64_t(-1))
		{
			Result[Entry.CorrelationId] = IoBuffer(IoBuffer::Wrap, View.GetData(), Entry.ChunkSize);
			View.MidInline(Entry.ChunkSize);
		}
	}
	return Result;
}

}  // namespace zen