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

#pragma once

#include <zencore/compactbinary.h>
#include <zencore/iobuffer.h>
#include <zencore/iohash.h>
#include <zencore/thread.h>
#include <zencore/uid.h>
#include <zenstore/basicfile.h>
#include <zenstore/cas.h>
#include <zenstore/caslog.h>
#include <zenstore/gc.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

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

#define ZEN_USE_CACHE_TRACKER 0

namespace zen {

class PathBuilderBase;
class CasStore;
class CasGc;
class ZenCacheTracker;

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

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

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

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

namespace access_tracking {

	struct KeyAccessTime
	{
		IoHash		  Key;
		GcClock::Tick LastAccess{};
	};

	struct AccessTimes
	{
		std::unordered_map<std::string, std::vector<KeyAccessTime>> Buckets;
	};
};	// namespace access_tracking

struct ZenCacheValue
{
	IoBuffer Value;
	CbObject IndexData;
};

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

#pragma pack(push)
#pragma pack(1)

struct DiskLocation
{
	inline DiskLocation() = default;

	inline DiskLocation(uint64_t Offset, uint64_t ValueSize, uint32_t IndexSize, uint64_t Flags)
	: OffsetAndFlags(CombineOffsetAndFlags(Offset, Flags))
	, LowerSize(ValueSize & 0xFFFFffff)
	, IndexDataSize(IndexSize)
	{
	}

	static const uint64_t kOffsetMask	  = 0x0000'ffFF'ffFF'ffFFull;
	static const uint64_t kSizeMask		  = 0x00FF'0000'0000'0000ull;  // Most significant bits of value size (lower 32 bits in LowerSize)
	static const uint64_t kFlagsMask	  = 0xff00'0000'0000'0000ull;
	static const uint64_t kStandaloneFile = 0x8000'0000'0000'0000ull;  // Stored as a separate file
	static const uint64_t kStructured	  = 0x4000'0000'0000'0000ull;  // Serialized as compact binary
	static const uint64_t kTombStone	  = 0x2000'0000'0000'0000ull;  // Represents a deleted key/value
	static const uint64_t kCompressed	  = 0x1000'0000'0000'0000ull;  // Stored in compressed buffer format

	static uint64_t CombineOffsetAndFlags(uint64_t Offset, uint64_t Flags) { return Offset | Flags; }

	inline uint64_t		  Offset() const { return OffsetAndFlags & kOffsetMask; }
	inline uint64_t		  Size() const { return LowerSize; }
	inline uint64_t		  IsFlagSet(uint64_t Flag) const { return OffsetAndFlags & Flag; }
	inline ZenContentType GetContentType() const
	{
		ZenContentType ContentType = ZenContentType::kBinary;

		if (IsFlagSet(DiskLocation::kStructured))
		{
			ContentType = ZenContentType::kCbObject;
		}

		if (IsFlagSet(DiskLocation::kCompressed))
		{
			ContentType = ZenContentType::kCompressedBinary;
		}

		return ContentType;
	}

private:
	uint64_t OffsetAndFlags = 0;
	uint32_t LowerSize		= 0;
	uint32_t IndexDataSize	= 0;
};

struct DiskIndexEntry
{
	IoHash		 Key;
	DiskLocation Location;
};

#pragma pack(pop)

static_assert(sizeof(DiskIndexEntry) == 36);

/** In-memory cache storage

	Intended for small values which are frequently accessed

	This should have a better memory management policy to maintain reasonable
	footprint.

 */
class ZenCacheMemoryLayer
{
public:
	ZenCacheMemoryLayer();
	~ZenCacheMemoryLayer();

	bool	 Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
	void	 Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value);
	bool	 DropBucket(std::string_view Bucket);
	void	 Scrub(ScrubContext& Ctx);
	void	 GatherAccessTimes(zen::access_tracking::AccessTimes& AccessTimes);
	void	 Reset();
	uint64_t TotalSize() const;

	struct Configuration
	{
		uint64_t TargetFootprintBytes = 16 * 1024 * 1024;
		uint64_t ScavengeThreshold	  = 4 * 1024 * 1024;
	};

	const Configuration& GetConfiguration() const { return m_Configuration; }
	void				 SetConfiguration(const Configuration& NewConfig) { m_Configuration = NewConfig; }

private:
	struct CacheBucket
	{
		struct BucketValue
		{
			IoBuffer			Payload;
			std::atomic_int64_t LastAccess;

			BucketValue() : Payload(), LastAccess() {}
			BucketValue(IoBuffer Value, const int64_t Timestamp) : Payload(Value), LastAccess(Timestamp) {}
			BucketValue(const BucketValue& V) : Payload(V.Payload), LastAccess(V.LastAccess.load(std::memory_order_relaxed)) {}
			BucketValue(BucketValue&& V) : Payload(std::move(V.Payload)), LastAccess(V.LastAccess.load(std::memory_order_relaxed)) {}

			BucketValue& operator=(const BucketValue& V) { return *this = BucketValue(V); }
			BucketValue& operator=(BucketValue&& V)
			{
				Payload = std::move(V.Payload);
				LastAccess.store(V.LastAccess.load(), std::memory_order_relaxed);
				return *this;
			}
		};

		RwLock								m_BucketLock;
		tsl::robin_map<IoHash, BucketValue> m_CacheMap;
		std::atomic_uint64_t				m_TotalSize{};

		bool			Get(const IoHash& HashKey, ZenCacheValue& OutValue);
		void			Put(const IoHash& HashKey, const ZenCacheValue& Value);
		void			Scrub(ScrubContext& Ctx);
		void			GatherAccessTimes(std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes);
		inline uint64_t TotalSize() const { return m_TotalSize; }
	};

	mutable RwLock								 m_Lock;
	std::unordered_map<std::string, CacheBucket> m_Buckets;
	Configuration								 m_Configuration;

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

class ZenCacheDiskLayer
{
public:
	explicit ZenCacheDiskLayer(const std::filesystem::path& RootDir);
	~ZenCacheDiskLayer();

	bool Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
	void Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value);
	bool DropBucket(std::string_view Bucket);
	void Flush();
	void Scrub(ScrubContext& Ctx);
	void GatherReferences(GcContext& GcCtx);
	void CollectGarbage(GcContext& GcCtx);
	void UpdateAccessTimes(const zen::access_tracking::AccessTimes& AccessTimes);

	void	 DiscoverBuckets();
	uint64_t TotalSize() const;

private:
	/** A cache bucket manages a single directory containing
		metadata and data for that bucket
	  */
	struct CacheBucket
	{
		CacheBucket(std::string BucketName);
		~CacheBucket();

		void		OpenOrCreate(std::filesystem::path BucketDir, bool AllowCreate = true);
		static bool Delete(std::filesystem::path BucketDir);
		bool		Get(const IoHash& HashKey, ZenCacheValue& OutValue);
		void		Put(const IoHash& HashKey, const ZenCacheValue& Value);
		void		Drop();
		void		Flush();
		void		SaveManifest();
		void		Scrub(ScrubContext& Ctx);
		void		GatherReferences(GcContext& GcCtx);
		void		CollectGarbage(GcContext& GcCtx);
		void		UpdateAccessTimes(const std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes);

		inline bool		IsOk() const { return m_IsOk; }
		inline uint64_t TotalSize() const { return m_TotalSize.load(std::memory_order::relaxed); }

	private:
		std::string			  m_BucketName;
		std::filesystem::path m_BucketDir;
		Oid					  m_BucketId;
		bool				  m_IsOk				 = false;
		uint64_t			  m_LargeObjectThreshold = 64 * 1024;

		// These files are used to manage storage of small objects for this bucket

		BasicFile					m_SobsFile;
		TCasLogFile<DiskIndexEntry> m_SlogFile;

		struct IndexEntry
		{
			DiskLocation		Location;
			std::atomic_int64_t LastAccess;

			IndexEntry() : Location(), LastAccess() {}
			IndexEntry(const DiskLocation& Loc, const int64_t Timestamp) : Location(Loc), LastAccess(Timestamp) {}
			IndexEntry(const IndexEntry& E) : Location(E.Location), LastAccess(E.LastAccess.load(std::memory_order_relaxed)) {}
			IndexEntry(IndexEntry&& E) : Location(std::move(E.Location)), LastAccess(E.LastAccess.load(std::memory_order_relaxed)) {}

			IndexEntry& operator=(const IndexEntry& E) { return *this = IndexEntry(E); }
			IndexEntry& operator=(IndexEntry&& E)
			{
				Location = std::move(E.Location);
				LastAccess.store(E.LastAccess.load(), std::memory_order_relaxed);
				return *this;
			}
		};

		using IndexMap = tsl::robin_map<IoHash, IndexEntry, IoHash::Hasher>;

		RwLock				 m_IndexLock;
		IndexMap			 m_Index;
		uint64_t			 m_SobsCursor = 0;
		std::atomic_uint64_t m_TotalSize{};

		void BuildPath(PathBuilderBase& Path, const IoHash& HashKey);
		void PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
		bool GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey, ZenCacheValue& OutValue);
		void DeleteStandaloneCacheValue(const DiskLocation&			 Loc,
										const IoHash&				 HashKey,
										const std::filesystem::path& Path,
										std::error_code&			 Ec);
		bool GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue);
		void OpenLog(const std::filesystem::path& BucketDir, const bool IsNew);

		// These locks are here to avoid contention on file creation, therefore it's sufficient
		// that we take the same lock for the same hash
		//
		// These locks are small and should really be spaced out so they don't share cache lines,
		// but we don't currently access them at particularly high frequency so it should not be
		// an issue in practice

		RwLock		   m_ShardedLocks[256];
		inline RwLock& LockForHash(const IoHash& Hash) { return m_ShardedLocks[Hash.Hash[19]]; }
	};

	std::filesystem::path						 m_RootDir;
	mutable RwLock								 m_Lock;
	std::unordered_map<std::string, CacheBucket> m_Buckets;	 // TODO: make this case insensitive

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

class ZenCacheStore final : public GcStorage, public GcContributor
{
public:
	ZenCacheStore(CasGc& Gc, const std::filesystem::path& RootDir);
	~ZenCacheStore();

	bool				  Get(std::string_view Bucket, const IoHash& HashKey, ZenCacheValue& OutValue);
	void				  Put(std::string_view Bucket, const IoHash& HashKey, const ZenCacheValue& Value);
	bool				  DropBucket(std::string_view Bucket);
	void				  Flush();
	void				  Scrub(ScrubContext& Ctx);
	uint64_t			  DiskLayerThreshold() const { return m_DiskLayerSizeThreshold; }
	virtual void		  GatherReferences(GcContext& GcCtx) override;
	virtual void		  CollectGarbage(GcContext& GcCtx) override;
	virtual GcStorageSize StorageSize() const override;

private:
	std::filesystem::path m_RootDir;
	ZenCacheMemoryLayer	  m_MemLayer;
	ZenCacheDiskLayer	  m_DiskLayer;
	uint64_t			  m_DiskLayerSizeThreshold = 1 * 1024;
	uint64_t			  m_LastScrubTime		   = 0;

#if ZEN_USE_CACHE_TRACKER
	std::unique_ptr<ZenCacheTracker> m_AccessTracker;
#endif

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

void z$_forcelink();

}  // namespace zen