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

#if ZEN_WITH_TESTS && ZEN_WITH_NOMAD
#	include "zenserver-test.h"
#	include <zencore/filesystem.h>
#	include <zencore/logging.h>
#	include <zencore/testing.h>
#	include <zencore/timer.h>
#	include <zenhttp/httpclient.h>
#	include <zennomad/nomadclient.h>
#	include <zennomad/nomadprocess.h>
#	include <zenutil/zenserverprocess.h>

#	include <fmt/format.h>

namespace zen::tests::nomad_tests {

using namespace std::literals;

TEST_CASE("nomad.client.lifecycle" * doctest::skip())
{
	zen::nomad::NomadProcess NomadProc;
	NomadProc.SpawnNomadAgent();

	zen::nomad::NomadTestClient Client("http://localhost:4646/");

	// Submit a simple batch job that sleeps briefly
#	if ZEN_PLATFORM_WINDOWS
	auto Job = Client.SubmitJob("zen-test-job", "cmd.exe", {"/C", "timeout /t 10 /nobreak"});
#	else
	auto Job = Client.SubmitJob("zen-test-job", "/bin/sleep", {"10"});
#	endif
	REQUIRE(!Job.Id.empty());
	CHECK_EQ(Job.Status, "pending");

	// Poll until the job is running (or dead)
	{
		Stopwatch Timer;
		bool	  FoundRunning = false;
		while (Timer.GetElapsedTimeMs() < 15000)
		{
			auto Status = Client.GetJobStatus("zen-test-job");
			if (Status.Status == "running")
			{
				FoundRunning = true;
				break;
			}
			if (Status.Status == "dead")
			{
				break;
			}
			Sleep(500);
		}
		CHECK(FoundRunning);
	}

	// Verify allocations exist
	auto Allocs = Client.GetAllocations("zen-test-job");
	CHECK(!Allocs.empty());

	// Stop the job
	Client.StopJob("zen-test-job");

	// Verify it reaches dead state
	{
		Stopwatch Timer;
		bool	  FoundDead = false;
		while (Timer.GetElapsedTimeMs() < 10000)
		{
			auto Status = Client.GetJobStatus("zen-test-job");
			if (Status.Status == "dead")
			{
				FoundDead = true;
				break;
			}
			Sleep(500);
		}
		CHECK(FoundDead);
	}

	NomadProc.StopNomadAgent();
}

TEST_CASE("nomad.provisioner.integration" * doctest::skip())
{
	zen::nomad::NomadProcess NomadProc;
	NomadProc.SpawnNomadAgent();

	// Spawn zenserver in compute mode with Nomad provisioning enabled
	ZenServerInstance Instance(TestEnv, ZenServerInstance::ServerMode::kComputeServer);

	Instance.SetDataDir(TestEnv.CreateNewTestDir());

	std::filesystem::path ZenServerPath = TestEnv.ProgramBaseDir() / "zenserver" ZEN_EXE_SUFFIX_LITERAL;

	std::string NomadArgs = fmt::format(
		"--nomad-enabled=true"
		" --nomad-server=http://localhost:4646"
		" --nomad-driver=raw_exec"
		" --nomad-binary-path={}"
		" --nomad-max-cores=32"
		" --nomad-cores-per-job=32",
		ZenServerPath.string());

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

	// Give the provisioner time to submit jobs.
	// The management thread has a 5s wait between cycles, and the HTTP client has
	// a 10s connect timeout, so we need to allow enough time for at least one full cycle.
	Sleep(15000);

	// Verify jobs were submitted to Nomad
	zen::nomad::NomadTestClient NomadClient("http://localhost:4646/");

	auto Jobs = NomadClient.ListJobs("zenserver-worker");

	ZEN_INFO("nomad.provisioner.integration: found {} jobs with prefix 'zenserver-worker'", Jobs.size());
	CHECK_MESSAGE(!Jobs.empty(), Instance.GetLogOutput());

	Instance.Shutdown();
	NomadProc.StopNomadAgent();
}

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