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

#if ZEN_WITH_TESTS
#	include "zenserver-test.h"
#	include <zencore/memoryview.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>
#	include <zenhttp/httpclient.h>
#	include <zenutil/cloud/s3client.h>
#	include <zenutil/zenserverprocess.h>

ZEN_THIRD_PARTY_INCLUDES_START
#	include <tsl/robin_set.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::tests {

using namespace std::literals;

TEST_SUITE_BEGIN("server.objectstore");

TEST_CASE("objectstore.blobs")
{
	std::string_view Bucket = "bkt"sv;

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

		const uint16_t PortNumber = Instance.SpawnServerAndWaitUntilReady(fmt::format("--objectstore-enabled"));
		CHECK(PortNumber != 0);

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

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

			std::string ObjectPath = fmt::format("{}/{}.utoc",
												 CompressedBlobsHashes.back().ToHexString().substr(0, 2),
												 CompressedBlobsHashes.back().ToHexString());

			HttpClient::Response Result = Client.Put(fmt::format("bucket/{}/{}.utoc", Bucket, ObjectPath), Payload);
			CHECK(Result);
		}

		for (size_t I = 0; I < 5; I++)
		{
			std::string ObjectPath =
				fmt::format("{}/{}.utoc", CompressedBlobsHashes[I].ToHexString().substr(0, 2), CompressedBlobsHashes[I].ToHexString());
			HttpClient::Response Result = Client.Get(fmt::format("bucket/{}/{}.utoc", Bucket, ObjectPath));
			CHECK(Result);
			CHECK_EQ(Result.ResponsePayload.GetSize(), CompressedBlobsSizes[I]);
			IoHash			 RawHash;
			uint64_t		 RawSize;
			CompressedBuffer Compressed =
				CompressedBuffer::FromCompressed(SharedBuffer(std::move(Result.ResponsePayload)), RawHash, RawSize);
			CHECK(Compressed);
			CHECK_EQ(RawHash, CompressedBlobsHashes[I]);
			CHECK_EQ(RawSize, BlobsSizes[I]);
		}
	}
}

TEST_CASE("objectstore.s3client")
{
	ZenServerInstance Instance(TestEnv);
	const uint16_t	  Port = Instance.SpawnServerAndWaitUntilReady("--objectstore-enabled");
	CHECK_MESSAGE(Port != 0, Instance.GetLogOutput());

	// S3Client in path-style builds paths as /{bucket}/{key}.
	// The objectstore routes objects at bucket/{bucket}/{key} relative to its base.
	// Point the S3Client endpoint at {server}/obj/bucket so the paths line up.
	S3ClientOptions Opts;
	Opts.BucketName					 = "s3test";
	Opts.Region						 = "us-east-1";
	Opts.Endpoint					 = fmt::format("http://localhost:{}/obj/bucket", Port);
	Opts.PathStyle					 = true;
	Opts.Credentials.AccessKeyId	 = "testkey";
	Opts.Credentials.SecretAccessKey = "testsecret";

	S3Client Client(Opts);

	// -- PUT + GET roundtrip --
	std::string_view TestData = "hello from s3client via objectstore"sv;
	IoBuffer		 Content  = IoBufferBuilder::MakeFromMemory(MakeMemoryView(TestData));
	S3Result		 PutRes	  = Client.PutObject("test/hello.txt", std::move(Content));
	REQUIRE_MESSAGE(PutRes.IsSuccess(), PutRes.Error);

	S3GetObjectResult GetRes = Client.GetObject("test/hello.txt");
	REQUIRE_MESSAGE(GetRes.IsSuccess(), GetRes.Error);
	CHECK(GetRes.AsText() == TestData);

	// -- PUT overwrites --
	IoBuffer Original  = IoBufferBuilder::MakeFromMemory(MakeMemoryView("original"sv));
	IoBuffer Overwrite = IoBufferBuilder::MakeFromMemory(MakeMemoryView("overwritten"sv));
	REQUIRE(Client.PutObject("overwrite/file.txt", std::move(Original)).IsSuccess());
	REQUIRE(Client.PutObject("overwrite/file.txt", std::move(Overwrite)).IsSuccess());

	S3GetObjectResult OverwriteGet = Client.GetObject("overwrite/file.txt");
	REQUIRE(OverwriteGet.IsSuccess());
	CHECK(OverwriteGet.AsText() == "overwritten"sv);

	// -- GET not found --
	S3GetObjectResult NotFoundGet = Client.GetObject("nonexistent/file.dat");
	CHECK_FALSE(NotFoundGet.IsSuccess());

	// -- HEAD found --
	std::string_view HeadData	 = "head test data"sv;
	IoBuffer		 HeadContent = IoBufferBuilder::MakeFromMemory(MakeMemoryView(HeadData));
	REQUIRE(Client.PutObject("head/meta.txt", std::move(HeadContent)).IsSuccess());

	S3HeadObjectResult HeadRes = Client.HeadObject("head/meta.txt");
	REQUIRE_MESSAGE(HeadRes.IsSuccess(), HeadRes.Error);
	CHECK(HeadRes.Status == HeadObjectResult::Found);
	CHECK(HeadRes.Info.Size == HeadData.size());

	// -- HEAD not found --
	S3HeadObjectResult HeadNotFound = Client.HeadObject("nonexistent/file.dat");
	CHECK(HeadNotFound.IsSuccess());
	CHECK(HeadNotFound.Status == HeadObjectResult::NotFound);

	// -- LIST objects --
	for (int i = 0; i < 3; ++i)
	{
		std::string Key		= fmt::format("listing/item-{}.txt", i);
		std::string Payload = fmt::format("content-{}", i);
		IoBuffer	Buf		= IoBufferBuilder::MakeFromMemory(MakeMemoryView(Payload));
		REQUIRE(Client.PutObject(Key, std::move(Buf)).IsSuccess());
	}

	S3ListObjectsResult ListRes = Client.ListObjects("listing/");
	REQUIRE_MESSAGE(ListRes.IsSuccess(), ListRes.Error);
	REQUIRE(ListRes.Objects.size() == 3);

	std::vector<std::string> Keys;
	for (const S3ObjectInfo& Obj : ListRes.Objects)
	{
		Keys.push_back(Obj.Key);
		CHECK(Obj.Size > 0);
	}
	std::sort(Keys.begin(), Keys.end());
	CHECK(Keys[0] == "listing/item-0.txt");
	CHECK(Keys[1] == "listing/item-1.txt");
	CHECK(Keys[2] == "listing/item-2.txt");

	// -- LIST empty prefix --
	S3ListObjectsResult EmptyList = Client.ListObjects("no-such-prefix/");
	REQUIRE(EmptyList.IsSuccess());
	CHECK(EmptyList.Objects.empty());
}

TEST_SUITE_END();

}  // namespace zen::tests
#endif