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

#include <zenutil/testartifactprovider.h>

#include <zencore/except_fmt.h>
#include <zencore/filesystem.h>
#include <zencore/logging.h>

#include <memory>
#include <system_error>

namespace zen {

namespace {

	std::string JoinKey(std::string_view Prefix, std::string_view RelativePath)
	{
		if (Prefix.empty())
		{
			return std::string(RelativePath);
		}
		std::string Result;
		Result.reserve(Prefix.size() + 1 + RelativePath.size());
		Result.append(Prefix);
		if (Result.back() != '/')
		{
			Result.push_back('/');
		}
		Result.append(RelativePath);
		return Result;
	}

	class LocalBackend
	{
	public:
		explicit LocalBackend(std::filesystem::path RootDir) : m_RootDir(std::move(RootDir)) {}

		const std::filesystem::path& RootDir() const { return m_RootDir; }

		bool Exists(std::string_view RelativePath) const
		{
			std::error_code Ec;
			return std::filesystem::exists(Resolve(RelativePath), Ec) && !Ec;
		}

		[[nodiscard]] TestArtifactFetchResult Fetch(std::string_view RelativePath) const
		{
			TestArtifactFetchResult Result;

			std::filesystem::path FullPath = Resolve(RelativePath);
			FileContents		  Contents = ReadFile(FullPath);
			if (!Contents)
			{
				Result.Error = fmt::format("failed to read '{}': {}", FullPath.string(), Contents.ErrorCode.message());
				return Result;
			}
			Result.Content = Contents.Flatten();
			return Result;
		}

		[[nodiscard]] TestArtifactResult Store(std::string_view RelativePath, IoBuffer Content) const
		{
			TestArtifactResult Result;

			std::filesystem::path FullPath	 = Resolve(RelativePath);
			std::filesystem::path ParentPath = FullPath.parent_path();
			if (!ParentPath.empty())
			{
				std::error_code Ec;
				std::filesystem::create_directories(ParentPath, Ec);
				if (Ec)
				{
					Result.Error = fmt::format("failed to create '{}': {}", ParentPath.string(), Ec.message());
					return Result;
				}
			}

			try
			{
				WriteFile(FullPath, std::move(Content));
			}
			catch (const std::exception& Ex)
			{
				Result.Error = fmt::format("failed to write '{}': {}", FullPath.string(), Ex.what());
			}
			return Result;
		}

		[[nodiscard]] TestArtifactListResult List(std::string_view Prefix) const
		{
			TestArtifactListResult Result;

			std::error_code Ec;
			if (!std::filesystem::exists(m_RootDir, Ec))
			{
				Result.Error = fmt::format("root directory does not exist: {}", m_RootDir.string());
				return Result;
			}

			std::filesystem::recursive_directory_iterator It(m_RootDir, Ec);
			if (Ec)
			{
				Result.Error = fmt::format("failed to iterate '{}': {}", m_RootDir.string(), Ec.message());
				return Result;
			}

			for (const std::filesystem::directory_entry& Entry : It)
			{
				if (!Entry.is_regular_file())
				{
					continue;
				}

				std::filesystem::path Relative = std::filesystem::relative(Entry.path(), m_RootDir, Ec);
				if (Ec)
				{
					continue;
				}

				std::string RelString = Relative.generic_string();
				if (!Prefix.empty() && RelString.compare(0, Prefix.size(), Prefix) != 0)
				{
					continue;
				}

				TestArtifactInfo Info;
				Info.RelativePath = std::move(RelString);
				Info.Size		  = Entry.file_size(Ec);
				Result.Artifacts.push_back(std::move(Info));
			}
			return Result;
		}

	private:
		std::filesystem::path Resolve(std::string_view RelativePath) const { return m_RootDir / std::filesystem::path(RelativePath); }

		std::filesystem::path m_RootDir;
	};

	class S3Backend
	{
	public:
		S3Backend(S3ClientOptions ClientOptions, std::string KeyPrefix)
		: m_ClientOptions(std::move(ClientOptions))
		, m_KeyPrefix(std::move(KeyPrefix))
		, m_Client(m_ClientOptions)
		{
		}

		std::string Describe() const { return fmt::format("s3:{}/{}", m_Client.BucketName(), m_KeyPrefix); }

		bool Exists(std::string_view RelativePath)
		{
			S3HeadObjectResult Head = m_Client.HeadObject(JoinKey(m_KeyPrefix, RelativePath));
			return Head.Status == HeadObjectResult::Found;
		}

		TestArtifactFetchResult Fetch(std::string_view RelativePath)
		{
			TestArtifactFetchResult Result;

			S3GetObjectResult Get = m_Client.GetObject(JoinKey(m_KeyPrefix, RelativePath));
			if (!Get)
			{
				Result.Error = std::move(Get.Error);
				return Result;
			}
			Result.Content = std::move(Get.Content);
			return Result;
		}

		TestArtifactListResult List(std::string_view Prefix)
		{
			TestArtifactListResult Result;

			std::string			FullPrefix = JoinKey(m_KeyPrefix, Prefix);
			S3ListObjectsResult Listing	   = m_Client.ListObjects(FullPrefix);
			if (!Listing)
			{
				Result.Error = std::move(Listing.Error);
				return Result;
			}

			size_t StripLen = m_KeyPrefix.size();
			if (StripLen > 0 && m_KeyPrefix.back() != '/')
			{
				StripLen += 1;	// account for the separator JoinKey inserts
			}

			Result.Artifacts.reserve(Listing.Objects.size());
			for (S3ObjectInfo& Obj : Listing.Objects)
			{
				TestArtifactInfo Info;
				Info.RelativePath = (StripLen <= Obj.Key.size()) ? Obj.Key.substr(StripLen) : std::move(Obj.Key);
				Info.Size		  = Obj.Size;
				Result.Artifacts.push_back(std::move(Info));
			}
			return Result;
		}

	private:
		S3ClientOptions m_ClientOptions;
		std::string		m_KeyPrefix;
		S3Client		m_Client;
	};

	class TestArtifactProviderImpl final : public TestArtifactProvider
	{
	public:
		TestArtifactProviderImpl(LocalBackend Cache, std::unique_ptr<S3Backend> Primary)
		: m_Cache(std::move(Cache))
		, m_Primary(std::move(Primary))
		{
		}

		std::string Describe() const override
		{
			if (m_Primary)
			{
				return fmt::format("cache:{} <- {}", m_Cache.RootDir().string(), m_Primary->Describe());
			}
			return fmt::format("local:{}", m_Cache.RootDir().string());
		}

		bool Exists(std::string_view RelativePath) override
		{
			if (m_Cache.Exists(RelativePath))
			{
				return true;
			}
			return m_Primary && m_Primary->Exists(RelativePath);
		}

		TestArtifactFetchResult Fetch(std::string_view RelativePath) override
		{
			if (m_Cache.Exists(RelativePath))
			{
				TestArtifactFetchResult Result = m_Cache.Fetch(RelativePath);
				if (Result)
				{
					ZEN_INFO("test artifact '{}' served from cache {} ({} bytes)",
							 RelativePath,
							 m_Cache.RootDir().string(),
							 Result.Content.GetSize());
				}
				return Result;
			}

			if (!m_Primary)
			{
				TestArtifactFetchResult Result;
				Result.Error = fmt::format("artifact '{}' not found in cache and no remote source configured", RelativePath);
				return Result;
			}

			ZEN_INFO("downloading test artifact '{}' from {}", RelativePath, m_Primary->Describe());

			TestArtifactFetchResult Fetched = m_Primary->Fetch(RelativePath);
			if (!Fetched)
			{
				return Fetched;
			}

			ZEN_INFO("test artifact '{}' fetched from {} ({} bytes)", RelativePath, m_Primary->Describe(), Fetched.Content.GetSize());

			TestArtifactResult StoreRes = m_Cache.Store(RelativePath, Fetched.Content);
			if (!StoreRes)
			{
				ZEN_WARN("failed to cache artifact '{}' into {}: {}", RelativePath, m_Cache.RootDir().string(), StoreRes.Error);
			}
			return Fetched;
		}

		TestArtifactListResult List(std::string_view Prefix) override
		{
			if (m_Primary)
			{
				return m_Primary->List(Prefix);
			}
			return m_Cache.List(Prefix);
		}

		TestArtifactResult Store(std::string_view RelativePath, IoBuffer Content) override
		{
			return m_Cache.Store(RelativePath, std::move(Content));
		}

	private:
		LocalBackend			   m_Cache;
		std::unique_ptr<S3Backend> m_Primary;
	};

	void ApplyS3UrlToOptions(std::string_view Url, S3ClientOptions& Client, std::string& KeyPrefix)
	{
		constexpr std::string_view kScheme = "s3://";
		if (Url.substr(0, kScheme.size()) == kScheme)
		{
			Url.remove_prefix(kScheme.size());
		}

		size_t			 Slash	= Url.find('/');
		std::string_view Bucket = (Slash == std::string_view::npos) ? Url : Url.substr(0, Slash);
		std::string_view Prefix = (Slash == std::string_view::npos) ? std::string_view{} : Url.substr(Slash + 1);

		if (Client.BucketName.empty() && !Bucket.empty())
		{
			Client.BucketName.assign(Bucket);
		}
		if (KeyPrefix.empty() && !Prefix.empty())
		{
			KeyPrefix.assign(Prefix);
		}
	}

	// Returns true when the caller has explicitly (or implicitly by platform) disabled
	// the EC2 Instance Metadata Service fallback. Honors the standard AWS env var
	// AWS_EC2_METADATA_DISABLED=true and skips by default on macOS, where Mac EC2
	// instances are rare and the link-local probe would just emit noise on failure.
	bool IsImdsDisabled()
	{
#if ZEN_PLATFORM_MAC
		return true;
#else
		std::string Disabled = GetEnvVariable("AWS_EC2_METADATA_DISABLED");
		return Disabled == "true" || Disabled == "TRUE" || Disabled == "1";
#endif
	}

	void ApplyAwsEnvDefaults(S3ClientOptions& Client)
	{
		if (std::string Region = GetEnvVariable("AWS_DEFAULT_REGION"); !Region.empty())
		{
			Client.Region = std::move(Region);
		}
		else if (std::string FallbackRegion = GetEnvVariable("AWS_REGION"); !FallbackRegion.empty())
		{
			Client.Region = std::move(FallbackRegion);
		}

		if (Client.Endpoint.empty())
		{
			Client.Endpoint = GetEnvVariable("AWS_ENDPOINT_URL");
		}

		if (Client.Credentials.AccessKeyId.empty() && !Client.CredentialProvider)
		{
			std::string AccessKeyId = GetEnvVariable("AWS_ACCESS_KEY_ID");
			if (!AccessKeyId.empty())
			{
				Client.Credentials.AccessKeyId	   = std::move(AccessKeyId);
				Client.Credentials.SecretAccessKey = GetEnvVariable("AWS_SECRET_ACCESS_KEY");
				Client.Credentials.SessionToken	   = GetEnvVariable("AWS_SESSION_TOKEN");
			}
			else if (!IsImdsDisabled())
			{
				// Fall back to the EC2 Instance Metadata Service so self-hosted runners
				// with an attached IAM role can sign S3 requests without static creds.
				Client.CredentialProvider = Ref<ImdsCredentialProvider>(new ImdsCredentialProvider({}));
			}
		}
	}

}  // namespace

std::filesystem::path
GetDefaultLocalTestArtifactPath()
{
	std::filesystem::path SystemRoot = PickDefaultSystemRootDirectory();
	if (SystemRoot.empty())
	{
		return {};
	}
	return SystemRoot / kDefaultLocalTestArtifactDirName;
}

bool
S3TestArtifactsAvailable()
{
	if (GetEnvVariable(kTestArtifactsS3EnvVar).empty())
	{
		return false;
	}
	if (!GetEnvVariable("AWS_ACCESS_KEY_ID").empty() || !GetEnvVariable("AWS_SESSION_TOKEN").empty())
	{
		return true;
	}
	return !IsImdsDisabled();
}

bool
TestArtifactsAvailable()
{
	if (!GetEnvVariable(kTestArtifactsPathEnvVar).empty())
	{
		return true;
	}
	return S3TestArtifactsAvailable();
}

Ref<TestArtifactProvider>
CreateTestArtifactProvider(TestArtifactProviderOptions Options)
{
	if (Options.CacheDir.empty())
	{
		std::string EnvValue = GetEnvVariable(kTestArtifactsPathEnvVar);
		if (!EnvValue.empty())
		{
			Options.CacheDir = std::filesystem::path(EnvValue);
		}
	}
	if (Options.CacheDir.empty())
	{
		Options.CacheDir = GetDefaultLocalTestArtifactPath();
	}
	if (Options.CacheDir.empty())
	{
		return {};
	}

	if (Options.S3Client.BucketName.empty() || Options.S3KeyPrefix.empty())
	{
		std::string EnvValue = GetEnvVariable(kTestArtifactsS3EnvVar);
		if (!EnvValue.empty())
		{
			ApplyS3UrlToOptions(EnvValue, Options.S3Client, Options.S3KeyPrefix);
		}
	}

	std::unique_ptr<S3Backend> Primary;
	if (!Options.S3Client.BucketName.empty())
	{
		ApplyAwsEnvDefaults(Options.S3Client);
		Primary = std::make_unique<S3Backend>(std::move(Options.S3Client), std::move(Options.S3KeyPrefix));
	}

	return Ref<TestArtifactProvider>(new TestArtifactProviderImpl(LocalBackend(std::move(Options.CacheDir)), std::move(Primary)));
}

void
testartifactprovider_forcelink()
{
}

}  // namespace zen

//////////////////////////////////////////////////////////////////////////
// Tests

#if ZEN_WITH_TESTS

#	include <zenutil/cloud/minioprocess.h>

#	include <zencore/memoryview.h>
#	include <zencore/testing.h>
#	include <zencore/testutils.h>

#	include <cstring>

namespace zen::tests {

using namespace std::literals;

namespace {

	bool ContentMatches(const IoBuffer& Buf, std::string_view Expected)
	{
		return Buf.GetSize() == Expected.size() && std::memcmp(Buf.GetData(), Expected.data(), Expected.size()) == 0;
	}

}  // namespace

TEST_SUITE_BEGIN("util.testartifactprovider");

TEST_CASE("local_only.store_fetch_list")
{
	ScopedTemporaryDirectory CacheDir;

	// Keep the test hermetic: ignore any S3 configuration the developer may have in the environment.
	ScopedEnvVar EnvS3(kTestArtifactsS3EnvVar, "");

	TestArtifactProviderOptions Opts;
	Opts.CacheDir					   = CacheDir.Path();
	Ref<TestArtifactProvider> Provider = CreateTestArtifactProvider(std::move(Opts));
	REQUIRE(Provider);

	CHECK_FALSE(Provider->Exists("missing.txt"));

	constexpr std::string_view kContent = "local payload"sv;
	TestArtifactResult		   StoreRes = Provider->Store("greet/hello.txt", IoBufferBuilder::MakeFromMemory(MakeMemoryView(kContent)));
	REQUIRE_MESSAGE(StoreRes.IsSuccess(), StoreRes.Error);

	CHECK(Provider->Exists("greet/hello.txt"));

	TestArtifactFetchResult Fetch = Provider->Fetch("greet/hello.txt");
	REQUIRE_MESSAGE(Fetch.IsSuccess(), Fetch.Error);
	CHECK(ContentMatches(Fetch.Content, kContent));

	TestArtifactFetchResult Missing = Provider->Fetch("nope.txt");
	CHECK_FALSE(Missing.IsSuccess());

	TestArtifactListResult ListRes = Provider->List("");
	REQUIRE_MESSAGE(ListRes.IsSuccess(), ListRes.Error);
	REQUIRE_EQ(ListRes.Artifacts.size(), 1u);
	CHECK_EQ(ListRes.Artifacts.front().RelativePath, "greet/hello.txt");
}

TEST_CASE("minio.s3_primary")
{
	MinioProcessOptions MinioOpts;
	MinioOpts.Port = 19020;
	MinioProcess Minio(MinioOpts);
	Minio.SpawnMinioServer();
	Minio.CreateBucket("artifacts-test");

	// Seed S3 directly so we can verify the provider's read path.
	S3ClientOptions SeedOpts;
	SeedOpts.BucketName					 = "artifacts-test";
	SeedOpts.Region						 = "us-east-1";
	SeedOpts.Endpoint					 = Minio.Endpoint();
	SeedOpts.PathStyle					 = true;
	SeedOpts.Credentials.AccessKeyId	 = std::string(Minio.RootUser());
	SeedOpts.Credentials.SecretAccessKey = std::string(Minio.RootPassword());
	S3Client SeedClient(SeedOpts);

	constexpr std::string_view kHello = "hello from minio"sv;
	REQUIRE(SeedClient.PutObject("artifacts/hello.txt", IoBufferBuilder::MakeFromMemory(MakeMemoryView(kHello))).IsSuccess());

	ScopedTemporaryDirectory CacheDir;

	// Configure everything via environment variables to exercise the env-based defaults.
	ScopedEnvVar EnvPath(kTestArtifactsPathEnvVar, CacheDir.Path().string());
	ScopedEnvVar EnvS3(kTestArtifactsS3EnvVar, "s3://artifacts-test/artifacts");
	ScopedEnvVar EnvEndpoint("AWS_ENDPOINT_URL", Minio.Endpoint());
	ScopedEnvVar EnvRegion("AWS_DEFAULT_REGION", "us-east-1");
	ScopedEnvVar EnvAccessKey("AWS_ACCESS_KEY_ID", Minio.RootUser());
	ScopedEnvVar EnvSecretKey("AWS_SECRET_ACCESS_KEY", Minio.RootPassword());

	TestArtifactProviderOptions Opts;
	Opts.S3Client.PathStyle			   = true;	// MinIO requires path-style addressing
	Ref<TestArtifactProvider> Provider = CreateTestArtifactProvider(std::move(Opts));
	REQUIRE(Provider);

	// -- exists checks consult primary when cache is cold --------------------
	CHECK(Provider->Exists("hello.txt"));
	CHECK_FALSE(Provider->Exists("nope.txt"));

	// -- first fetch pulls from primary and populates the cache --------------
	TestArtifactFetchResult First = Provider->Fetch("hello.txt");
	REQUIRE_MESSAGE(First.IsSuccess(), First.Error);
	CHECK(ContentMatches(First.Content, kHello));
	CHECK(std::filesystem::exists(CacheDir.Path() / "hello.txt"));

	// -- delete from S3, then fetch again: must come from cache --------------
	REQUIRE(SeedClient.DeleteObject("artifacts/hello.txt").IsSuccess());
	TestArtifactFetchResult Cached = Provider->Fetch("hello.txt");
	REQUIRE_MESSAGE(Cached.IsSuccess(), Cached.Error);
	CHECK(ContentMatches(Cached.Content, kHello));

	// -- list reflects the primary (S3) source ------------------------------
	REQUIRE(SeedClient.PutObject("artifacts/sub/file-a.bin", IoBufferBuilder::MakeFromMemory(MakeMemoryView("A"sv))).IsSuccess());
	REQUIRE(SeedClient.PutObject("artifacts/sub/file-b.bin", IoBufferBuilder::MakeFromMemory(MakeMemoryView("B"sv))).IsSuccess());

	TestArtifactListResult ListRes = Provider->List("");
	REQUIRE_MESSAGE(ListRes.IsSuccess(), ListRes.Error);

	auto HasPath = [&](std::string_view Rel) {
		for (const TestArtifactInfo& Info : ListRes.Artifacts)
		{
			if (Info.RelativePath == Rel)
			{
				return true;
			}
		}
		return false;
	};
	CHECK(HasPath("sub/file-a.bin"));
	CHECK(HasPath("sub/file-b.bin"));

	// -- fetching a missing artifact surfaces a remote error -----------------
	TestArtifactFetchResult MissingFetch = Provider->Fetch("does-not-exist.bin");
	CHECK_FALSE(MissingFetch.IsSuccess());
}

TEST_SUITE_END();

}  // namespace zen::tests

#endif