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

#include <zencore/compositebuffer.h>

#include <zencore/sharedbuffer.h>
#include <zencore/testing.h>

namespace zen {

const CompositeBuffer CompositeBuffer::Null;

void
CompositeBuffer::Reset()
{
	m_Segments.clear();
}

uint64_t
CompositeBuffer::GetSize() const
{
	uint64_t Accum = 0;

	for (const SharedBuffer& It : m_Segments)
	{
		Accum += It.GetSize();
	}

	return Accum;
}

bool
CompositeBuffer::IsOwned() const
{
	for (const SharedBuffer& It : m_Segments)
	{
		if (It.IsOwned() == false)
		{
			return false;
		}
	}
	return true;
}

CompositeBuffer
CompositeBuffer::MakeOwned() const&
{
	return CompositeBuffer(*this).MakeOwned();
}

CompositeBuffer
CompositeBuffer::MakeOwned() &&
{
	for (SharedBuffer& Segment : m_Segments)
	{
		Segment = std::move(Segment).MakeOwned();
	}
	return std::move(*this);
}

SharedBuffer
CompositeBuffer::Flatten() const&
{
	switch (m_Segments.size())
	{
		case 0:
			return SharedBuffer();
		case 1:
			return m_Segments[0];
		default:
			UniqueBuffer	  Buffer  = UniqueBuffer::Alloc(GetSize());
			MutableMemoryView OutView = Buffer.GetMutableView();

			for (const SharedBuffer& Segment : m_Segments)
			{
				OutView.CopyFrom(Segment.GetView());
				OutView += Segment.GetSize();
			}

			return Buffer.MoveToShared();
	}
}

SharedBuffer
CompositeBuffer::Flatten() &&
{
	return m_Segments.size() == 1 ? std::move(m_Segments[0]) : std::as_const(*this).Flatten();
}

CompositeBuffer
CompositeBuffer::Mid(uint64_t Offset, uint64_t Size) const
{
	const uint64_t BufferSize = GetSize();
	Offset					  = Min(Offset, BufferSize);
	Size					  = Min(Size, BufferSize - Offset);
	CompositeBuffer Buffer;
	IterateRange(Offset, Size, [&Buffer](MemoryView View, const SharedBuffer& ViewOuter) {
		Buffer.m_Segments.push_back(SharedBuffer::MakeView(View, ViewOuter));
	});
	return Buffer;
}

MemoryView
CompositeBuffer::ViewOrCopyRange(uint64_t Offset, uint64_t Size, UniqueBuffer& CopyBuffer) const
{
	MemoryView View;
	IterateRange(Offset, Size, [Size, &View, &CopyBuffer, WriteView = MutableMemoryView()](MemoryView Segment) mutable {
		if (Size == Segment.GetSize())
		{
			View = Segment;
		}
		else
		{
			if (WriteView.IsEmpty())
			{
				if (CopyBuffer.GetSize() < Size)
				{
					CopyBuffer = UniqueBuffer::Alloc(Size);
				}
				View = WriteView = CopyBuffer.GetMutableView().Left(Size);
			}
			WriteView = WriteView.CopyFrom(Segment);
		}
	});
	return View;
}

void
CompositeBuffer::CopyTo(MutableMemoryView Target, uint64_t Offset) const
{
	IterateRange(Offset, Target.GetSize(), [Target](MemoryView View, [[maybe_unused]] const SharedBuffer& ViewOuter) mutable {
		Target = Target.CopyFrom(View);
	});
}

void
CompositeBuffer::IterateRange(uint64_t Offset, uint64_t Size, std::function<void(MemoryView View)> Visitor) const
{
	IterateRange(Offset, Size, [Visitor](MemoryView View, [[maybe_unused]] const SharedBuffer& ViewOuter) { Visitor(View); });
}

void
CompositeBuffer::IterateRange(uint64_t															  Offset,
							  uint64_t															  Size,
							  std::function<void(MemoryView View, const SharedBuffer& ViewOuter)> Visitor) const
{
	ZEN_ASSERT(Offset + Size <= GetSize());
	for (const SharedBuffer& Segment : m_Segments)
	{
		if (const uint64_t SegmentSize = Segment.GetSize(); Offset <= SegmentSize)
		{
			const MemoryView View = Segment.GetView().Mid(Offset, Size);
			Offset				  = 0;
			if (Size == 0 || !View.IsEmpty())
			{
				Visitor(View, Segment);
			}
			Size -= View.GetSize();
			if (Size == 0)
			{
				break;
			}
		}
		else
		{
			Offset -= SegmentSize;
		}
	}
}

#if ZEN_WITH_TESTS
TEST_CASE("CompositeBuffer Null")
{
	CompositeBuffer Buffer;
	CHECK(Buffer.IsNull());
	CHECK(Buffer.IsOwned());
	CHECK(Buffer.MakeOwned().IsNull());
	CHECK(Buffer.Flatten().IsNull());
	CHECK(Buffer.Mid(0, 0).IsNull());
	CHECK(Buffer.GetSize() == 0);
	CHECK(Buffer.GetSegments().size() == 0);

	UniqueBuffer CopyBuffer;
	CHECK(Buffer.ViewOrCopyRange(0, 0, CopyBuffer).IsEmpty());
	CHECK(CopyBuffer.IsNull());

	MutableMemoryView CopyView;
	Buffer.CopyTo(CopyView);

	uint32_t VisitCount = 0;
	Buffer.IterateRange(0, 0, [&VisitCount](MemoryView) { ++VisitCount; });
	CHECK(VisitCount == 0);
}

TEST_CASE("CompositeBuffer Empty")
{
	const uint8_t	   EmptyArray[]{0};
	const SharedBuffer EmptyView = SharedBuffer::MakeView(EmptyArray, 0);
	CompositeBuffer	   Buffer(EmptyView);
	CHECK(Buffer.IsNull() == false);
	CHECK(Buffer.IsOwned() == false);
	CHECK(Buffer.MakeOwned().IsNull() == false);
	CHECK(Buffer.MakeOwned().IsOwned() == true);
	CHECK(Buffer.Flatten() == EmptyView);
	CHECK(Buffer.Mid(0, 0).Flatten() == EmptyView);
	CHECK(Buffer.GetSize() == 0);
	CHECK(Buffer.GetSegments().size() == 1);
	CHECK(Buffer.GetSegments()[0] == EmptyView);

	UniqueBuffer CopyBuffer;
	CHECK(Buffer.ViewOrCopyRange(0, 0, CopyBuffer) == EmptyView.GetView());
	CHECK(CopyBuffer.IsNull());

	MutableMemoryView CopyView;
	Buffer.CopyTo(CopyView);

	uint32_t VisitCount = 0;
	Buffer.IterateRange(0, 0, [&VisitCount](MemoryView) { ++VisitCount; });
	CHECK(VisitCount == 1);
}

TEST_CASE("CompositeBuffer Empty[1]")
{
	const uint8_t	   EmptyArray[1]{};
	const SharedBuffer EmptyView1 = SharedBuffer::MakeView(EmptyArray, 0);
	const SharedBuffer EmptyView2 = SharedBuffer::MakeView(EmptyArray + 1, 0);
	CompositeBuffer	   Buffer(EmptyView1, EmptyView2);
	CHECK(Buffer.Mid(0, 0).Flatten() == EmptyView1);
	CHECK(Buffer.GetSize() == 0);
	CHECK(Buffer.GetSegments().size() == 2);
	CHECK(Buffer.GetSegments()[0] == EmptyView1);
	CHECK(Buffer.GetSegments()[1] == EmptyView2);

	UniqueBuffer CopyBuffer;
	CHECK(Buffer.ViewOrCopyRange(0, 0, CopyBuffer) == EmptyView1.GetView());
	CHECK(CopyBuffer.IsNull());

	MutableMemoryView CopyView;
	Buffer.CopyTo(CopyView);

	uint32_t VisitCount = 0;
	Buffer.IterateRange(0, 0, [&VisitCount](MemoryView) { ++VisitCount; });
	CHECK(VisitCount == 1);
}

TEST_CASE("CompositeBuffer Flat")
{
	const uint8_t	   FlatArray[]{1, 2, 3, 4, 5, 6, 7, 8};
	const SharedBuffer FlatView = SharedBuffer::Clone(MakeMemoryView(FlatArray));
	CompositeBuffer	   Buffer(FlatView);

	CHECK(Buffer.IsNull() == false);
	CHECK(Buffer.IsOwned() == true);
	CHECK(Buffer.Flatten() == FlatView);
	CHECK(Buffer.MakeOwned().Flatten() == FlatView);
	CHECK(Buffer.Mid(0).Flatten() == FlatView);
	CHECK(Buffer.Mid(4).Flatten().GetView() == FlatView.GetView().Mid(4));
	CHECK(Buffer.Mid(8).Flatten().GetView() == FlatView.GetView().Mid(8));
	CHECK(Buffer.Mid(4, 2).Flatten().GetView() == FlatView.GetView().Mid(4, 2));
	CHECK(Buffer.Mid(8, 0).Flatten().GetView() == FlatView.GetView().Mid(8, 0));
	CHECK(Buffer.GetSize() == sizeof(FlatArray));
	CHECK(Buffer.GetSegments().size() == 1);
	CHECK(Buffer.GetSegments()[0] == FlatView);

	UniqueBuffer CopyBuffer;
	CHECK(Buffer.ViewOrCopyRange(0, sizeof(FlatArray), CopyBuffer) == FlatView.GetView());
	CHECK(CopyBuffer.IsNull());

	uint8_t CopyArray[sizeof(FlatArray) - 3];
	Buffer.CopyTo(MakeMutableMemoryView(CopyArray), 3);
	CHECK(MakeMemoryView(CopyArray).EqualBytes(MakeMemoryView(FlatArray) + 3));

	uint32_t VisitCount = 0;
	Buffer.IterateRange(0, sizeof(FlatArray), [&VisitCount](MemoryView) { ++VisitCount; });
	CHECK(VisitCount == 1);
}

TEST_CASE("CompositeBuffer Composite")
{
	const uint8_t	   FlatArray[]{1, 2, 3, 4, 5, 6, 7, 8};
	const SharedBuffer FlatView1 = SharedBuffer::MakeView(MakeMemoryView(FlatArray).Left(4));
	const SharedBuffer FlatView2 = SharedBuffer::MakeView(MakeMemoryView(FlatArray).Right(4));
	CompositeBuffer	   Buffer(FlatView1, FlatView2);

	CHECK(Buffer.IsNull() == false);
	CHECK(Buffer.IsOwned() == false);
	CHECK(Buffer.Flatten().GetView().EqualBytes(MakeMemoryView(FlatArray)));
	CHECK(Buffer.Mid(2, 4).Flatten().GetView().EqualBytes(MakeMemoryView(FlatArray).Mid(2, 4)));
	CHECK(Buffer.Mid(0, 4).Flatten() == FlatView1);
	CHECK(Buffer.Mid(4, 4).Flatten() == FlatView2);
	CHECK(Buffer.GetSize() == sizeof(FlatArray));
	CHECK(Buffer.GetSegments().size() == 2);
	CHECK(Buffer.GetSegments()[0] == FlatView1);
	CHECK(Buffer.GetSegments()[1] == FlatView2);

	UniqueBuffer CopyBuffer;

	CHECK(Buffer.ViewOrCopyRange(0, 4, CopyBuffer) == FlatView1.GetView());
	CHECK(CopyBuffer.IsNull() == true);
	CHECK(Buffer.ViewOrCopyRange(4, 4, CopyBuffer) == FlatView2.GetView());
	CHECK(CopyBuffer.IsNull() == true);
	CHECK(Buffer.ViewOrCopyRange(3, 2, CopyBuffer).EqualBytes(MakeMemoryView(FlatArray).Mid(3, 2)));
	CHECK(CopyBuffer.GetSize() == 2);
	CHECK(Buffer.ViewOrCopyRange(1, 6, CopyBuffer).EqualBytes(MakeMemoryView(FlatArray).Mid(1, 6)));
	CHECK(CopyBuffer.GetSize() == 6);
	CHECK(Buffer.ViewOrCopyRange(2, 4, CopyBuffer).EqualBytes(MakeMemoryView(FlatArray).Mid(2, 4)));
	CHECK(CopyBuffer.GetSize() == 6);

	uint8_t CopyArray[4];
	Buffer.CopyTo(MakeMutableMemoryView(CopyArray), 2);
	CHECK(MakeMemoryView(CopyArray).EqualBytes(MakeMemoryView(FlatArray).Mid(2, 4)));

	uint32_t VisitCount = 0;
	Buffer.IterateRange(0, sizeof(FlatArray), [&VisitCount](MemoryView) { ++VisitCount; });
	CHECK(VisitCount == 2);

	const auto TestIterateRange =
		[&Buffer](uint64_t Offset, uint64_t Size, MemoryView ExpectedView, const SharedBuffer& ExpectedViewOuter) {
			uint32_t	 VisitCount = 0;
			MemoryView	 ActualView;
			SharedBuffer ActualViewOuter;
			Buffer.IterateRange(Offset, Size, [&VisitCount, &ActualView, &ActualViewOuter](MemoryView View, const SharedBuffer& ViewOuter) {
				++VisitCount;
				ActualView		= View;
				ActualViewOuter = ViewOuter;
			});
			CHECK(VisitCount == 1);
			CHECK(ActualView == ExpectedView);
			CHECK(ActualViewOuter == ExpectedViewOuter);
		};
	TestIterateRange(0, 4, MakeMemoryView(FlatArray).Mid(0, 4), FlatView1);
	TestIterateRange(4, 0, MakeMemoryView(FlatArray).Mid(4, 0), FlatView1);
	TestIterateRange(4, 4, MakeMemoryView(FlatArray).Mid(4, 4), FlatView2);
	TestIterateRange(8, 0, MakeMemoryView(FlatArray).Mid(8, 0), FlatView2);
}

void
compositebuffer_forcelink()
{
}
#endif

}  // namespace zen