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

#pragma once

#include "cachedisklayer.h"
#include "stats/statsreporter.h"

#include <zencore/compactbinary.h>
#include <zencore/iohash.h>
#include <zencore/stats.h>
#include <zenstore/gc.h>
#include <zenutil/cache/cache.h>

#include <atomic>
#include <compare>
#include <filesystem>
#include <string_view>
#include <unordered_map>

namespace zen {

class StatsDaemonClient;

/******************************************************************************

  /$$$$$$$$                        /$$$$$$                   /$$
 |_____ $$                        /$$__  $$                 | $$
	  /$$/  /$$$$$$ /$$$$$$$     | $$  \__/ /$$$$$$  /$$$$$$| $$$$$$$  /$$$$$$
	 /$$/  /$$__  $| $$__  $$    | $$      |____  $$/$$_____| $$__  $$/$$__  $$
	/$$/  | $$$$$$$| $$  \ $$    | $$       /$$$$$$| $$     | $$  \ $| $$$$$$$$
   /$$/   | $$_____| $$  | $$    | $$    $$/$$__  $| $$     | $$  | $| $$_____/
  /$$$$$$$|  $$$$$$| $$  | $$    |  $$$$$$|  $$$$$$|  $$$$$$| $$  | $|  $$$$$$$
 |________/\_______|__/  |__/     \______/ \_______/\_______|__/  |__/\_______/

  Cache store for UE5. Restricts keys to "{bucket}/{hash}" pairs where the hash
  is 40 (hex) chars in size. Values may be opaque blobs or structured objects
  which can in turn contain references to other objects (or blobs). Buckets are
  organized in namespaces to enable project isolation.

******************************************************************************/

class WorkerThreadPool;
class DiskWriteBlocker;
class JobQueue;

/* Z$ namespace

   A namespace scopes a set of buckets, and would typically be used to isolate
   projects from each other.

 */
class ZenCacheNamespace final : public GcStorage, public GcContributor
{
public:
	struct Configuration
	{
		ZenCacheDiskLayer::Configuration DiskLayerConfig;
	};
	struct BucketInfo
	{
		ZenCacheDiskLayer::BucketInfo DiskLayerInfo;
	};
	struct Info
	{
		std::filesystem::path	 RootDir;
		Configuration			 Config;
		std::vector<std::string> BucketNames;
		ZenCacheDiskLayer::Info	 DiskLayerInfo;
	};

	struct NamespaceStats
	{
		uint64_t					  HitCount;
		uint64_t					  MissCount;
		uint64_t					  WriteCount;
		metrics::RequestStatsSnapshot PutOps;
		metrics::RequestStatsSnapshot GetOps;
		ZenCacheDiskLayer::DiskStats  DiskStats;
	};

	ZenCacheNamespace(GcManager& Gc, JobQueue& JobQueue, const std::filesystem::path& RootDir, const Configuration& Config);
	~ZenCacheNamespace();

	bool Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
	void Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value, std::span<IoHash> References);

	bool DropBucket(std::string_view Bucket);
	void EnumerateBucketContents(std::string_view																		 Bucket,
								 std::function<void(const IoHash& Key, const CacheValueDetails::ValueDetails& Details)>& Fn) const;

	bool Drop();
	void Flush();

	// GcContributor
	virtual void GatherReferences(GcContext& GcCtx) override;

	// GcStorage
	virtual void		  ScrubStorage(ScrubContext& ScrubCtx) override;
	virtual void		  CollectGarbage(GcContext& GcCtx) override;
	virtual GcStorageSize StorageSize() const override;

	Configuration			  GetConfig() const { return m_Configuration; }
	Info					  GetInfo() const;
	std::optional<BucketInfo> GetBucketInfo(std::string_view Bucket) const;
	NamespaceStats			  Stats();

	CacheValueDetails::NamespaceDetails GetValueDetails(const std::string_view BucketFilter, const std::string_view ValueFilter) const;

#if ZEN_WITH_TESTS
	void SetAccessTime(std::string_view Bucket, const IoHash& HashKey, GcClock::TimePoint Time);
#endif	// ZEN_WITH_TESTS

private:
	GcManager&			  m_Gc;
	JobQueue&			  m_JobQueue;
	std::filesystem::path m_RootDir;
	Configuration		  m_Configuration;
	ZenCacheDiskLayer	  m_DiskLayer;
	std::atomic<uint64_t> m_HitCount{};
	std::atomic<uint64_t> m_MissCount{};
	std::atomic<uint64_t> m_WriteCount{};
	metrics::RequestStats m_PutOps;
	metrics::RequestStats m_GetOps;
	uint64_t			  m_LastScrubTime = 0;

	ZenCacheNamespace(const ZenCacheNamespace&) = delete;
	ZenCacheNamespace& operator=(const ZenCacheNamespace&) = delete;
};

/** Cache store interface

	This manages a set of namespaces used for derived data caching purposes.

  */

class ZenCacheStore final : public RefCounted, public StatsProvider
{
public:
	static constexpr std::string_view DefaultNamespace =
		"!default!";  // This is intentionally not a valid namespace name and will only be used for mapping when no namespace is given
	static constexpr std::string_view NamespaceDiskPrefix = "ns_";

	struct Configuration
	{
		ZenCacheNamespace::Configuration NamespaceConfig;
		bool							 AllowAutomaticCreationOfNamespaces = false;
		struct LogConfig
		{
			bool EnableWriteLog	 = true;
			bool EnableAccessLog = true;
		} Logging;
	};

	struct Info
	{
		std::filesystem::path	 BasePath;
		Configuration			 Config;
		std::vector<std::string> NamespaceNames;
		uint64_t				 DiskEntryCount = 0;
		GcStorageSize			 StorageSize;
	};

	struct NamedNamespaceStats
	{
		std::string						  NamespaceName;
		ZenCacheNamespace::NamespaceStats Stats;
	};

	struct CacheStoreStats
	{
		uint64_t						 HitCount			= 0;
		uint64_t						 MissCount			= 0;
		uint64_t						 WriteCount			= 0;
		uint64_t						 RejectedWriteCount = 0;
		uint64_t						 RejectedReadCount	= 0;
		metrics::RequestStatsSnapshot	 PutOps;
		metrics::RequestStatsSnapshot	 GetOps;
		std::vector<NamedNamespaceStats> NamespaceStats;
	};

	ZenCacheStore(GcManager&				   Gc,
				  JobQueue&					   JobQueue,
				  const std::filesystem::path& BasePath,
				  const Configuration&		   Configuration,
				  const DiskWriteBlocker*	   InDiskWriteBlocker);
	~ZenCacheStore();

	bool Get(const CacheRequestContext& Context,
			 std::string_view			Namespace,
			 std::string_view			Bucket,
			 const IoHash&				HashKey,
			 ZenCacheValue&				OutValue);
	void Put(const CacheRequestContext& Context,
			 std::string_view			Namespace,
			 std::string_view			Bucket,
			 const IoHash&				HashKey,
			 const ZenCacheValue&		Value,
			 std::span<IoHash>			References);
	bool DropBucket(std::string_view Namespace, std::string_view Bucket);
	bool DropNamespace(std::string_view Namespace);
	void Flush();
	void ScrubStorage(ScrubContext& Ctx);

	CacheValueDetails GetValueDetails(const std::string_view NamespaceFilter,
									  const std::string_view BucketFilter,
									  const std::string_view ValueFilter) const;

	GcStorageSize	StorageSize() const;
	CacheStoreStats Stats(bool IncludeNamespaceStats = true);

	Configuration								 GetConfiguration() const { return m_Configuration; }
	void										 SetLoggingConfig(const Configuration::LogConfig& Loggingconfig);
	Info										 GetInfo() const;
	std::optional<ZenCacheNamespace::Info>		 GetNamespaceInfo(std::string_view Namespace);
	std::optional<ZenCacheNamespace::BucketInfo> GetBucketInfo(std::string_view Namespace, std::string_view Bucket);
	std::vector<std::string>					 GetNamespaces();

	void EnumerateBucketContents(std::string_view																		  Namespace,
								 std::string_view																		  Bucket,
								 std::function<void(const IoHash& Key, const CacheValueDetails::ValueDetails& Details)>&& Fn);

	// StatsProvider
	virtual void ReportMetrics(StatsMetrics& Statsd) override;

private:
	const ZenCacheNamespace* FindNamespace(std::string_view Namespace) const;
	ZenCacheNamespace*		 GetNamespace(std::string_view Namespace);
	void IterateNamespaces(const std::function<void(std::string_view Namespace, ZenCacheNamespace& Store)>& Callback) const;

	typedef std::unordered_map<std::string, std::unique_ptr<ZenCacheNamespace>> NamespaceMap;

	CacheStoreStats									m_LastReportedMetrics;
	const DiskWriteBlocker*							m_DiskWriteBlocker = nullptr;
	mutable RwLock									m_NamespacesLock;
	NamespaceMap									m_Namespaces;
	std::vector<std::unique_ptr<ZenCacheNamespace>> m_DroppedNamespaces;

	GcManager&			  m_Gc;
	JobQueue&			  m_JobQueue;
	std::filesystem::path m_BasePath;
	Configuration		  m_Configuration;
	std::atomic<uint64_t> m_HitCount{};
	std::atomic<uint64_t> m_MissCount{};
	std::atomic<uint64_t> m_WriteCount{};
	std::atomic<uint64_t> m_RejectedWriteCount{};
	std::atomic<uint64_t> m_RejectedReadCount{};
	metrics::RequestStats m_PutOps;
	metrics::RequestStats m_GetOps;

	struct AccessLogItem
	{
		const char*			Op;
		CacheRequestContext Context;
		std::string			Namespace;
		std::string			Bucket;
		IoHash				HashKey;
		ZenCacheValue		Value;
	};

	void					   LogWorker();
	RwLock					   m_LogQueueLock;
	std::vector<AccessLogItem> m_LogQueue;
	std::atomic_bool		   m_ExitLogging;
	Event					   m_LogEvent;
	std::thread				   m_AsyncLoggingThread;
	std::atomic_bool		   m_WriteLogEnabled;
	std::atomic_bool		   m_AccessLogEnabled;
};

void z$_forcelink();

}  // namespace zen