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

#pragma once

#include <zencore/compactbinary.h>
#include <zencore/compress.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/stats.h>
#include <zencore/zencore.h>
#include <zenstore/cache/upstreamcacheclient.h>
#include <zenutil/cache/cache.h>

#include <atomic>
#include <chrono>
#include <functional>
#include <memory>
#include <vector>

namespace zen {

class CbObjectView;
class AuthMgr;
class CbObjectView;
class CbPackage;
class CbObjectWriter;
class ZenCacheStore;
struct JupiterClientOptions;
class JupiterAccessTokenProvider;
struct ZenStructuredCacheClientOptions;

struct UpstreamEndpointStats
{
	metrics::OperationTiming CacheGetRequestTiming;
	metrics::OperationTiming CachePutRequestTiming;
	metrics::Counter		 CacheGetTotalBytes;
	metrics::Counter		 CachePutTotalBytes;
	metrics::Counter		 CacheGetCount;
	metrics::Counter		 CacheHitCount;
	metrics::Counter		 CacheErrorCount;
};

enum class UpstreamEndpointState : uint32_t
{
	kDisabled,
	kUnauthorized,
	kError,
	kOk
};

inline std::string_view
ToString(UpstreamEndpointState State)
{
	using namespace std::literals;

	switch (State)
	{
		case UpstreamEndpointState::kDisabled:
			return "Disabled"sv;
		case UpstreamEndpointState::kUnauthorized:
			return "Unauthorized"sv;
		case UpstreamEndpointState::kError:
			return "Error"sv;
		case UpstreamEndpointState::kOk:
			return "Ok"sv;
		default:
			return "Unknown"sv;
	}
}

struct UpstreamAuthConfig
{
	std::string_view OAuthUrl;
	std::string_view OAuthClientId;
	std::string_view OAuthClientSecret;
	std::string_view OpenIdProvider;
	std::string_view AccessToken;
};

struct UpstreamEndpointStatus
{
	std::string			  Reason;
	UpstreamEndpointState State;
};

struct GetUpstreamCacheSingleResult
{
	GetUpstreamCacheResult		Status;
	IoBuffer					Value;
	const UpstreamEndpointInfo* Source = nullptr;
};

/**
 * The upstream endpoint is responsible for handling upload/downloading of cache records.
 */
class UpstreamEndpoint
{
public:
	virtual ~UpstreamEndpoint() = default;

	virtual UpstreamEndpointStatus Initialize() = 0;

	virtual const UpstreamEndpointInfo& GetEndpointInfo() const = 0;

	virtual UpstreamEndpointState  GetState()  = 0;
	virtual UpstreamEndpointStatus GetStatus() = 0;

	virtual GetUpstreamCacheSingleResult GetCacheRecord(std::string_view Namespace, const CacheKey& CacheKey, ZenContentType Type) = 0;
	virtual GetUpstreamCacheResult		 GetCacheRecords(std::string_view			 Namespace,
														 std::span<CacheKeyRequest*> Requests,
														 OnCacheRecordGetComplete&&	 OnComplete)									   = 0;

	virtual GetUpstreamCacheResult GetCacheValues(std::string_view				Namespace,
												  std::span<CacheValueRequest*> CacheValueRequests,
												  OnCacheValueGetComplete&&		OnComplete) = 0;

	virtual GetUpstreamCacheSingleResult GetCacheChunk(std::string_view Namespace, const CacheKey& CacheKey, const IoHash& PayloadId) = 0;
	virtual GetUpstreamCacheResult		 GetCacheChunks(std::string_view			  Namespace,
														std::span<CacheChunkRequest*> CacheChunkRequests,
														OnCacheChunksGetComplete&&	  OnComplete)										  = 0;

	virtual PutUpstreamCacheResult PutCacheRecord(const UpstreamCacheRecord& CacheRecord,
												  IoBuffer					 RecordValue,
												  std::span<IoBuffer const>	 Payloads) = 0;

	virtual UpstreamEndpointStats& Stats() = 0;

	static std::unique_ptr<UpstreamEndpoint> CreateZenEndpoint(const ZenStructuredCacheClientOptions& Options);

	static std::unique_ptr<UpstreamEndpoint> CreateJupiterEndpoint(const JupiterClientOptions& Options,
																   const UpstreamAuthConfig&   AuthConfig,
																   AuthMgr&					   Mgr);
};

/**
 * Manages one or more upstream cache endpoints.
 */

class UpstreamCache : public UpstreamCacheClient
{
public:
	virtual void Initialize() = 0;

	virtual void RegisterEndpoint(std::unique_ptr<UpstreamEndpoint> Endpoint)  = 0;
	virtual void IterateEndpoints(std::function<bool(UpstreamEndpoint&)>&& Fn) = 0;

	virtual GetUpstreamCacheSingleResult GetCacheRecord(std::string_view Namespace, const CacheKey& CacheKey, ZenContentType Type) = 0;

	virtual GetUpstreamCacheSingleResult GetCacheChunk(std::string_view Namespace,
													   const CacheKey&	CacheKey,
													   const IoHash&	ValueContentId) = 0;

	virtual void GetStatus(CbObjectWriter& CbO) = 0;
};

struct UpstreamCacheOptions
{
	std::chrono::seconds HealthCheckInterval{5};
	uint32_t			 ThreadCount   = 4;
	bool				 ReadUpstream  = true;
	bool				 WriteUpstream = true;
};

std::unique_ptr<UpstreamCache> CreateUpstreamCache(const UpstreamCacheOptions& Options, ZenCacheStore& CacheStore);

}  // namespace zen