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

#include "httpshared.h"

#include <zencore/compactbinarypackage.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/stream.h>

#include <span>
#include <vector>

namespace zen {

std::vector<IoBuffer>
FormatPackageMessage(const CbPackage& Data)
{
	const std::span<const CbAttachment>& Attachments = Data.GetAttachments();

	std::vector<IoBuffer> ResponseBuffers;
	ResponseBuffers.reserve(3 + Attachments.size());  // TODO: may want to use an additional fudge factor here to avoid growing since each
													  // attachment is likely to consist of several buffers

	uint64_t TotalAttachmentsSize = 0;

	// Fixed size header

	CbPackageHeader Hdr{.HeaderMagic = kCbPkgMagic, .AttachmentCount = gsl::narrow<uint32_t>(Attachments.size())};

	ResponseBuffers.push_back(IoBufferBuilder::MakeCloneFromMemory(&Hdr, sizeof Hdr));

	// Attachment metadata array

	IoBuffer AttachmentMetadataBuffer = IoBuffer{sizeof(CbAttachmentEntry) * (Attachments.size() + /* root */ 1)};

	CbAttachmentEntry* AttachmentInfo = reinterpret_cast<CbAttachmentEntry*>(AttachmentMetadataBuffer.MutableData());

	ResponseBuffers.push_back(AttachmentMetadataBuffer);  // Attachment metadata

	// Root object

	IoBuffer RootIoBuffer = Data.GetObject().GetBuffer().AsIoBuffer();
	ResponseBuffers.push_back(RootIoBuffer);  // Root object

	*AttachmentInfo++ = {.AttachmentSize = RootIoBuffer.Size(), .AttachmentHash = Data.GetObjectHash()};

	// Attachment payloads

	for (const CbAttachment& Attachment : Attachments)
	{
		CompressedBuffer AttachmentBuffer = Attachment.AsCompressedBinary();
		CompositeBuffer	 Compressed		  = AttachmentBuffer.GetCompressed();

		*AttachmentInfo++ = {.AttachmentSize = AttachmentBuffer.GetCompressedSize(),
							 .AttachmentHash = IoHash::FromBLAKE3(AttachmentBuffer.GetRawHash())};

		for (const SharedBuffer& Segment : Compressed.GetSegments())
		{
			ResponseBuffers.push_back(Segment.AsIoBuffer());
			TotalAttachmentsSize += Segment.GetSize();
		}
	}

	return std::move(ResponseBuffers);
}

CbPackage
ParsePackageMessage(IoBuffer Payload)
{
	MemoryInStream InStream(Payload);
	BinaryReader   Reader(InStream);

	if (!Payload)
	{
		return {};
	}

	CbPackage Package;

	CbPackageHeader Hdr;
	Reader.Read(&Hdr, sizeof Hdr);

	if (Hdr.HeaderMagic != kCbPkgMagic)
	{
		// report error
		return {};
	}

	uint32_t ChunkCount = Hdr.AttachmentCount + 1;

	std::unique_ptr<CbAttachmentEntry[]> AttachmentEntries{new CbAttachmentEntry[ChunkCount]};

	Reader.Read(AttachmentEntries.get(), sizeof(CbAttachmentEntry) * ChunkCount);

	for (uint32_t i = 0; i < ChunkCount; ++i)
	{
		const uint64_t AttachmentSize = AttachmentEntries[i].AttachmentSize;
		IoBuffer	   AttachmentBuffer{AttachmentSize};
		Reader.Read(AttachmentBuffer.MutableData(), AttachmentSize);
		CompressedBuffer CompBuf(CompressedBuffer::FromCompressed(SharedBuffer(AttachmentBuffer)));

		if (i == 0)
		{
			Package.SetObject(LoadCompactBinaryObject(CompBuf));
		}
		else
		{
			CbAttachment Attachment(CompBuf);
			Package.AddAttachment(Attachment);
		}
	}

	return Package;
}

}  // namespace zen