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

#pragma once

#include "cachedisklayer.h"
#include "cachememorylayer.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 {

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

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

  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
	{
		std::filesystem::path RootDir;
		uint64_t			  DiskLayerThreshold = 0;
	};
	struct BucketInfo
	{
		ZenCacheDiskLayer::BucketInfo	DiskLayerInfo;
		ZenCacheMemoryLayer::BucketInfo MemoryLayerInfo;
	};
	struct Info
	{
		Configuration			  Config;
		std::vector<std::string>  BucketNames;
		ZenCacheDiskLayer::Info	  DiskLayerInfo;
		ZenCacheMemoryLayer::Info MemoryLayerInfo;
	};

	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,
					  bool									   EnableReferenceCaching,
					  const ZenCacheMemoryLayer::Configuration MemLayerConfig = {});
	~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();
	uint64_t DiskLayerThreshold() const { return m_DiskLayerSizeThreshold; }

	// 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;

	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;

private:
	std::filesystem::path m_RootDir;
	JobQueue&			  m_JobQueue;
	ZenCacheMemoryLayer	  m_MemLayer;
	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_DiskLayerSizeThreshold = 1 * 1024;
	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:
	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
	{
		std::filesystem::path			   BasePath;
		bool							   AllowAutomaticCreationOfNamespaces = false;
		bool							   EnableReferenceCaching			  = false;
		ZenCacheMemoryLayer::Configuration MemLayerConfig;
		struct LogConfig
		{
			bool EnableWriteLog	 = true;
			bool EnableAccessLog = true;
		} Logging;
	};

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

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

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

	ZenCacheStore(GcManager& Gc, JobQueue& JobQueue, 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();

	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);

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;

	spdlog::logger&									m_Log;
	spdlog::logger&									Log() { return m_Log; }
	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;
	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