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

#pragma once

#include "cacheshared.h"

#include <zencore/iohash.h>
#include <zenstore/gc.h>

#include <inttypes.h>
#include <string_view>
#include <vector>

ZEN_THIRD_PARTY_INCLUDES_START
#include <tsl/robin_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

/** 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:
	struct Configuration
	{
		uint64_t TargetFootprintBytes = 16 * 1024 * 1024;
		uint64_t ScavengeThreshold	  = 4 * 1024 * 1024;
	};

	struct BucketInfo
	{
		uint64_t EntryCount = 0;
		uint64_t TotalSize	= 0;
	};

	struct Info
	{
		Configuration			 Config;
		std::vector<std::string> BucketNames;
		uint64_t				 EntryCount = 0;
		uint64_t				 TotalSize	= 0;
	};

	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);
	void	 Drop();
	bool	 DropBucket(std::string_view Bucket);
	void	 ScrubStorage(ScrubContext& Ctx);
	void	 GatherAccessTimes(zen::access_tracking::AccessTimes& AccessTimes);
	void	 Reset();
	uint64_t TotalSize() const;

	Info					  GetInfo() const;
	std::optional<BucketInfo> GetBucketInfo(std::string_view Bucket) const;

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

private:
	struct CacheBucket
	{
#pragma pack(push)
#pragma pack(1)
		struct BucketPayload
		{
			IoBuffer Payload;  // 8
			uint32_t RawSize;  // 4
			IoHash	 RawHash;  // 20
		};
#pragma pack(pop)
		static_assert(sizeof(BucketPayload) == 32u);
		static_assert(sizeof(AccessTime) == 4u);

		mutable RwLock					 m_BucketLock;
		std::vector<AccessTime>			 m_AccessTimes;
		std::vector<BucketPayload>		 m_Payloads;
		tsl::robin_map<IoHash, uint32_t> m_CacheMap;

		std::atomic_uint64_t m_TotalSize{};

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

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

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

}  // namespace zen