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

#if ZEN_WITH_TESTS
#	include "zenserver-test.h"
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#	include <zencore/workthreadpool.h>
#	include <zencore/compactbinarybuilder.h>
#	include <zencore/compactbinarypackage.h>
#	include <zencore/compress.h>
#	include <zencore/filesystem.h>
#	include <zencore/stream.h>
#	include <zencore/string.h>
#	include <zencore/fmtutils.h>
#	include <zencore/scopeguard.h>
#	include <zenhttp/packageformat.h>
#	include <zenremotestore/builds/buildstoragecache.h>
#	include <zenutil/workerpools.h>
#	include <zenutil/zenserverprocess.h>
#	include <zenhttp/httpclient.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <tsl/robin_set.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::tests {

using namespace std::literals;

TEST_CASE("buildstore.blobs")
{
	std::filesystem::path SystemRootPath = TestEnv.CreateNewTestDir();
	auto				  _				 = MakeGuard([&SystemRootPath]() { DeleteDirectories(SystemRootPath); });

	std::string_view Namespace = "ns"sv;
	std::string_view Bucket	   = "bkt"sv;
	Oid				 BuildId   = Oid::NewOid();

	std::vector<IoHash> CompressedBlobsHashes;
	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri() + "/builds/");

		for (size_t I = 0; I < 5; I++)
		{
			IoBuffer		 Blob			= CreateSemiRandomBlob(4711 + I * 7);
			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
			Payload.SetContentType(ZenContentType::kCompressedBinary);

			HttpClient::Response Result =
				Client.Put(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, CompressedBlobsHashes.back()), Payload);
			CHECK(Result);
		}

		for (const IoHash& RawHash : CompressedBlobsHashes)
		{
			HttpClient::Response Result = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
													 HttpClient::Accept(ZenContentType::kCompressedBinary));
			REQUIRE(Result);
			IoBuffer Payload = Result.ResponsePayload;
			REQUIRE(Payload.GetContentType() == ZenContentType::kCompressedBinary);
			IoHash			 VerifyRawHash;
			uint64_t		 VerifyRawSize;
			CompressedBuffer CompressedBlob =
				CompressedBuffer::FromCompressed(SharedBuffer(std::move(Payload)), VerifyRawHash, VerifyRawSize);
			REQUIRE(CompressedBlob);
			REQUIRE(VerifyRawHash == RawHash);
			IoBuffer Decompressed = CompressedBlob.Decompress().AsIoBuffer();
			CHECK(IoHash::HashBuffer(Decompressed) == RawHash);
		}
	}
	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri() + "/builds/");

		for (const IoHash& RawHash : CompressedBlobsHashes)
		{
			HttpClient::Response Result = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
													 HttpClient::Accept(ZenContentType::kCompressedBinary));
			REQUIRE(Result);
			IoBuffer Payload = Result.ResponsePayload;
			REQUIRE(Payload.GetContentType() == ZenContentType::kCompressedBinary);
			IoHash			 VerifyRawHash;
			uint64_t		 VerifyRawSize;
			CompressedBuffer CompressedBlob =
				CompressedBuffer::FromCompressed(SharedBuffer(std::move(Payload)), VerifyRawHash, VerifyRawSize);
			REQUIRE(CompressedBlob);
			REQUIRE(VerifyRawHash == RawHash);
			IoBuffer Decompressed = CompressedBlob.Decompress().AsIoBuffer();
			CHECK(IoHash::HashBuffer(Decompressed) == RawHash);
		}

		for (size_t I = 0; I < 5; I++)
		{
			IoBuffer		 Blob			= CreateSemiRandomBlob(5713 + I * 7);
			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
			CompressedBlobsHashes.push_back(CompressedBlob.DecodeRawHash());
			IoBuffer Payload = std::move(CompressedBlob).GetCompressed().Flatten().AsIoBuffer();
			Payload.SetContentType(ZenContentType::kCompressedBinary);

			HttpClient::Response Result =
				Client.Put(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, CompressedBlobsHashes.back()), Payload);
			CHECK(Result);
		}
	}
	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri() + "/builds/");

		for (const IoHash& RawHash : CompressedBlobsHashes)
		{
			HttpClient::Response Result = Client.Get(fmt::format("{}/{}/{}/blobs/{}", Namespace, Bucket, BuildId, RawHash),
													 HttpClient::Accept(ZenContentType::kCompressedBinary));
			REQUIRE(Result);
			IoBuffer Payload = Result.ResponsePayload;
			REQUIRE(Payload.GetContentType() == ZenContentType::kCompressedBinary);
			IoHash			 VerifyRawHash;
			uint64_t		 VerifyRawSize;
			CompressedBuffer CompressedBlob =
				CompressedBuffer::FromCompressed(SharedBuffer(std::move(Payload)), VerifyRawHash, VerifyRawSize);
			REQUIRE(CompressedBlob);
			REQUIRE(VerifyRawHash == RawHash);
			IoBuffer Decompressed = CompressedBlob.Decompress().AsIoBuffer();
			CHECK(IoHash::HashBuffer(Decompressed) == RawHash);
		}
	}
}

namespace {
	CbObject MakeMetadata(const IoHash& BlobHash, const std::vector<std::pair<std::string, std::string>>& KeyValues)
	{
		CbObjectWriter Writer;
		Writer.AddHash("rawHash"sv, BlobHash);
		Writer.BeginObject("values");
		{
			for (const auto& V : KeyValues)
			{
				Writer.AddString(V.first, V.second);
			}
		}
		Writer.EndObject();	 // values
		return Writer.Save();
	};

}  // namespace

TEST_CASE("buildstore.metadata")
{
	std::filesystem::path SystemRootPath = TestEnv.CreateNewTestDir();
	auto				  _				 = MakeGuard([&SystemRootPath]() { DeleteDirectories(SystemRootPath); });

	std::string_view Namespace = "ns"sv;
	std::string_view Bucket	   = "bkt"sv;
	Oid				 BuildId   = Oid::NewOid();

	std::vector<IoHash>	  BlobHashes;
	std::vector<CbObject> Metadatas;
	std::vector<IoHash>	  MetadataHashes;

	auto GetMetadatas =
		[](HttpClient& Client, std::string_view Namespace, std::string_view Bucket, const Oid& BuildId, std::vector<IoHash> BlobHashes) {
			CbObjectWriter Request;

			Request.BeginArray("blobHashes"sv);
			for (const IoHash& BlobHash : BlobHashes)
			{
				Request.AddHash(BlobHash);
			}
			Request.EndArray();

			IoBuffer Payload = Request.Save().GetBuffer().AsIoBuffer();
			Payload.SetContentType(ZenContentType::kCbObject);

			HttpClient::Response Result = Client.Post(fmt::format("{}/{}/{}/blobs/getBlobMetadata", Namespace, Bucket, BuildId),
													  Payload,
													  HttpClient::Accept(ZenContentType::kCbObject));
			CHECK(Result);

			std::vector<CbObject> ResultMetadatas;

			CbPackage ResponsePackage = ParsePackageMessage(Result.ResponsePayload);
			CbObject  ResponseObject  = ResponsePackage.GetObject();

			CbArrayView BlobHashArray  = ResponseObject["blobHashes"sv].AsArrayView();
			CbArrayView MetadatasArray = ResponseObject["metadatas"sv].AsArrayView();
			ResultMetadatas.reserve(MetadatasArray.Num());
			auto BlobHashesIt	 = BlobHashes.begin();
			auto BlobHashArrayIt = begin(BlobHashArray);
			auto MetadataArrayIt = begin(MetadatasArray);
			while (MetadataArrayIt != end(MetadatasArray))
			{
				const IoHash BlobHash = (*BlobHashArrayIt).AsHash();
				while (BlobHash != *BlobHashesIt)
				{
					ZEN_ASSERT(BlobHashesIt != BlobHashes.end());
					BlobHashesIt++;
				}

				ZEN_ASSERT(BlobHash == *BlobHashesIt);

				const IoHash		MetaHash	   = (*MetadataArrayIt).AsAttachment();
				const CbAttachment* MetaAttachment = ResponsePackage.FindAttachment(MetaHash);
				ZEN_ASSERT(MetaAttachment);

				CbObject Metadata = MetaAttachment->AsObject();
				ResultMetadatas.emplace_back(std::move(Metadata));

				BlobHashArrayIt++;
				MetadataArrayIt++;
				BlobHashesIt++;
			}
			return ResultMetadatas;
		};

	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri() + "/builds/");

		const size_t BlobCount = 5;

		for (size_t I = 0; I < BlobCount; I++)
		{
			BlobHashes.push_back(IoHash::HashBuffer(&I, sizeof(I)));
			Metadatas.push_back(MakeMetadata(BlobHashes.back(), {{"index", fmt::format("{}", I)}}));
			MetadataHashes.push_back(IoHash::HashBuffer(Metadatas.back().GetBuffer().AsIoBuffer()));
		}

		{
			CbPackage							   RequestPackage;
			std::vector<CbAttachment>			   Attachments;
			tsl::robin_set<IoHash, IoHash::Hasher> AttachmentHashes;
			Attachments.reserve(BlobCount);
			AttachmentHashes.reserve(BlobCount);
			{
				CbObjectWriter RequestWriter;
				RequestWriter.BeginArray("blobHashes");
				for (size_t BlockHashIndex = 0; BlockHashIndex < BlobHashes.size(); BlockHashIndex++)
				{
					RequestWriter.AddHash(BlobHashes[BlockHashIndex]);
				}
				RequestWriter.EndArray();  // blobHashes

				RequestWriter.BeginArray("metadatas");
				for (size_t BlockHashIndex = 0; BlockHashIndex < BlobHashes.size(); BlockHashIndex++)
				{
					const IoHash ObjectHash = Metadatas[BlockHashIndex].GetHash();
					RequestWriter.AddBinaryAttachment(ObjectHash);
					if (!AttachmentHashes.contains(ObjectHash))
					{
						Attachments.push_back(CbAttachment(Metadatas[BlockHashIndex], ObjectHash));
						AttachmentHashes.insert(ObjectHash);
					}
				}

				RequestWriter.EndArray();  // metadatas

				RequestPackage.SetObject(RequestWriter.Save());
			}
			RequestPackage.AddAttachments(Attachments);

			CompositeBuffer RpcRequestBuffer = FormatPackageMessageBuffer(RequestPackage);

			HttpClient::Response Result = Client.Post(fmt::format("{}/{}/{}/blobs/putBlobMetadata", Namespace, Bucket, BuildId),
													  RpcRequestBuffer,
													  ZenContentType::kCbPackage);
			CHECK(Result);
		}

		{
			std::vector<CbObject> ResultMetadatas = GetMetadatas(Client, Namespace, Bucket, BuildId, BlobHashes);

			for (size_t Index = 0; Index < MetadataHashes.size(); Index++)
			{
				const IoHash& ExpectedHash = MetadataHashes[Index];
				IoHash		  Hash		   = IoHash::HashBuffer(ResultMetadatas[Index].GetBuffer().AsIoBuffer());
				CHECK_EQ(ExpectedHash, Hash);
			}
		}
	}
	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri() + "/builds/");

		std::vector<CbObject> ResultMetadatas = GetMetadatas(Client, Namespace, Bucket, BuildId, BlobHashes);

		for (size_t Index = 0; Index < MetadataHashes.size(); Index++)
		{
			const IoHash& ExpectedHash = MetadataHashes[Index];
			IoHash		  Hash		   = IoHash::HashBuffer(ResultMetadatas[Index].GetBuffer().AsIoBuffer());
			CHECK_EQ(ExpectedHash, Hash);
		}
	}
}

TEST_CASE("buildstore.cache")
{
	std::filesystem::path SystemRootPath = TestEnv.CreateNewTestDir();
	std::filesystem::path TempDir		 = TestEnv.CreateNewTestDir();
	auto				  _				 = MakeGuard([&SystemRootPath, &TempDir]() {
		  DeleteDirectories(SystemRootPath);
		  DeleteDirectories(TempDir);
	  });

	std::string_view Namespace = "ns"sv;
	std::string_view Bucket	   = "bkt"sv;
	Oid				 BuildId   = Oid::NewOid();

	std::vector<IoHash>	  BlobHashes;
	std::vector<CbObject> Metadatas;
	std::vector<IoHash>	  MetadataHashes;

	const size_t BlobCount = 5;
	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri());

		BuildStorageCache::Statistics	   Stats;
		std::unique_ptr<BuildStorageCache> Cache(
			CreateZenBuildStorageCache(Client, Stats, Namespace, Bucket, TempDir, GetTinyWorkerPool(EWorkloadType::Background)));

		{
			IoHash											 NoneBlob	= IoHash::HashBuffer("data", 4);
			std::vector<BuildStorageCache::BlobExistsResult> NoneExists = Cache->BlobsExists(BuildId, std::vector<IoHash>{NoneBlob});
			REQUIRE(NoneExists.size() == 1);
			CHECK(!NoneExists[0].HasBody);
			CHECK(!NoneExists[0].HasMetadata);
		}

		for (size_t I = 0; I < BlobCount; I++)
		{
			IoBuffer		 Blob			= CreateSemiRandomBlob(4711 + I * 7);
			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
			BlobHashes.push_back(CompressedBlob.DecodeRawHash());
			Cache->PutBuildBlob(BuildId, BlobHashes.back(), ZenContentType::kCompressedBinary, CompressedBlob.GetCompressed());
		}

		Cache->Flush(500);

		Cache = CreateZenBuildStorageCache(Client, Stats, Namespace, Bucket, TempDir, GetTinyWorkerPool(EWorkloadType::Background));

		{
			std::vector<BuildStorageCache::BlobExistsResult> Exists = Cache->BlobsExists(BuildId, BlobHashes);
			CHECK(Exists.size() == BlobHashes.size());
			for (size_t I = 0; I < BlobCount; I++)
			{
				CHECK(Exists[I].HasBody);
				CHECK(!Exists[I].HasMetadata);
			}

			std::vector<CbObject> FetchedMetadatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
			CHECK_EQ(0, FetchedMetadatas.size());
		}

		{
			for (size_t I = 0; I < BlobCount; I++)
			{
				IoBuffer BuildBlob = Cache->GetBuildBlob(BuildId, BlobHashes[I]);
				CHECK(BuildBlob);
				CHECK_EQ(BlobHashes[I],
						 IoHash::HashBuffer(CompressedBuffer::FromCompressedNoValidate(std::move(BuildBlob)).Decompress().AsIoBuffer()));
			}
		}

		{
			for (size_t I = 0; I < BlobCount; I++)
			{
				CbObject Metadata = MakeMetadata(BlobHashes[I],
												 {{"key", fmt::format("{}", I)},
												  {"key_plus_one", fmt::format("{}", I + 1)},
												  {"block_hash", fmt::format("{}", BlobHashes[I])}});
				Metadatas.push_back(Metadata);
				MetadataHashes.push_back(IoHash::HashBuffer(Metadata.GetBuffer().AsIoBuffer()));
			}
			Cache->PutBlobMetadatas(BuildId, BlobHashes, Metadatas);
		}

		Cache->Flush(500);
		Cache = CreateZenBuildStorageCache(Client, Stats, Namespace, Bucket, TempDir, GetTinyWorkerPool(EWorkloadType::Background));

		{
			std::vector<BuildStorageCache::BlobExistsResult> Exists = Cache->BlobsExists(BuildId, BlobHashes);
			CHECK(Exists.size() == BlobHashes.size());
			for (size_t I = 0; I < BlobCount; I++)
			{
				CHECK(Exists[I].HasBody);
				CHECK(Exists[I].HasMetadata);
			}

			std::vector<CbObject> FetchedMetadatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
			CHECK_EQ(BlobCount, FetchedMetadatas.size());

			for (size_t I = 0; I < BlobCount; I++)
			{
				CHECK_EQ(MetadataHashes[I], IoHash::HashBuffer(FetchedMetadatas[I].GetBuffer().AsIoBuffer()));
			}
		}

		for (size_t I = 0; I < BlobCount; I++)
		{
			IoBuffer		 Blob			= CreateSemiRandomBlob(4711 + I * 7);
			CompressedBuffer CompressedBlob = CompressedBuffer::Compress(SharedBuffer(std::move(Blob)));
			BlobHashes.push_back(CompressedBlob.DecodeRawHash());
			Cache->PutBuildBlob(BuildId, BlobHashes.back(), ZenContentType::kCompressedBinary, CompressedBlob.GetCompressed());
		}

		Cache->Flush(500);
		Cache = CreateZenBuildStorageCache(Client, Stats, Namespace, Bucket, TempDir, GetTinyWorkerPool(EWorkloadType::Background));

		{
			std::vector<BuildStorageCache::BlobExistsResult> Exists = Cache->BlobsExists(BuildId, BlobHashes);
			CHECK(Exists.size() == BlobHashes.size());
			for (size_t I = 0; I < BlobCount * 2; I++)
			{
				CHECK(Exists[I].HasBody);
				CHECK_EQ(I < BlobCount, Exists[I].HasMetadata);
			}

			std::vector<CbObject> MetaDatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
			CHECK_EQ(BlobCount, MetaDatas.size());

			std::vector<CbObject> FetchedMetadatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
			CHECK_EQ(BlobCount, FetchedMetadatas.size());

			for (size_t I = 0; I < BlobCount; I++)
			{
				CHECK_EQ(MetadataHashes[I], IoHash::HashBuffer(FetchedMetadatas[I].GetBuffer().AsIoBuffer()));
			}
		}
	}

	{
		ZenServerInstance Instance(TestEnv);

		const uint16_t PortNumber =
			Instance.SpawnServerAndWaitUntilReady(fmt::format("--buildstore-enabled --system-dir {}", SystemRootPath));
		CHECK(PortNumber != 0);

		HttpClient Client(Instance.GetBaseUri());

		BuildStorageCache::Statistics	   Stats;
		std::unique_ptr<BuildStorageCache> Cache(
			CreateZenBuildStorageCache(Client, Stats, Namespace, Bucket, TempDir, GetTinyWorkerPool(EWorkloadType::Background)));

		std::vector<BuildStorageCache::BlobExistsResult> Exists = Cache->BlobsExists(BuildId, BlobHashes);
		CHECK(Exists.size() == BlobHashes.size());
		for (size_t I = 0; I < BlobCount * 2; I++)
		{
			CHECK(Exists[I].HasBody);
			CHECK_EQ(I < BlobCount, Exists[I].HasMetadata);
		}

		for (size_t I = 0; I < BlobCount * 2; I++)
		{
			IoBuffer BuildBlob = Cache->GetBuildBlob(BuildId, BlobHashes[I]);
			CHECK(BuildBlob);
			CHECK_EQ(BlobHashes[I],
					 IoHash::HashBuffer(CompressedBuffer::FromCompressedNoValidate(std::move(BuildBlob)).Decompress().AsIoBuffer()));
		}

		std::vector<CbObject> MetaDatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
		CHECK_EQ(BlobCount, MetaDatas.size());

		std::vector<CbObject> FetchedMetadatas = Cache->GetBlobMetadatas(BuildId, BlobHashes);
		CHECK_EQ(BlobCount, FetchedMetadatas.size());

		for (size_t I = 0; I < BlobCount; I++)
		{
			CHECK_EQ(MetadataHashes[I], IoHash::HashBuffer(FetchedMetadatas[I].GetBuffer().AsIoBuffer()));
		}
	}
}

}  // namespace zen::tests
#endif