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

#pragma once

#include <zencore/compositebuffer.h>

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

namespace zen {

#pragma pack(push)
#pragma pack(1)

struct ReferenceMetaDataHeader
{
	static constexpr uint32_t Version1		 = 1;
	static constexpr uint32_t CurrentVersion = Version1;

	ReferenceMetaDataHeader(uint32_t InMagic, uint32_t InEntryCount, uint64_t InAttachmentCount)
	: Magic(InMagic)
	, EntryCount(InEntryCount)
	, AttachmentCount(InAttachmentCount)
	, Checksum(ComputeChecksum())
	{
	}

	uint32_t Magic	 = 0xffffffffu;
	uint32_t Version = CurrentVersion;

	uint32_t EntryCount		 = 0;
	uint64_t AttachmentCount = 0;
	uint64_t Padding		 = 0;

	uint32_t Checksum = 0;

	bool IsValid(uint32_t ExpectedMagic) const;

	template<typename KeyType, typename AttachmentType>
	static uint64_t ExpectedSize(uint32_t EntryCount, uint64_t AttachmentCount)
	{
		return sizeof(ReferenceMetaDataHeader) + sizeof(KeyType) * EntryCount + sizeof(uint32_t) * EntryCount +
			   sizeof(AttachmentType) * AttachmentCount;
	}

	ReferenceMetaDataHeader() = delete;

private:
	uint32_t ComputeChecksum() const;
};

static_assert(sizeof(ReferenceMetaDataHeader) == 32);

#pragma pack(pop)

template<typename KeyType, typename AttachmentType>
CompositeBuffer
BuildReferenceMetaData(uint32_t						   HeaderMagic,
					   std::span<const KeyType>		   Keys,
					   std::span<const uint32_t>	   AttachmentCounts,
					   std::span<const AttachmentType> Attachments)
{
	uint32_t KeyCount = gsl::narrow<uint32_t>(Keys.size());

	ReferenceMetaDataHeader Header(HeaderMagic, KeyCount, Attachments.size());
	return CompositeBuffer(std::vector<SharedBuffer>{
		SharedBuffer(IoBuffer(IoBuffer::Clone, &Header, sizeof(ReferenceMetaDataHeader))),
		SharedBuffer(IoBuffer(IoBuffer::Wrap, Keys.data(), sizeof(KeyType) * Keys.size())),
		SharedBuffer(IoBuffer(IoBuffer::Wrap, AttachmentCounts.data(), sizeof(uint32_t) * AttachmentCounts.size())),
		SharedBuffer(IoBuffer(IoBuffer::Wrap, Attachments.data(), sizeof(AttachmentType) * Attachments.size()))});
}

template<typename KeyType, typename AttachmentType>
bool
GetAttachmentsFromMetaData(const IoBuffer&													  MetaData,
						   const uint32_t													  ExpectedHeaderMagic,
						   std::function<void(std::span<const KeyType>		  Keys,
											  std::span<const uint32_t>		  AttachmentCounts,
											  std::span<const AttachmentType> Attachments)>&& Parser)
{
	MemoryView PayloadView = MetaData.GetView();
	if (PayloadView.GetSize() >= sizeof(ReferenceMetaDataHeader))
	{
		const ReferenceMetaDataHeader* Header = (const ReferenceMetaDataHeader*)PayloadView.GetData();
		if (Header->IsValid(ExpectedHeaderMagic))
		{
			if (MetaData.GetSize() ==
				ReferenceMetaDataHeader::ExpectedSize<KeyType, AttachmentType>(Header->EntryCount, Header->AttachmentCount))
			{
				PayloadView.MidInline(sizeof(ReferenceMetaDataHeader));
				std::span<const KeyType> PayloadKeys = {(const KeyType*)PayloadView.GetData(), Header->EntryCount};
				PayloadView.MidInline(sizeof(KeyType) * Header->EntryCount);
				std::span<const uint32_t> PayloadAttachmentCounts = {(const uint32_t*)PayloadView.GetData(), Header->EntryCount};
				PayloadView.MidInline(sizeof(uint32_t) * Header->EntryCount);
				std::span<const AttachmentType> PayloadAttachments = {(const AttachmentType*)PayloadView.GetData(),
																	  Header->AttachmentCount};
				PayloadView.MidInline(sizeof(AttachmentType) * Header->AttachmentCount);
				ZEN_ASSERT(PayloadView.GetSize() == 0);

				Parser(PayloadKeys, PayloadAttachmentCounts, PayloadAttachments);

				return true;
			}
		}
	}
	return false;
}
}  // namespace zen