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

#include "upstreamcache.h"
#include "jupiter.h"
#include "zen.h"

#include <zencore/blockingqueue.h>
#include <zencore/compactbinary.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/compactbinaryvalidation.h>
#include <zencore/fmtutils.h>
#include <zencore/stats.h>
#include <zencore/stream.h>
#include <zencore/timer.h>

#include <zenstore/cas.h>
#include <zenstore/cidstore.h>

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

#include <fmt/format.h>

#include <algorithm>
#include <atomic>
#include <thread>
#include <unordered_map>

namespace zen {

using namespace std::literals;

namespace detail {

	class JupiterUpstreamEndpoint final : public UpstreamEndpoint
	{
	public:
		JupiterUpstreamEndpoint(const CloudCacheClientOptions& Options)
		: m_Log(zen::logging::Get("upstream"))
		, m_UseLegacyDdc(Options.UseLegacyDdc)
		{
			using namespace fmt::literals;
			m_DisplayName = "Jupiter - '{}'"_format(Options.ServiceUrl);
			m_Client	  = new CloudCacheClient(Options);
		}

		virtual ~JupiterUpstreamEndpoint() = default;

		virtual UpstreamEndpointHealth Initialize() override { return CheckHealth(); }

		virtual bool IsHealthy() const override { return m_HealthOk.load(); }

		virtual UpstreamEndpointHealth CheckHealth() override
		{
			try
			{
				CloudCacheSession	   Session(m_Client);
				const CloudCacheResult Result = Session.Authenticate();

				m_HealthOk = Result.ErrorCode == 0;

				return {.Reason = std::move(Result.Reason), .Ok = Result.Success};
			}
			catch (std::exception& Err)
			{
				return {.Reason = Err.what(), .Ok = false};
			}
		}

		virtual std::string_view DisplayName() const override { return m_DisplayName; }

		virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
		{
			try
			{
				CloudCacheSession Session(m_Client);
				CloudCacheResult  Result;

				if (m_UseLegacyDdc && Type == ZenContentType::kBinary)
				{
					Result = Session.GetDerivedData(CacheKey.Bucket, CacheKey.Hash);
				}
				else
				{
					const ZenContentType AcceptType = Type == ZenContentType::kCbPackage ? ZenContentType::kCbObject : Type;
					Result							= Session.GetRef(CacheKey.Bucket, CacheKey.Hash, AcceptType);

					if (Result.Success && Type == ZenContentType::kCbPackage)
					{
						CbPackage Package;

						const CbValidateError ValidationResult = ValidateCompactBinary(Result.Response, CbValidateMode::All);
						if (Result.Success = ValidationResult == CbValidateError::None; Result.Success)
						{
							CbObject CacheRecord = LoadCompactBinaryObject(Result.Response);

							CacheRecord.IterateAttachments([&Session, &Result, &Package](CbFieldView AttachmentHash) {
								CloudCacheResult AttachmentResult = Session.GetCompressedBlob(AttachmentHash.AsHash());
								Result.Bytes += AttachmentResult.Bytes;
								Result.ElapsedSeconds += AttachmentResult.ElapsedSeconds;
								Result.ErrorCode = AttachmentResult.ErrorCode;

								if (CompressedBuffer Chunk = CompressedBuffer::FromCompressed(SharedBuffer(AttachmentResult.Response)))
								{
									Package.AddAttachment(CbAttachment(Chunk));
								}
								else
								{
									Result.Success = false;
								}
							});

							Package.SetObject(CacheRecord);
						}

						if (Result.Success)
						{
							BinaryWriter MemStream;
							Package.Save(MemStream);

							Result.Response = IoBuffer(IoBuffer::Clone, MemStream.Data(), MemStream.Size());
						}
					}
				}

				if (Result.ErrorCode == 0)
				{
					return {.Value			= Result.Response,
							.Bytes			= Result.Bytes,
							.ElapsedSeconds = Result.ElapsedSeconds,
							.Success		= Result.Success};
				}
				else
				{
					m_HealthOk = false;
					return {.Error{.ErrorCode = Result.ErrorCode, .Reason = std::move(Result.Reason)}};
				}
			}
			catch (std::exception& Err)
			{
				m_HealthOk = false;
				return {.Error{.ErrorCode = -1, .Reason = Err.what()}};
			}
		}

		virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
		{
			try
			{
				CloudCacheSession	   Session(m_Client);
				const CloudCacheResult Result = Session.GetCompressedBlob(PayloadKey.PayloadId);

				if (Result.ErrorCode == 0)
				{
					return {.Value			= Result.Response,
							.Bytes			= Result.Bytes,
							.ElapsedSeconds = Result.ElapsedSeconds,
							.Success		= Result.Success};
				}
				else
				{
					m_HealthOk = false;
					return {.Error{.ErrorCode = Result.ErrorCode, .Reason = std::move(Result.Reason)}};
				}
			}
			catch (std::exception& Err)
			{
				m_HealthOk = false;
				return {.Error{.ErrorCode = -1, .Reason = Err.what()}};
			}
		}

		virtual PutUpstreamCacheResult PutCacheRecord(const UpstreamCacheRecord& CacheRecord,
													  IoBuffer					 RecordValue,
													  std::span<IoBuffer const>	 Payloads) override
		{
			ZEN_ASSERT(CacheRecord.PayloadIds.size() == Payloads.size());
			const uint32_t MaxAttempts = 3;

			try
			{
				CloudCacheSession Session(m_Client);

				if (CacheRecord.Type == ZenContentType::kBinary)
				{
					CloudCacheResult Result;
					for (int32_t Attempt = 0; Attempt < MaxAttempts && !Result.Success; Attempt++)
					{
						if (m_UseLegacyDdc)
						{
							Result = Session.PutDerivedData(CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash, RecordValue);
						}
						else
						{
							Result = Session.PutRef(CacheRecord.CacheKey.Bucket,
													CacheRecord.CacheKey.Hash,
													RecordValue,
													ZenContentType::kBinary);
						}
					}

					return {.Bytes = Result.Bytes, .ElapsedSeconds = Result.ElapsedSeconds, .Success = Result.Success};
				}
				else
				{
					bool	Success				= false;
					int64_t TotalBytes			= 0ull;
					double	TotalElapsedSeconds = 0.0;

					for (size_t Idx = 0, Count = Payloads.size(); Idx < Count; Idx++)
					{
						Success = false;
						for (int32_t Attempt = 0; Attempt < MaxAttempts; Attempt++)
						{
							if (CloudCacheResult Result = Session.PutCompressedBlob(CacheRecord.PayloadIds[Idx], Payloads[Idx]);
								Result.Success)
							{
								TotalBytes += Result.Bytes;
								TotalElapsedSeconds += Result.ElapsedSeconds;
								Success = true;
								break;
							}
						}

						if (!Success)
						{
							return {.Reason			= "Failed to upload payload",
									.Bytes			= TotalBytes,
									.ElapsedSeconds = TotalElapsedSeconds,
									.Success		= false};
						}
					}

					Success = false;
					for (int32_t Attempt = 0; Attempt < MaxAttempts; Attempt++)
					{
						if (PutRefResult Result = Session.PutRef(CacheRecord.CacheKey.Bucket,
																 CacheRecord.CacheKey.Hash,
																 RecordValue,
																 ZenContentType::kCbObject);
							Result.Success)
						{
							TotalBytes += Result.Bytes;
							TotalElapsedSeconds += Result.ElapsedSeconds;
							Success = true;

							if (!Result.Needs.empty())
							{
								for (const IoHash& NeededHash : Result.Needs)
								{
									Success = false;

									if (auto It =
											std::find(std::begin(CacheRecord.PayloadIds), std::end(CacheRecord.PayloadIds), NeededHash);
										It != std::end(CacheRecord.PayloadIds))
									{
										const size_t Idx = It - std::begin(CacheRecord.PayloadIds);

										if (CloudCacheResult BlobResult =
												Session.PutCompressedBlob(CacheRecord.PayloadIds[Idx], Payloads[Idx]);
											BlobResult.Success)
										{
											TotalBytes += BlobResult.Bytes;
											TotalElapsedSeconds += BlobResult.ElapsedSeconds;
											Success = true;
										}
										else
										{
											ZEN_WARN("upload missing payload '{}/{}/{}' FAILED",
													 CacheRecord.CacheKey.Bucket,
													 CacheRecord.CacheKey.Hash,
													 NeededHash);
										}
									}
									else
									{
										ZEN_WARN("needed payload '{}/{}/{}' MISSING",
												 CacheRecord.CacheKey.Bucket,
												 CacheRecord.CacheKey.Hash,
												 NeededHash);
									}
								}

								const IoHash RefHash = IoHash::HashBuffer(RecordValue);

								if (FinalizeRefResult FinalizeResult =
										Session.FinalizeRef(CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash, RefHash);
									FinalizeResult.Success)
								{
									TotalBytes += FinalizeResult.Bytes;
									TotalElapsedSeconds += FinalizeResult.ElapsedSeconds;
									Success = true;

									for (const IoHash& MissingHash : FinalizeResult.Needs)
									{
										ZEN_WARN("finalize '{}/{}' FAILED, missing '{}'",
												 CacheRecord.CacheKey.Bucket,
												 CacheRecord.CacheKey.Hash,
												 MissingHash);
									}
								}
								else
								{
									ZEN_WARN("finalize '{}/{}' FAILED", CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash);
									Success = false;
								}
							}

							if (Success)
							{
								break;
							}
						}
					}

					return {.Bytes = TotalBytes, .ElapsedSeconds = TotalElapsedSeconds, .Success = Success};
				}
			}
			catch (std::exception& Err)
			{
				return {.Reason = std::string(Err.what()), .Success = false};
			}
		}

		virtual UpstreamEndpointStats& Stats() override { return m_Stats; }

	private:
		spdlog::logger& Log() { return m_Log; }

		spdlog::logger&			 m_Log;
		bool					 m_UseLegacyDdc;
		std::string				 m_DisplayName;
		RefPtr<CloudCacheClient> m_Client;
		UpstreamEndpointStats	 m_Stats;
		std::atomic_bool		 m_HealthOk{false};
	};

	class ZenUpstreamEndpoint final : public UpstreamEndpoint
	{
		struct ZenEndpoint
		{
			std::string Url;
			std::string Reason;
			double		Latency{};
			bool		Ok = false;

			bool operator<(const ZenEndpoint& RHS) const { return Ok && RHS.Ok ? Latency < RHS.Latency : Ok; }
		};

	public:
		ZenUpstreamEndpoint(std::span<std::string const> Urls) : m_Log(zen::logging::Get("upstream")), m_DisplayName("ZEN")
		{
			for (const auto& Url : Urls)
			{
				m_Endpoints.push_back({.Url = Url});
			}
		}

		~ZenUpstreamEndpoint() = default;

		virtual UpstreamEndpointHealth Initialize() override
		{
			using namespace fmt::literals;

			const ZenEndpoint& Ep = GetEndpoint();
			if (Ep.Ok)
			{
				m_ServiceUrl  = Ep.Url;
				m_DisplayName = "ZEN - {}"_format(m_ServiceUrl);
				m_Client	  = new ZenStructuredCacheClient(m_ServiceUrl);

				m_HealthOk = true;
				return {.Ok = true};
			}

			m_HealthOk = false;
			return {.Reason = Ep.Reason};
		}

		virtual bool IsHealthy() const override { return m_HealthOk; }

		virtual UpstreamEndpointHealth CheckHealth() override
		{
			using namespace fmt::literals;

			try
			{
				if (m_Client.IsNull())
				{
					const ZenEndpoint& Ep = GetEndpoint();
					if (Ep.Ok)
					{
						m_ServiceUrl  = Ep.Url;
						m_DisplayName = "ZEN - {}"_format(m_ServiceUrl);
						m_Client	  = new ZenStructuredCacheClient(m_ServiceUrl);

						m_HealthOk = true;
						return {.Ok = true};
					}

					return {.Reason = Ep.Reason};
				}

				ZenStructuredCacheSession Session(*m_Client);
				ZenCacheResult			  Result;

				for (int32_t Attempt = 0, MaxAttempts = 2; Attempt < MaxAttempts && !Result.Success; ++Attempt)
				{
					Result = Session.CheckHealth();
				}

				m_HealthOk = Result.ErrorCode == 0;

				return {.Reason = std::move(Result.Reason), .Ok = m_HealthOk};
			}
			catch (std::exception& Err)
			{
				return {.Reason = Err.what(), .Ok = false};
			}
		}

		virtual std::string_view DisplayName() const override { return m_DisplayName; }

		virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
		{
			try
			{
				ZenStructuredCacheSession Session(*m_Client);
				const ZenCacheResult	  Result = Session.GetCacheRecord(CacheKey.Bucket, CacheKey.Hash, Type);

				if (Result.ErrorCode == 0)
				{
					return {.Value			= Result.Response,
							.Bytes			= Result.Bytes,
							.ElapsedSeconds = Result.ElapsedSeconds,
							.Success		= Result.Success};
				}
				else
				{
					m_HealthOk = false;
					return {.Error{.ErrorCode = Result.ErrorCode, .Reason = std::move(Result.Reason)}};
				}
			}
			catch (std::exception& Err)
			{
				m_HealthOk = false;
				return {.Error{.ErrorCode = -1, .Reason = Err.what()}};
			}
		}

		virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
		{
			try
			{
				ZenStructuredCacheSession Session(*m_Client);
				const ZenCacheResult	  Result =
					Session.GetCachePayload(PayloadKey.CacheKey.Bucket, PayloadKey.CacheKey.Hash, PayloadKey.PayloadId);

				if (Result.ErrorCode == 0)
				{
					return {.Value			= Result.Response,
							.Bytes			= Result.Bytes,
							.ElapsedSeconds = Result.ElapsedSeconds,
							.Success		= Result.Success};
				}
				else
				{
					m_HealthOk = false;
					return {.Error{.ErrorCode = Result.ErrorCode, .Reason = std::move(Result.Reason)}};
				}
			}
			catch (std::exception& Err)
			{
				m_HealthOk = false;
				return {.Error{.ErrorCode = -1, .Reason = Err.what()}};
			}
		}

		virtual PutUpstreamCacheResult PutCacheRecord(const UpstreamCacheRecord& CacheRecord,
													  IoBuffer					 RecordValue,
													  std::span<IoBuffer const>	 Payloads) override
		{
			ZEN_ASSERT(CacheRecord.PayloadIds.size() == Payloads.size());
			const uint32_t MaxAttempts = 3;

			try
			{
				ZenStructuredCacheSession Session(*m_Client);
				ZenCacheResult			  Result;
				int64_t					  TotalBytes		  = 0ull;
				double					  TotalElapsedSeconds = 0.0;

				if (CacheRecord.Type == ZenContentType::kCbPackage)
				{
					CbPackage Package;
					Package.SetObject(CbObject(SharedBuffer(RecordValue)));

					for (const IoBuffer& Payload : Payloads)
					{
						if (CompressedBuffer AttachmentBuffer = CompressedBuffer::FromCompressed(SharedBuffer(Payload)))
						{
							Package.AddAttachment(CbAttachment(AttachmentBuffer));
						}
						else
						{
							return {.Reason = std::string("invalid payload buffer"), .Success = false};
						}
					}

					BinaryWriter MemStream;
					Package.Save(MemStream);
					IoBuffer PackagePayload(IoBuffer::Wrap, MemStream.Data(), MemStream.Size());

					for (int32_t Attempt = 0; Attempt < MaxAttempts && !Result.Success; Attempt++)
					{
						Result = Session.PutCacheRecord(CacheRecord.CacheKey.Bucket,
														CacheRecord.CacheKey.Hash,
														PackagePayload,
														CacheRecord.Type);

						m_HealthOk = Result.ErrorCode == 0;
					}

					TotalBytes			= Result.Bytes;
					TotalElapsedSeconds = Result.ElapsedSeconds;
				}
				else
				{
					for (size_t Idx = 0, Count = Payloads.size(); Idx < Count; Idx++)
					{
						Result.Success = false;
						for (int32_t Attempt = 0; Attempt < MaxAttempts && !Result.Success; Attempt++)
						{
							Result = Session.PutCachePayload(CacheRecord.CacheKey.Bucket,
															 CacheRecord.CacheKey.Hash,
															 CacheRecord.PayloadIds[Idx],
															 Payloads[Idx]);

							m_HealthOk = Result.ErrorCode == 0;
						}

						TotalBytes += Result.Bytes;
						TotalElapsedSeconds += Result.ElapsedSeconds;

						if (!Result.Success)
						{
							return {.Reason			= "Failed to upload payload",
									.Bytes			= TotalBytes,
									.ElapsedSeconds = TotalElapsedSeconds,
									.Success		= false};
						}
					}

					Result.Success = false;
					for (int32_t Attempt = 0; Attempt < MaxAttempts && !Result.Success; Attempt++)
					{
						Result =
							Session.PutCacheRecord(CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash, RecordValue, CacheRecord.Type);

						m_HealthOk = Result.ErrorCode == 0;
					}

					TotalBytes += Result.Bytes;
					TotalElapsedSeconds += Result.ElapsedSeconds;
				}

				return {.Bytes = TotalBytes, .ElapsedSeconds = TotalElapsedSeconds, .Success = Result.Success};
			}
			catch (std::exception& e)
			{
				m_HealthOk = false;
				return {.Reason = std::string(e.what()), .Success = false};
			}
		}

		virtual UpstreamEndpointStats& Stats() override { return m_Stats; }

	private:
		const ZenEndpoint& GetEndpoint()
		{
			for (ZenEndpoint& Ep : m_Endpoints)
			{
				ZenStructuredCacheClient  Client(Ep.Url);
				ZenStructuredCacheSession Session(Client);
				const int32_t			  SampleCount = 2;

				Ep.Ok	   = false;
				Ep.Latency = {};

				for (int32_t Sample = 0; Sample < SampleCount; ++Sample)
				{
					ZenCacheResult Result = Session.CheckHealth();
					Ep.Ok				  = Result.Success;
					Ep.Reason			  = std::move(Result.Reason);
					Ep.Latency += Result.ElapsedSeconds;
				}
				Ep.Latency /= double(SampleCount);
			}

			std::sort(std::begin(m_Endpoints), std::end(m_Endpoints));

			for (const auto& Ep : m_Endpoints)
			{
				ZEN_INFO("ping ZEN endpoint '{}' latency '{:.3}s' {}", Ep.Url, Ep.Latency, Ep.Ok ? "OK" : Ep.Reason);
			}

			return m_Endpoints.front();
		}

		spdlog::logger& Log() { return m_Log; }

		spdlog::logger&					 m_Log;
		std::string						 m_ServiceUrl;
		std::vector<ZenEndpoint>		 m_Endpoints;
		std::string						 m_DisplayName;
		RefPtr<ZenStructuredCacheClient> m_Client;
		UpstreamEndpointStats			 m_Stats;
		std::atomic_bool				 m_HealthOk{false};
	};

}  // namespace detail

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

struct UpstreamStats
{
	static constexpr uint64_t MaxSampleCount = 1000ull;

	UpstreamStats(bool Enabled) : m_Enabled(Enabled) {}

	void Add(spdlog::logger&									   Logger,
			 UpstreamEndpoint&									   Endpoint,
			 const GetUpstreamCacheResult&						   Result,
			 const std::vector<std::unique_ptr<UpstreamEndpoint>>& Endpoints)
	{
		UpstreamEndpointStats& Stats = Endpoint.Stats();

		if (Result.Error)
		{
			Stats.ErrorCount++;
		}
		else if (Result.Success)
		{
			Stats.HitCount++;
			Stats.DownBytes.fetch_add(double(Result.Bytes) / 1024.0 / 1024.0);
			Stats.SecondsDown.fetch_add(Result.ElapsedSeconds);
		}
		else
		{
			Stats.MissCount++;
		}

		if (m_Enabled && m_SampleCount++ % MaxSampleCount)
		{
			Dump(Logger, Endpoints);
		}
	}

	void Add(spdlog::logger&									   Logger,
			 UpstreamEndpoint&									   Endpoint,
			 const PutUpstreamCacheResult&						   Result,
			 const std::vector<std::unique_ptr<UpstreamEndpoint>>& Endpoints)
	{
		UpstreamEndpointStats& Stats = Endpoint.Stats();
		if (Result.Success)
		{
			Stats.UpCount++;
			Stats.UpBytes.fetch_add(double(Result.Bytes) / 1024.0 / 1024.0);
			Stats.SecondsUp.fetch_add(Result.ElapsedSeconds);
		}
		else
		{
			Stats.ErrorCount++;
		}

		if (m_Enabled && m_SampleCount++ % MaxSampleCount)
		{
			Dump(Logger, Endpoints);
		}
	}

	void Dump(spdlog::logger& Logger, const std::vector<std::unique_ptr<UpstreamEndpoint>>& Endpoints)
	{
		for (auto& Ep : Endpoints)
		{
			// These stats will not be totally correct as the numbers are not captured atomically

			UpstreamEndpointStats& Stats	   = Ep->Stats();
			const uint64_t		   HitCount	   = Stats.HitCount;
			const uint64_t		   MissCount   = Stats.MissCount;
			const double		   DownBytes   = Stats.DownBytes;
			const double		   SecondsDown = Stats.SecondsDown;
			const double		   UpBytes	   = Stats.UpBytes;
			const double		   SecondsUp   = Stats.SecondsUp;

			const double   UpSpeed	  = UpBytes > 0 ? UpBytes / SecondsUp : 0.0;
			const double   DownSpeed  = DownBytes > 0 ? DownBytes / SecondsDown : 0.0;
			const uint64_t TotalCount = HitCount + MissCount;
			const double   HitRate	  = TotalCount > 0 ? (double(HitCount) / double(TotalCount)) : 0.0;

			Logger.debug("STATS - '{}', Hit rate: {:.2f}%, DOWN: '{:.2f} MiB {:.2f} MiB/s', UP: '{:.2f} MiB {:.2f} MiB/s'",
						 Ep->DisplayName(),
						 HitRate,
						 DownBytes,
						 DownSpeed,
						 UpBytes,
						 UpSpeed);
		}
	}

	bool				 m_Enabled;
	std::atomic_uint64_t m_SampleCount = {};
};

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

class DefaultUpstreamCache final : public UpstreamCache
{
public:
	DefaultUpstreamCache(const UpstreamCacheOptions& Options, ZenCacheStore& CacheStore, CidStore& CidStore)
	: m_Log(logging::Get("upstream"))
	, m_Options(Options)
	, m_CacheStore(CacheStore)
	, m_CidStore(CidStore)
	, m_Stats(Options.StatsEnabled)
	{
	}

	virtual ~DefaultUpstreamCache() { Shutdown(); }

	virtual bool Initialize() override
	{
		for (auto& Endpoint : m_Endpoints)
		{
			const UpstreamEndpointHealth Health = Endpoint->Initialize();
			if (Health.Ok)
			{
				ZEN_INFO("initialize endpoint '{}' OK", Endpoint->DisplayName());
			}
			else
			{
				ZEN_WARN("initialize endpoint '{}' FAILED, reason '{}'", Endpoint->DisplayName(), Health.Reason);
			}
		}

		m_RunState.IsRunning = !m_Endpoints.empty();

		if (m_RunState.IsRunning)
		{
			for (uint32_t Idx = 0; Idx < m_Options.ThreadCount; Idx++)
			{
				m_UpstreamThreads.emplace_back(&DefaultUpstreamCache::ProcessUpstreamQueue, this);
			}

			m_EndpointMonitorThread = std::thread(&DefaultUpstreamCache::MonitorEndpoints, this);
		}

		return m_RunState.IsRunning;
	}

	virtual void RegisterEndpoint(std::unique_ptr<UpstreamEndpoint> Endpoint) override { m_Endpoints.emplace_back(std::move(Endpoint)); }

	virtual GetUpstreamCacheResult GetCacheRecord(UpstreamCacheKey CacheKey, ZenContentType Type) override
	{
		if (m_Options.ReadUpstream)
		{
			for (auto& Endpoint : m_Endpoints)
			{
				if (Endpoint->IsHealthy())
				{
					const GetUpstreamCacheResult Result = Endpoint->GetCacheRecord(CacheKey, Type);
					m_Stats.Add(m_Log, *Endpoint, Result, m_Endpoints);

					if (Result.Success)
					{
						return Result;
					}
				}
			}
		}

		return {};
	}

	virtual GetUpstreamCacheResult GetCachePayload(UpstreamPayloadKey PayloadKey) override
	{
		if (m_Options.ReadUpstream)
		{
			for (auto& Endpoint : m_Endpoints)
			{
				if (Endpoint->IsHealthy())
				{
					const GetUpstreamCacheResult Result = Endpoint->GetCachePayload(PayloadKey);
					m_Stats.Add(m_Log, *Endpoint, Result, m_Endpoints);

					if (Result.Success)
					{
						return Result;
					}
				}
			}
		}

		return {};
	}

	virtual EnqueueResult EnqueueUpstream(UpstreamCacheRecord CacheRecord) override
	{
		if (m_RunState.IsRunning && m_Options.WriteUpstream)
		{
			if (!m_UpstreamThreads.empty())
			{
				m_UpstreamQueue.Enqueue(std::move(CacheRecord));
			}
			else
			{
				ProcessCacheRecord(std::move(CacheRecord));
			}

			return {.Success = true};
		}

		return {};
	}

	virtual void GetStatus(CbObjectWriter& Status) override
	{
		Status << "reading" << m_Options.ReadUpstream;
		Status << "writing" << m_Options.WriteUpstream;
		Status << "worker_threads" << m_Options.ThreadCount;
		Status << "queue_count" << m_UpstreamQueue.Size();

		Status.BeginArray("endpoints");
		for (const auto& Ep : m_Endpoints)
		{
			Status.BeginObject();
			Status << "name" << Ep->DisplayName();
			Status << "health" << (Ep->IsHealthy() ? "ok"sv : "inactive"sv);

			UpstreamEndpointStats& Stats	  = Ep->Stats();
			const uint64_t		   HitCount	  = Stats.HitCount;
			const uint64_t		   MissCount  = Stats.MissCount;
			const uint64_t		   TotalCount = HitCount + MissCount;
			const double		   HitRate	  = TotalCount > 0 ? (double(HitCount) / double(TotalCount)) : 0.0;

			Status << "hit_ratio" << HitRate;
			Status << "downloaded_mb" << Stats.DownBytes;
			Status << "uploaded_mb" << Stats.UpBytes;
			Status << "error_count" << Stats.ErrorCount;

			Status.EndObject();
		}
		Status.EndArray();
	}

private:
	void ProcessCacheRecord(UpstreamCacheRecord CacheRecord)
	{
		ZenCacheValue		  CacheValue;
		std::vector<IoBuffer> Payloads;

		if (!m_CacheStore.Get(CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash, CacheValue))
		{
			ZEN_WARN("process upstream FAILED, '{}/{}', cache record doesn't exist",
					 CacheRecord.CacheKey.Bucket,
					 CacheRecord.CacheKey.Hash);
			return;
		}

		for (const IoHash& PayloadId : CacheRecord.PayloadIds)
		{
			if (IoBuffer Payload = m_CidStore.FindChunkByCid(PayloadId))
			{
				Payloads.push_back(Payload);
			}
			else
			{
				ZEN_WARN("process upstream FAILED, '{}/{}/{}', payload doesn't exist in CAS",
						 CacheRecord.CacheKey.Bucket,
						 CacheRecord.CacheKey.Hash,
						 PayloadId);
				return;
			}
		}

		for (auto& Endpoint : m_Endpoints)
		{
			if (Endpoint->IsHealthy())
			{
				const PutUpstreamCacheResult Result = Endpoint->PutCacheRecord(CacheRecord, CacheValue.Value, std::span(Payloads));
				m_Stats.Add(m_Log, *Endpoint, Result, m_Endpoints);
			}
		}
	}

	void ProcessUpstreamQueue()
	{
		for (;;)
		{
			UpstreamCacheRecord CacheRecord;
			if (m_UpstreamQueue.WaitAndDequeue(CacheRecord))
			{
				try
				{
					ProcessCacheRecord(std::move(CacheRecord));
				}
				catch (std::exception& e)
				{
					ZEN_WARN("process upstream ({}/{}) FAILED '{}'", CacheRecord.CacheKey.Bucket, CacheRecord.CacheKey.Hash, e.what());
				}
			}

			if (!m_RunState.IsRunning)
			{
				break;
			}
		}
	}

	void MonitorEndpoints()
	{
		for (;;)
		{
			{
				std::unique_lock lk(m_RunState.Mutex);
				if (m_RunState.ExitSignal.wait_for(lk, m_Options.HealthCheckInterval, [this]() { return !m_RunState.IsRunning.load(); }))
				{
					break;
				}
			}

			for (auto& Endpoint : m_Endpoints)
			{
				if (!Endpoint->IsHealthy())
				{
					if (const UpstreamEndpointHealth Health = Endpoint->CheckHealth(); Health.Ok)
					{
						ZEN_INFO("health check endpoint '{}' OK", Endpoint->DisplayName(), Health.Reason);
					}
					else
					{
						ZEN_WARN("health check endpoint '{}' FAILED, reason '{}'", Endpoint->DisplayName(), Health.Reason);
					}
				}
			}
		}
	}

	void Shutdown()
	{
		if (m_RunState.Stop())
		{
			m_UpstreamQueue.CompleteAdding();
			for (std::thread& Thread : m_UpstreamThreads)
			{
				Thread.join();
			}

			m_EndpointMonitorThread.join();
			m_UpstreamThreads.clear();
			m_Endpoints.clear();
		}
	}

	spdlog::logger& Log() { return m_Log; }

	using UpstreamQueue = BlockingQueue<UpstreamCacheRecord>;

	struct RunState
	{
		std::mutex				Mutex;
		std::condition_variable ExitSignal;
		std::atomic_bool		IsRunning{false};

		bool Stop()
		{
			bool Stopped = false;
			{
				std::lock_guard _(Mutex);
				Stopped = IsRunning.exchange(false);
			}
			if (Stopped)
			{
				ExitSignal.notify_all();
			}
			return Stopped;
		}
	};

	spdlog::logger&								   m_Log;
	UpstreamCacheOptions						   m_Options;
	ZenCacheStore&								   m_CacheStore;
	CidStore&									   m_CidStore;
	UpstreamQueue								   m_UpstreamQueue;
	UpstreamStats								   m_Stats;
	std::vector<std::unique_ptr<UpstreamEndpoint>> m_Endpoints;
	std::vector<std::thread>					   m_UpstreamThreads;
	std::thread									   m_EndpointMonitorThread;
	RunState									   m_RunState;
};

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

std::unique_ptr<UpstreamCache>
MakeUpstreamCache(const UpstreamCacheOptions& Options, ZenCacheStore& CacheStore, CidStore& CidStore)
{
	return std::make_unique<DefaultUpstreamCache>(Options, CacheStore, CidStore);
}

std::unique_ptr<UpstreamEndpoint>
MakeJupiterUpstreamEndpoint(const CloudCacheClientOptions& Options)
{
	return std::make_unique<detail::JupiterUpstreamEndpoint>(Options);
}

std::unique_ptr<UpstreamEndpoint>
MakeZenUpstreamEndpoint(std::span<std::string const> Urls)
{
	return std::make_unique<detail::ZenUpstreamEndpoint>(Urls);
}

}  // namespace zen