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

#include <zenutil/cloud/s3requestbuilder.h>

#include <zencore/fmtutils.h>

#include <algorithm>

namespace zen {

S3RequestBuilder::S3RequestBuilder(std::string Region, std::string BucketName, std::string Endpoint, bool PathStyle)
: m_Region(std::move(Region))
, m_BucketName(std::move(BucketName))
, m_Endpoint(DeriveEndpoint(Endpoint, m_Region, m_BucketName, PathStyle))
, m_Host(HostFromEndpoint(m_Endpoint))
, m_PathStyle(PathStyle)
{
}

std::string
S3RequestBuilder::DeriveEndpoint(std::string_view Endpoint, std::string_view Region, std::string_view BucketName, bool PathStyle)
{
	if (!Endpoint.empty())
	{
		return std::string(Endpoint);
	}
	if (PathStyle)
	{
		return fmt::format("https://s3.{}.amazonaws.com", Region);
	}
	return fmt::format("https://{}.s3.{}.amazonaws.com", BucketName, Region);
}

std::string
S3RequestBuilder::HostFromEndpoint(std::string_view Endpoint)
{
	std::string_view Ep = Endpoint;
	if (size_t Pos = Ep.find("://"); Pos != std::string_view::npos)
	{
		Ep = Ep.substr(Pos + 3);
	}
	if (!Ep.empty() && Ep.back() == '/')
	{
		Ep = Ep.substr(0, Ep.size() - 1);
	}
	return std::string(Ep);
}

std::string
S3RequestBuilder::KeyToPath(std::string_view Key) const
{
	if (m_PathStyle)
	{
		return fmt::format("/{}/{}", m_BucketName, Key);
	}
	return fmt::format("/{}", Key);
}

std::string
S3RequestBuilder::BucketRootPath() const
{
	if (m_PathStyle)
	{
		return fmt::format("/{}/", m_BucketName);
	}
	return "/";
}

// Cached for the (DateStamp, AccessKeyId) tuple. DateStamp rolls over at UTC
// midnight; concurrent callers around the rollover may each compute a fresh
// key once before the cache settles on the new day's value. Harmless extra
// HMACs, no signature corruption.
Sha256Digest
S3RequestBuilder::GetSigningKey(std::string_view DateStamp, const SigV4Credentials& Credentials)
{
	{
		RwLock::SharedLockScope _(m_SigningKeyLock);
		if (m_CachedDateStamp == DateStamp && m_CachedAccessKeyId == Credentials.AccessKeyId)
		{
			return m_CachedSigningKey;
		}
	}

	RwLock::ExclusiveLockScope _(m_SigningKeyLock);
	if (m_CachedDateStamp == DateStamp && m_CachedAccessKeyId == Credentials.AccessKeyId)
	{
		return m_CachedSigningKey;
	}

	std::string	 SecretPrefix = fmt::format("AWS4{}", Credentials.SecretAccessKey);
	Sha256Digest DateKey	  = ComputeHmacSha256(SecretPrefix.data(), SecretPrefix.size(), DateStamp.data(), DateStamp.size());
	SecureZeroSecret(SecretPrefix.data(), SecretPrefix.size());

	Sha256Digest RegionKey	= ComputeHmacSha256(DateKey, m_Region);
	Sha256Digest ServiceKey = ComputeHmacSha256(RegionKey, "s3");
	m_CachedSigningKey		= ComputeHmacSha256(ServiceKey, "aws4_request");
	m_CachedDateStamp		= std::string(DateStamp);
	m_CachedAccessKeyId		= Credentials.AccessKeyId;

	return m_CachedSigningKey;
}

HttpClient::KeyValueMap
S3RequestBuilder::SignRequest(const SigV4Credentials&							   Credentials,
							  std::string_view									   Method,
							  std::string_view									   Path,
							  std::string_view									   CanonicalQueryString,
							  std::string_view									   PayloadHash,
							  std::span<const std::pair<std::string, std::string>> ExtraSignedHeaders)
{
	std::string AmzDate = GetAmzTimestamp();

	std::vector<std::pair<std::string, std::string>> HeadersToSign;
	HeadersToSign.reserve(4 + ExtraSignedHeaders.size());
	HeadersToSign.emplace_back("host", m_Host);
	HeadersToSign.emplace_back("x-amz-content-sha256", std::string(PayloadHash));
	HeadersToSign.emplace_back("x-amz-date", AmzDate);
	if (!Credentials.SessionToken.empty())
	{
		HeadersToSign.emplace_back("x-amz-security-token", Credentials.SessionToken);
	}
	for (const auto& [K, V] : ExtraSignedHeaders)
	{
		HeadersToSign.emplace_back(K, V);
	}
	std::sort(HeadersToSign.begin(), HeadersToSign.end());

	std::string_view DateStamp(AmzDate.data(), 8);
	Sha256Digest	 SigningKey = GetSigningKey(DateStamp, Credentials);

	SigV4SignedHeaders Signed =
		SignRequestV4(Credentials, Method, Path, CanonicalQueryString, m_Region, "s3", AmzDate, HeadersToSign, PayloadHash, &SigningKey);

	HttpClient::KeyValueMap Result;
	Result->emplace("Authorization", std::move(Signed.Authorization));
	Result->emplace("x-amz-date", std::move(Signed.AmzDate));
	Result->emplace("x-amz-content-sha256", std::move(Signed.PayloadHash));
	if (!Credentials.SessionToken.empty())
	{
		Result->emplace("x-amz-security-token", Credentials.SessionToken);
	}
	for (const auto& [K, V] : ExtraSignedHeaders)
	{
		Result->emplace(K, V);
	}
	return Result;
}

}  // namespace zen