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

// Simple test bed for exercising the zens3 module against a real S3 bucket.
//
// Usage:
//   zens3-testbed --bucket <name> --region <region> [command] [args...]
//
// Credentials are read from environment variables:
//   AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY
//
// Commands:
//   put <key> <file>            Upload a local file
//   get <key> [file]            Download an object (prints to stdout if no file given)
//   head <key>                  Check if object exists, show metadata
//   delete <key>                Delete an object
//   list [prefix]               List objects with optional prefix
//   multipart-put <key> <file> [part-size-mb]  Upload via multipart
//   roundtrip <key>             Upload test data, download, verify, delete

#include <zenutil/cloud/imdscredentials.h>
#include <zenutil/cloud/s3client.h>

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

#include <zencore/memory/newdelete.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>
#include <cxxopts.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

#include <cstdlib>
#include <fstream>
#include <iostream>

namespace {

using namespace zen;

std::string
GetEnvVar(const char* Name)
{
	const char* Value = std::getenv(Name);
	return Value ? std::string(Value) : std::string();
}

IoBuffer
ReadFileToBuffer(const std::filesystem::path& Path)
{
	return zen::ReadFile(Path).Flatten();
}

void
WriteBufferToFile(const IoBuffer& Buffer, const std::filesystem::path& Path)
{
	std::ofstream File(Path, std::ios::binary);
	if (!File)
	{
		throw zen::runtime_error("failed to open '{}' for writing", Path.string());
	}
	File.write(reinterpret_cast<const char*>(Buffer.GetData()), static_cast<std::streamsize>(Buffer.GetSize()));
}

S3Client
CreateClient(const cxxopts::ParseResult& Args)
{
	S3ClientOptions Options;
	Options.BucketName = Args["bucket"].as<std::string>();
	Options.Region	   = Args["region"].as<std::string>();

	if (Args.count("imds"))
	{
		// Use IMDS credential provider for EC2 instances
		ImdsCredentialProviderOptions ImdsOpts;
		if (Args.count("imds-endpoint"))
		{
			ImdsOpts.Endpoint = Args["imds-endpoint"].as<std::string>();
		}
		Options.CredentialProvider = Ref<ImdsCredentialProvider>(new ImdsCredentialProvider(ImdsOpts));
	}
	else
	{
		std::string AccessKey	 = GetEnvVar("AWS_ACCESS_KEY_ID");
		std::string SecretKey	 = GetEnvVar("AWS_SECRET_ACCESS_KEY");
		std::string SessionToken = GetEnvVar("AWS_SESSION_TOKEN");

		if (AccessKey.empty() || SecretKey.empty())
		{
			throw zen::runtime_error("AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables must be set");
		}

		Options.Credentials.AccessKeyId		= std::move(AccessKey);
		Options.Credentials.SecretAccessKey = std::move(SecretKey);
		Options.Credentials.SessionToken	= std::move(SessionToken);
	}

	if (Args.count("endpoint"))
	{
		Options.Endpoint = Args["endpoint"].as<std::string>();
	}

	if (Args.count("path-style"))
	{
		Options.PathStyle = true;
	}

	if (Args.count("timeout"))
	{
		Options.Timeout = std::chrono::milliseconds(Args["timeout"].as<int>() * 1000);
	}

	return S3Client(Options);
}

int
CmdPut(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 3)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... put <key> <file>\n");
		return 1;
	}

	const auto& Key		 = Positional[1];
	const auto& FilePath = Positional[2];

	IoBuffer Content = ReadFileToBuffer(FilePath);
	fmt::print("Uploading '{}' ({} bytes) to s3://{}/{}\n", FilePath, Content.GetSize(), Client.BucketName(), Key);

	S3Result Result = Client.PutObject(Key, Content);
	if (!Result)
	{
		fmt::print(stderr, "PUT failed: {}\n", Result.Error);
		return 1;
	}

	fmt::print("OK\n");
	return 0;
}

int
CmdGet(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 2)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... get <key> [file]\n");
		return 1;
	}

	const auto& Key = Positional[1];

	S3GetObjectResult Result = Client.GetObject(Key);
	if (!Result)
	{
		fmt::print(stderr, "GET failed: {}\n", Result.Error);
		return 1;
	}

	if (Positional.size() >= 3)
	{
		const auto& FilePath = Positional[2];
		WriteBufferToFile(Result.Content, FilePath);
		fmt::print("Downloaded {} bytes to '{}'\n", Result.Content.GetSize(), FilePath);
	}
	else
	{
		// Print to stdout
		std::string_view Text = Result.AsText();
		std::cout.write(Text.data(), static_cast<std::streamsize>(Text.size()));
		std::cout << std::endl;
	}

	return 0;
}

int
CmdHead(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 2)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... head <key>\n");
		return 1;
	}

	const auto& Key = Positional[1];

	S3HeadObjectResult Result = Client.HeadObject(Key);

	if (!Result)
	{
		fmt::print(stderr, "HEAD failed: {}\n", Result.Error);
		return 1;
	}

	if (Result.Status == HeadObjectResult::NotFound)
	{
		fmt::print("Object '{}' does not exist\n", Key);
		return 1;
	}

	fmt::print("Key:           {}\n", Result.Info.Key);
	fmt::print("Size:          {} bytes\n", Result.Info.Size);
	fmt::print("ETag:          {}\n", Result.Info.ETag);
	fmt::print("Last-Modified: {}\n", Result.Info.LastModified);
	return 0;
}

int
CmdDelete(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 2)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... delete <key>\n");
		return 1;
	}

	const auto& Key = Positional[1];

	S3Result Result = Client.DeleteObject(Key);
	if (!Result)
	{
		fmt::print(stderr, "DELETE failed: {}\n", Result.Error);
		return 1;
	}

	fmt::print("Deleted '{}'\n", Key);
	return 0;
}

int
CmdList(S3Client& Client, const std::vector<std::string>& Positional)
{
	std::string Prefix;
	if (Positional.size() >= 2)
	{
		Prefix = Positional[1];
	}

	S3ListObjectsResult Result = Client.ListObjects(Prefix);
	if (!Result)
	{
		fmt::print(stderr, "LIST failed: {}\n", Result.Error);
		return 1;
	}

	fmt::print("{} objects found:\n", Result.Objects.size());
	for (const auto& Obj : Result.Objects)
	{
		fmt::print("  {:>12}  {}  {}\n", Obj.Size, Obj.LastModified, Obj.Key);
	}

	return 0;
}

int
CmdMultipartPut(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 3)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... multipart-put <key> <file> [part-size-mb]\n");
		return 1;
	}

	const auto& Key		 = Positional[1];
	const auto& FilePath = Positional[2];

	uint64_t PartSize = 8 * 1024 * 1024;  // 8 MB default
	if (Positional.size() >= 4)
	{
		PartSize = std::stoull(Positional[3]) * 1024 * 1024;
	}

	IoBuffer Content = ReadFileToBuffer(FilePath);
	fmt::print("Multipart uploading '{}' ({} bytes, part size {} MB) to s3://{}/{}\n",
			   FilePath,
			   Content.GetSize(),
			   PartSize / (1024 * 1024),
			   Client.BucketName(),
			   Key);

	S3Result Result = Client.PutObjectMultipart(Key, Content, PartSize);
	if (!Result)
	{
		fmt::print(stderr, "Multipart PUT failed: {}\n", Result.Error);
		return 1;
	}

	fmt::print("OK\n");
	return 0;
}

int
CmdRoundtrip(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 2)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... roundtrip <key>\n");
		return 1;
	}

	const auto& Key = Positional[1];

	// Generate test data
	const size_t		 TestSize = 1024 * 64;	// 64 KB
	std::vector<uint8_t> TestData(TestSize);
	for (size_t i = 0; i < TestSize; ++i)
	{
		TestData[i] = static_cast<uint8_t>(i & 0xFF);
	}

	IoBuffer UploadContent(IoBuffer::Clone, TestData.data(), TestData.size());

	fmt::print("=== Roundtrip test for key '{}' ===\n\n", Key);

	// PUT
	fmt::print("[1/4] PUT {} bytes...\n", TestSize);
	S3Result Result = Client.PutObject(Key, UploadContent);
	if (!Result)
	{
		fmt::print(stderr, "  FAILED: {}\n", Result.Error);
		return 1;
	}
	fmt::print("  OK\n");

	// HEAD
	fmt::print("[2/4] HEAD...\n");
	S3HeadObjectResult HeadResult = Client.HeadObject(Key);
	if (HeadResult.Status != HeadObjectResult::Found)
	{
		fmt::print(stderr, "  FAILED: {}\n", !HeadResult ? HeadResult.Error : "not found");
		return 1;
	}
	fmt::print("  OK (size={}, etag={})\n", HeadResult.Info.Size, HeadResult.Info.ETag);

	if (HeadResult.Info.Size != TestSize)
	{
		fmt::print(stderr, "  SIZE MISMATCH: expected {}, got {}\n", TestSize, HeadResult.Info.Size);
		return 1;
	}

	// GET
	fmt::print("[3/4] GET and verify...\n");
	S3GetObjectResult GetResult = Client.GetObject(Key);
	if (!GetResult)
	{
		fmt::print(stderr, "  FAILED: {}\n", GetResult.Error);
		return 1;
	}

	if (GetResult.Content.GetSize() != TestSize)
	{
		fmt::print(stderr, "  SIZE MISMATCH: expected {}, got {}\n", TestSize, GetResult.Content.GetSize());
		return 1;
	}

	if (memcmp(GetResult.Content.GetData(), TestData.data(), TestSize) != 0)
	{
		fmt::print(stderr, "  DATA MISMATCH\n");
		return 1;
	}
	fmt::print("  OK (verified {} bytes)\n", TestSize);

	// DELETE
	fmt::print("[4/4] DELETE...\n");
	Result = Client.DeleteObject(Key);
	if (!Result)
	{
		fmt::print(stderr, "  FAILED: {}\n", Result.Error);
		return 1;
	}
	fmt::print("  OK\n");

	fmt::print("\n=== Roundtrip test PASSED ===\n");
	return 0;
}

int
CmdPresign(S3Client& Client, const std::vector<std::string>& Positional)
{
	if (Positional.size() < 2)
	{
		fmt::print(stderr, "Usage: zens3-testbed ... presign <key> [method] [expires-seconds]\n");
		return 1;
	}

	const auto& Key = Positional[1];

	std::string Method = "GET";
	if (Positional.size() >= 3)
	{
		Method = Positional[2];
	}

	std::chrono::seconds ExpiresIn(3600);
	if (Positional.size() >= 4)
	{
		ExpiresIn = std::chrono::seconds(std::stoul(Positional[3]));
	}

	std::string Url;
	if (Method == "PUT")
	{
		Url = Client.GeneratePresignedPutUrl(Key, ExpiresIn);
	}
	else
	{
		Url = Client.GeneratePresignedGetUrl(Key, ExpiresIn);
	}

	fmt::print("{}\n", Url);
	return 0;
}

}  // namespace

int
main(int argc, char* argv[])
{
	using namespace zen;

	logging::InitializeLogging();

	cxxopts::Options Options("zens3-testbed", "Test bed for exercising S3 operations via the zens3 module");

	// clang-format off
	Options.add_options()
		("b,bucket",     "S3 bucket name",							cxxopts::value<std::string>())
		("r,region",     "AWS region",								cxxopts::value<std::string>()->default_value("us-east-1"))
		("e,endpoint",   "Custom S3 endpoint URL",					cxxopts::value<std::string>())
		("path-style",   "Use path-style addressing (for MinIO, etc.)")
		("imds",         "Use EC2 IMDS for credentials instead of env vars")
		("imds-endpoint", "Custom IMDS endpoint URL (for testing)",	cxxopts::value<std::string>())
		("timeout",      "Request timeout in seconds",				cxxopts::value<int>()->default_value("30"))
		("v,verbose",    "Enable verbose logging")
		("h,help",       "Show help")
		("positional",   "Command and arguments",					cxxopts::value<std::vector<std::string>>());
	// clang-format on

	Options.parse_positional({"positional"});
	Options.positional_help("<command> [args...]");

	try
	{
		auto Result = Options.parse(argc, argv);

		if (Result.count("help") || !Result.count("positional"))
		{
			fmt::print("{}\n", Options.help());
			fmt::print("Commands:\n");
			fmt::print("  put <key> <file>                        Upload a local file\n");
			fmt::print("  get <key> [file]                        Download (to file or stdout)\n");
			fmt::print("  head <key>                              Show object metadata\n");
			fmt::print("  delete <key>                            Delete an object\n");
			fmt::print("  list [prefix]                           List objects\n");
			fmt::print("  multipart-put <key> <file> [part-mb]    Multipart upload\n");
			fmt::print("  roundtrip <key>                         Upload/download/verify/delete\n");
			fmt::print("  presign <key> [method] [expires-sec]     Generate pre-signed URL\n");
			fmt::print("\nCredentials via AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY env vars,\n");
			fmt::print("or use --imds to fetch from EC2 Instance Metadata Service.\n");
			return 0;
		}

		if (!Result.count("bucket"))
		{
			fmt::print(stderr, "Error: --bucket is required\n");
			return 1;
		}

		if (Result.count("verbose"))
		{
			logging::SetLogLevel(logging::Debug);
		}

		auto Client = CreateClient(Result);

		const auto& Positional = Result["positional"].as<std::vector<std::string>>();
		const auto& Command	   = Positional[0];

		if (Command == "put")
		{
			return CmdPut(Client, Positional);
		}
		else if (Command == "get")
		{
			return CmdGet(Client, Positional);
		}
		else if (Command == "head")
		{
			return CmdHead(Client, Positional);
		}
		else if (Command == "delete")
		{
			return CmdDelete(Client, Positional);
		}
		else if (Command == "list")
		{
			return CmdList(Client, Positional);
		}
		else if (Command == "multipart-put")
		{
			return CmdMultipartPut(Client, Positional);
		}
		else if (Command == "roundtrip")
		{
			return CmdRoundtrip(Client, Positional);
		}
		else if (Command == "presign")
		{
			return CmdPresign(Client, Positional);
		}
		else
		{
			fmt::print(stderr, "Unknown command: '{}'\n", Command);
			return 1;
		}
	}
	catch (const std::exception& Ex)
	{
		fmt::print(stderr, "Error: {}\n", Ex.what());
		return 1;
	}
}