// 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>
#	include <zenutil/consul.h>
#	include <zencore/thread.h>
#	include <zencore/timer.h>

namespace zen::tests::hub {

using namespace std::literals;

TEST_SUITE_BEGIN("server.hub");

TEST_CASE("hub.lifecycle.basic")
{
	{
		ZenServerInstance Instance(TestEnv, ZenServerInstance::ServerMode::kHubServer);

		const uint16_t PortNumber = Instance.SpawnServerAndWaitUntilReady();
		CHECK(PortNumber != 0);

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

		HttpClient::Response Result = Client.Get("status");
		CHECK(Result);
	}
}

TEST_CASE("hub.lifecycle.children")
{
	ZenServerInstance Instance(TestEnv, ZenServerInstance::ServerMode::kHubServer);

	const uint16_t PortNumber = Instance.SpawnServerAndWaitUntilReady();
	REQUIRE(PortNumber != 0);

	SUBCASE("spawn")
	{
		HttpClient Client(Instance.GetBaseUri() + "/hub/");

		HttpClient::Response Result = Client.Get("status");
		REQUIRE(Result);

		{
			Result = Client.Post("modules/abc/provision");
			REQUIRE(Result);

			CbObject AbcResult = Result.AsObject();
			CHECK(AbcResult["moduleId"].AsString() == "abc"sv);
			const uint16_t AbcPort = AbcResult["port"].AsUInt16(0);
			CHECK_NE(AbcPort, 0);

			// This should be a fresh instance with no contents

			HttpClient AbcClient(fmt::format("http://localhost:{}", AbcPort));

			Result = AbcClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::NotFound);

			Result = AbcClient.Put("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567",
								   IoBufferBuilder::MakeFromMemory(MakeMemoryView("abcdef"sv)));
			CHECK_EQ(Result.StatusCode, HttpResponseCode::Created);
		}

		{
			Result = Client.Post("modules/def/provision");
			REQUIRE(Result);

			CbObject DefResult = Result.AsObject();
			CHECK(DefResult["moduleId"].AsString() == "def"sv);
			const uint16_t DefPort = DefResult["port"].AsUInt16(0);
			REQUIRE_NE(DefPort, 0);

			// This should be a fresh instance with no contents

			HttpClient DefClient(fmt::format("http://localhost:{}", DefPort));

			Result = DefClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::NotFound);

			Result = DefClient.Put("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567",
								   IoBufferBuilder::MakeFromMemory(MakeMemoryView("AbcDef"sv)));
			CHECK_EQ(Result.StatusCode, HttpResponseCode::Created);
		}

		// this should be rejected because of the invalid module id
		Result = Client.Post("modules/!!!!!/provision");
		CHECK(!Result);

		Result = Client.Post("modules/ghi/provision");
		REQUIRE(Result);

		// Tear down instances

		Result = Client.Post("modules/abc/deprovision");
		REQUIRE(Result);

		Result = Client.Post("modules/def/deprovision");
		REQUIRE(Result);

		Result = Client.Post("modules/ghi/deprovision");
		REQUIRE(Result);

		// re-provision to verify that (de)hydration preserved state
		{
			Result = Client.Post("modules/abc/provision");
			REQUIRE(Result);

			CbObject AbcResult = Result.AsObject();
			CHECK(AbcResult["moduleId"].AsString() == "abc"sv);
			const uint16_t AbcPort = AbcResult["port"].AsUInt16(0);
			REQUIRE_NE(AbcPort, 0);

			// This should contain the content from the previous run

			HttpClient AbcClient(fmt::format("http://localhost:{}", AbcPort));

			Result = AbcClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "abcdef"sv);

			Result = AbcClient.Put("/z$/ns1/b/1123456789abcdef0123456789abcdef01234567",
								   IoBufferBuilder::MakeFromMemory(MakeMemoryView("ghijklmnop"sv)));
			CHECK_EQ(Result.StatusCode, HttpResponseCode::Created);
		}

		{
			Result = Client.Post("modules/def/provision");
			REQUIRE(Result);

			CbObject DefResult = Result.AsObject();
			CHECK(DefResult["moduleId"].AsString() == "def"sv);
			const uint16_t DefPort = DefResult["port"].AsUInt16(0);
			REQUIRE_NE(DefPort, 0);

			// This should contain the content from the previous run

			HttpClient DefClient(fmt::format("http://localhost:{}", DefPort));

			Result = DefClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "AbcDef"sv);

			Result = DefClient.Put("/z$/ns1/b/1123456789abcdef0123456789abcdef01234567",
								   IoBufferBuilder::MakeFromMemory(MakeMemoryView("GhijklmNop"sv)));
			CHECK_EQ(Result.StatusCode, HttpResponseCode::Created);
		}

		Result = Client.Post("modules/abc/deprovision");
		REQUIRE(Result);

		Result = Client.Post("modules/def/deprovision");
		REQUIRE(Result);

		// re-provision to verify that (de)hydration preserved state, including
		// state which was generated after the very first dehydration
		{
			Result = Client.Post("modules/abc/provision");
			REQUIRE(Result);

			CbObject AbcResult = Result.AsObject();
			CHECK(AbcResult["moduleId"].AsString() == "abc"sv);
			const uint16_t AbcPort = AbcResult["port"].AsUInt16(0);
			REQUIRE_NE(AbcPort, 0);

			// This should contain the content from the previous two runs

			HttpClient AbcClient(fmt::format("http://localhost:{}", AbcPort));

			Result = AbcClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "abcdef"sv);

			Result = AbcClient.Get("/z$/ns1/b/1123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "ghijklmnop"sv);
		}

		{
			Result = Client.Post("modules/def/provision");
			REQUIRE(Result);

			CbObject DefResult = Result.AsObject();
			REQUIRE(DefResult["moduleId"].AsString() == "def"sv);
			const uint16_t DefPort = DefResult["port"].AsUInt16(0);
			REQUIRE_NE(DefPort, 0);

			// This should contain the content from the previous two runs

			HttpClient DefClient(fmt::format("http://localhost:{}", DefPort));

			Result = DefClient.Get("/z$/ns1/b/0123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "AbcDef"sv);

			Result = DefClient.Get("/z$/ns1/b/1123456789abcdef0123456789abcdef01234567");
			CHECK_EQ(Result.StatusCode, HttpResponseCode::OK);

			CHECK_EQ(Result.AsText(), "GhijklmNop"sv);
		}

		Result = Client.Post("modules/abc/deprovision");
		REQUIRE(Result);

		Result = Client.Post("modules/def/deprovision");
		REQUIRE(Result);

		// final sanity check that the hub is still responsive
		Result = Client.Get("status");
		CHECK(Result);
	}
}

static bool
WaitForConsulService(consul::ConsulClient& Client, std::string_view ServiceId, bool ExpectedState, int TimeoutMs)
{
	Stopwatch Timer;
	uint64_t  Index = 0;
	while (Timer.GetElapsedTimeMs() < static_cast<uint64_t>(TimeoutMs))
	{
		uint64_t RemainingMs = static_cast<uint64_t>(TimeoutMs) - Timer.GetElapsedTimeMs();
		int		 WaitSeconds = std::max(1, static_cast<int>(RemainingMs / 1000));
		if (Client.WatchService(ServiceId, Index, WaitSeconds) == ExpectedState)
		{
			return true;
		}
		if (Index == 0)
		{
			Sleep(100);	 // error path only: avoid tight loop on persistent connection failure
		}
	}
	return Client.HasService(ServiceId) == ExpectedState;
}

TEST_CASE("hub.consul.kv")
{
	consul::ConsulProcess ConsulProc;
	ConsulProc.SpawnConsulAgent();

	consul::ConsulClient Client("http://localhost:8500/");
	Client.SetKeyValue("zen/hub/testkey", "testvalue");

	std::string RetrievedValue = Client.GetKeyValue("zen/hub/testkey");
	CHECK_EQ(RetrievedValue, "testvalue");

	Client.DeleteKey("zen/hub/testkey");

	ConsulProc.StopConsulAgent();
}

TEST_CASE("hub.consul.hub.registration")
{
	consul::ConsulProcess ConsulProc;
	ConsulProc.SpawnConsulAgent();

	ZenServerInstance Instance(TestEnv, ZenServerInstance::ServerMode::kHubServer);
	const uint16_t	  PortNumber =
		Instance.SpawnServerAndWaitUntilReady("--consul-endpoint=http://localhost:8500/ --instance-id=test-instance");
	REQUIRE(PortNumber != 0);

	consul::ConsulClient Client("http://localhost:8500/");

	REQUIRE(WaitForConsulService(Client, "zen-hub-test-instance", true, 5000));

	Instance.Shutdown();

	CHECK(!Client.HasService("zen-hub-test-instance"));

	ConsulProc.StopConsulAgent();
}

TEST_CASE("hub.consul.provision.registration")
{
	consul::ConsulProcess ConsulProc;
	ConsulProc.SpawnConsulAgent();

	ZenServerInstance Instance(TestEnv, ZenServerInstance::ServerMode::kHubServer);
	const uint16_t	  PortNumber =
		Instance.SpawnServerAndWaitUntilReady("--consul-endpoint=http://localhost:8500/ --instance-id=test-instance");
	REQUIRE(PortNumber != 0);

	consul::ConsulClient Client("http://localhost:8500/");

	REQUIRE(WaitForConsulService(Client, "zen-hub-test-instance", true, 5000));

	HttpClient HubClient(Instance.GetBaseUri() + "/hub/");

	HttpClient::Response Result = HubClient.Post("modules/testmod/provision");
	REQUIRE(Result);

	CHECK(Client.HasService("testmod"));
	{
		const uint16_t ModulePort = Result.AsObject()["port"].AsUInt16(0);
		REQUIRE(ModulePort != 0);

		std::string		JsonError;
		CbFieldIterator ServicesRoot = LoadCompactBinaryFromJson(Client.GetAgentServicesJson(), JsonError);
		REQUIRE(JsonError.empty());

		CbObjectView ServicesMap;
		for (CbFieldView F : ServicesRoot)
		{
			if (F.IsObject())
			{
				ServicesMap = F.AsObjectView();
			}
		}
		REQUIRE(ServicesMap);

		// Verify fields registered by OnProvisioned
		{
			CbObjectView ModService = ServicesMap["testmod"].AsObjectView();
			CHECK_EQ(ModService["ID"sv].AsString(), "testmod"sv);
			CHECK_EQ(ModService["Service"sv].AsString(), "zen-storage"sv);
			CHECK_EQ(ModService["Port"sv].AsDouble(0), double(ModulePort));
			bool FoundModuleTag	 = false;
			bool FoundHubTag	 = false;
			bool FoundVersionTag = false;
			for (CbFieldView Tag : ModService["Tags"].AsArrayView())
			{
				std::string_view TagStr = Tag.AsString();
				if (TagStr == "module:testmod"sv)
				{
					FoundModuleTag = true;
				}
				else if (TagStr == "zen-hub:zen-hub-test-instance"sv)
				{
					FoundHubTag = true;
				}
				else if (TagStr.substr(0, 8) == "version:"sv)
				{
					FoundVersionTag = true;
				}
			}
			CHECK(FoundModuleTag);
			CHECK(FoundHubTag);
			CHECK(FoundVersionTag);
		}

		// Verify fields registered by InitializeConsulRegistration
		{
			CbObjectView HubService = ServicesMap["zen-hub-test-instance"].AsObjectView();
			CHECK_EQ(HubService["ID"sv].AsString(), "zen-hub-test-instance"sv);
			CHECK_EQ(HubService["Service"sv].AsString(), "zen-hub"sv);
			CHECK_EQ(HubService["Port"sv].AsDouble(0), double(PortNumber));
		}
	}

	Result = HubClient.Post("modules/testmod/deprovision");
	REQUIRE(Result);

	CHECK(!Client.HasService("testmod"));

	Instance.Shutdown();

	ConsulProc.StopConsulAgent();
}

TEST_SUITE_END();

}  // namespace zen::tests::hub
#endif