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

#include "jupiter.h"

#include "cache/structuredcachestore.h"
#include "diag/formatters.h"
#include "diag/logging.h"

#include <fmt/format.h>
#include <zencore/compactbinary.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/string.h>
#include <zencore/thread.h>

// For some reason, these don't seem to stick, so we disable the warnings
//#	define _SILENCE_CXX17_C_HEADER_DEPRECATION_WARNING 1
//#	define _SILENCE_ALL_CXX17_DEPRECATION_WARNINGS		1
#pragma warning(push)
#pragma warning(disable : 4004)
#pragma warning(disable : 4996)
#include <cpr/cpr.h>
#pragma warning(pop)

#if ZEN_PLATFORM_WINDOWS
#	pragma comment(lib, "Crypt32.lib")
#	pragma comment(lib, "Wldap32.lib")
#endif

#include <json11.hpp>

using namespace std::literals;
using namespace fmt::literals;

namespace zen {

namespace detail {
	struct CloudCacheSessionState
	{
		CloudCacheSessionState(CloudCacheClient& Client) : OwnerClient(Client) {}
		~CloudCacheSessionState() {}

		void Reset()
		{
			std::string Auth;
			OwnerClient.AcquireAccessToken(Auth);

			Session.SetBody({});
			Session.SetOption(cpr::Header{{"Authorization", Auth}});
		}

		CloudCacheClient& OwnerClient;
		cpr::Session	  Session;
	};

}  // namespace detail

CloudCacheSession::CloudCacheSession(CloudCacheClient* OuterClient) : m_Log(OuterClient->Logger()), m_CacheClient(OuterClient)
{
	m_SessionState = m_CacheClient->AllocSessionState();
}

CloudCacheSession::~CloudCacheSession()
{
	m_CacheClient->FreeSessionState(m_SessionState);
}

CloudCacheResult
CloudCacheSession::Authenticate()
{
	std::string Auth;
	const bool	Success = m_CacheClient->AcquireAccessToken(Auth);
	return {.Success = Success};
}

CloudCacheResult
CloudCacheSession::GetDerivedData(std::string_view BucketId, std::string_view Key)
{
	std::string Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/c/ddc/" << m_CacheClient->DdcNamespace() << "/" << BucketId << "/" << Key << ".raw";

	cpr::Session& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(cpr::Header{{"Authorization", Auth}});

	cpr::Response Response = Session.Get();
	ZEN_DEBUG("GET {}", Response);

	const bool	   Success = Response.status_code == 200;
	const IoBuffer Buffer  = Success ? IoBufferBuilder::MakeCloneFromMemory(Response.text.data(), Response.text.size()) : IoBuffer();

	return {.Response = Buffer, .Bytes = Response.downloaded_bytes, .ElapsedSeconds = Response.elapsed, .Success = Success};
}

CloudCacheResult
CloudCacheSession::GetDerivedData(std::string_view BucketId, const IoHash& Key)
{
	return GetDerivedData(BucketId, Key.ToHexString());
}

CloudCacheResult
CloudCacheSession::GetRef(std::string_view BucketId, const IoHash& Key, ZenContentType RefType)
{
	const std::string ContentType = RefType == ZenContentType::kCbObject ? "application/x-ue-cb" : "application/octet-stream";
	std::string		  Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/refs/" << m_CacheClient->BlobStoreNamespace() << "/" << BucketId << "/"
		<< Key.ToHexString();

	cpr::Session& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(cpr::Header{{"Authorization", Auth}, {"Accept", ContentType}});

	cpr::Response Response = Session.Get();
	ZEN_DEBUG("GET {}", Response);

	const bool	   Success = Response.status_code == 200;
	const IoBuffer Buffer  = Success ? IoBufferBuilder::MakeCloneFromMemory(Response.text.data(), Response.text.size()) : IoBuffer();

	return {.Response = Buffer, .Bytes = Response.downloaded_bytes, .ElapsedSeconds = Response.elapsed, .Success = Success};
}

CloudCacheResult
CloudCacheSession::GetCompressedBlob(const IoHash& Key)
{
	std::string Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/compressed-blobs/" << m_CacheClient->BlobStoreNamespace() << "/" << Key.ToHexString();

	cpr::Session& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(cpr::Header{{"Authorization", Auth}, {"Accept", "application/x-ue-comp"}});

	cpr::Response Response = Session.Get();
	ZEN_DEBUG("GET {}", Response);

	const bool	   Success = Response.status_code == 200;
	const IoBuffer Buffer  = Success ? IoBufferBuilder::MakeCloneFromMemory(Response.text.data(), Response.text.size()) : IoBuffer();

	return {.Response = Buffer, .Bytes = Response.downloaded_bytes, .ElapsedSeconds = Response.elapsed, .Success = Success};
}

CloudCacheResult
CloudCacheSession::PutDerivedData(std::string_view BucketId, std::string_view Key, IoBuffer DerivedData)
{
	IoHash Hash = IoHash::HashBuffer(DerivedData.Data(), DerivedData.Size());

	std::string Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/c/ddc/" << m_CacheClient->DdcNamespace() << "/" << BucketId << "/" << Key;

	auto& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(
		cpr::Header{{"Authorization", Auth}, {"X-Jupiter-IoHash", Hash.ToHexString()}, {"Content-Type", "application/octet-stream"}});
	Session.SetBody(cpr::Body{(const char*)DerivedData.Data(), DerivedData.Size()});

	cpr::Response Response = Session.Put();
	ZEN_DEBUG("PUT {}", Response);

	return {.Bytes			= Response.uploaded_bytes,
			.ElapsedSeconds = Response.elapsed,
			.Success		= (Response.status_code == 200 || Response.status_code == 201)};
}

CloudCacheResult
CloudCacheSession::PutDerivedData(std::string_view BucketId, const IoHash& Key, IoBuffer DerivedData)
{
	return PutDerivedData(BucketId, Key.ToHexString(), DerivedData);
}

CloudCacheResult
CloudCacheSession::PutRef(std::string_view BucketId, const IoHash& Key, IoBuffer Ref, ZenContentType RefType)
{
	IoHash Hash = IoHash::HashBuffer(Ref.Data(), Ref.Size());

	const std::string ContentType = RefType == ZenContentType::kCbObject ? "application/x-ue-cb" : "application/octet-stream";
	std::string		  Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/refs/" << m_CacheClient->BlobStoreNamespace() << "/" << BucketId << "/"
		<< Key.ToHexString();

	cpr::Session& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(cpr::Header{{"Authorization", Auth}, {"X-Jupiter-IoHash", Hash.ToHexString()}, {"Content-Type", ContentType}});
	Session.SetBody(cpr::Body{(const char*)Ref.Data(), Ref.Size()});

	cpr::Response Response = Session.Put();
	ZEN_DEBUG("PUT {}", Response);

	return {.Bytes			= Response.uploaded_bytes,
			.ElapsedSeconds = Response.elapsed,
			.Success		= (Response.status_code == 200 || Response.status_code == 201)};
}

CloudCacheResult
CloudCacheSession::PutCompressedBlob(const IoHash& Key, IoBuffer Blob)
{
	std::string Auth;
	m_CacheClient->AcquireAccessToken(Auth);

	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl() << "/api/v1/compressed-blobs/" << m_CacheClient->BlobStoreNamespace() << "/" << Key.ToHexString();

	cpr::Session& Session = m_SessionState->Session;

	Session.SetOption(cpr::Url{Uri.c_str()});
	Session.SetOption(cpr::Header{{"Authorization", Auth}, {"Content-Type", "application/x-ue-comp"}});
	Session.SetBody(cpr::Body{(const char*)Blob.Data(), Blob.Size()});

	cpr::Response Response = Session.Put();
	ZEN_DEBUG("PUT {}", Response);

	return {.Bytes			= Response.uploaded_bytes,
			.ElapsedSeconds = Response.elapsed,
			.Success		= (Response.status_code == 200 || Response.status_code == 201)};
}

std::vector<IoHash>
CloudCacheSession::Filter(std::string_view BucketId, const std::vector<IoHash>& ChunkHashes)
{
	ExtendableStringBuilder<256> Uri;
	Uri << m_CacheClient->ServiceUrl();
	Uri << "/api/v1/s/" << m_CacheClient->DdcNamespace();

	ZEN_UNUSED(BucketId, ChunkHashes);

	return {};
}

//////////////////////////////////////////////////////////////////////////

std::string
CloudCacheAccessToken::GetAuthorizationHeaderValue()
{
	RwLock::SharedLockScope _(m_Lock);

	return "Bearer {}"_format(m_Token);
}

inline void
CloudCacheAccessToken::SetToken(std::string_view Token)
{
	RwLock::ExclusiveLockScope _(m_Lock);
	m_Token = Token;
	++m_Serial;
}

//////////////////////////////////////////////////////////////////////////
//
// ServiceUrl:		https://jupiter.devtools.epicgames.com
// DdcNamespace:		ue4.ddc
// OAuthClientId:	0oao91lrhqPiAlaGD0x7
// OAuthProvider:	https://epicgames.okta.com/oauth2/auso645ojjWVdRI3d0x7/v1/token
// OAuthSecret:		-GBWjjenhCgOwhxL5yBKNJECVIoDPH0MK4RDuN7d
//

CloudCacheClient::CloudCacheClient(const CloudCacheClientOptions& Options)
: m_Log(zen::logging::Get("jupiter"))
, m_ServiceUrl(Options.ServiceUrl)
, m_OAuthFullUri(Options.OAuthProvider)
, m_DdcNamespace(Options.DdcNamespace)
, m_BlobStoreNamespace(Options.BlobStoreNamespace)
, m_OAuthClientId(Options.OAuthClientId)
, m_OAuthSecret(Options.OAuthSecret)
{
	if (!Options.OAuthProvider.starts_with("http://"sv) && !Options.OAuthProvider.starts_with("https://"sv))
	{
		ZEN_WARN("bad provider specification: '{}' - must be fully qualified", Options.OAuthProvider);
		m_IsValid = false;

		return;
	}

	// Split into host and Uri substrings

	auto SchemePos = Options.OAuthProvider.find("://"sv);

	if (SchemePos == std::string::npos)
	{
		ZEN_WARN("Bad service URL passed to cloud cache client: '{}'", Options.ServiceUrl);
		m_IsValid = false;

		return;
	}

	auto DomainEnd = Options.OAuthProvider.find('/', /* also skip the :// */ SchemePos + 3);

	if (DomainEnd == std::string::npos)
	{
		ZEN_WARN("Bad service URL passed to cloud cache client: '{}' no path delimiter found", Options.ServiceUrl);
		m_IsValid = false;

		return;
	}

	m_OAuthDomain  = Options.OAuthProvider.substr(SchemePos + 3, DomainEnd - SchemePos - 3);  // epicgames.okta.com
	m_OAuthUriPath = Options.OAuthProvider.substr(DomainEnd + 1);							  // oauth2/..../v1/token
}

CloudCacheClient::~CloudCacheClient()
{
	RwLock::ExclusiveLockScope _(m_SessionStateLock);

	for (auto State : m_SessionStateCache)
	{
		delete State;
	}
}

bool
CloudCacheClient::AcquireAccessToken(std::string& AuthorizationHeaderValue)
{
	// TODO: check for expiration

	if (!m_IsValid)
	{
		ExtendableStringBuilder<128> OAuthFormData;
		OAuthFormData << "client_id=" << m_OAuthClientId
					  << "&scope=cache_access&grant_type=client_credentials&client_secret=" << m_OAuthSecret;

		const uint32_t CurrentSerial = m_AccessToken.GetSerial();

		static RwLock			   AuthMutex;
		RwLock::ExclusiveLockScope _(AuthMutex);

		// Protect against redundant authentication operations
		if (m_AccessToken.GetSerial() != CurrentSerial)
		{
			// TODO: this could verify that the token is actually valid and retry if not?

			return true;
		}

		std::string data{OAuthFormData};

		cpr::Response Response =
			cpr::Post(cpr::Url{m_OAuthFullUri}, cpr::Header{{"Content-Type", "application/x-www-form-urlencoded"}}, cpr::Body{data});

		std::string Body{std::move(Response.text)};

		// Parse JSON response

		std::string	 JsonError;
		json11::Json JsonResponse = json11::Json::parse(Body, /* out */ JsonError);
		if (!JsonError.empty())
		{
			ZEN_WARN("failed to parse OAuth response: '{}'", JsonError);

			return false;
		}

		std::string AccessToken		  = JsonResponse["access_token"].string_value();
		int			ExpiryTimeSeconds = JsonResponse["expires_in"].int_value();
		ZEN_UNUSED(ExpiryTimeSeconds);

		m_AccessToken.SetToken(AccessToken);

		m_IsValid = true;
	}

	AuthorizationHeaderValue = m_AccessToken.GetAuthorizationHeaderValue();

	return true;
}

detail::CloudCacheSessionState*
CloudCacheClient::AllocSessionState()
{
	detail::CloudCacheSessionState* State = nullptr;

	if (RwLock::ExclusiveLockScope _(m_SessionStateLock); !m_SessionStateCache.empty())
	{
		State = m_SessionStateCache.front();
		m_SessionStateCache.pop_front();
	}

	if (State == nullptr)
	{
		State = new detail::CloudCacheSessionState(*this);
	}

	State->Reset();

	return State;
}

void
CloudCacheClient::FreeSessionState(detail::CloudCacheSessionState* State)
{
	RwLock::ExclusiveLockScope _(m_SessionStateLock);
	m_SessionStateCache.push_front(State);
}

}  // namespace zen