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

#include "structuredcachestore.h"

#include <zencore/except.h>
#include <zencore/windows.h>

#include <zencore/compactbinary.h>
#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/compactbinaryvalidation.h>
#include <zencore/compress.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/logging.h>
#include <zencore/string.h>
#include <zencore/testing.h>
#include <zencore/testutils.h>
#include <zencore/thread.h>
#include <zenstore/basicfile.h>
#include <zenstore/cas.h>
#include <zenstore/caslog.h>
#include <zenstore/cidstore.h>

#include <concepts>
#include <filesystem>
#include <unordered_map>
#include <ranges>

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/core.h>
#include <gsl/gsl-lite.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

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

namespace zen {

using namespace fmt::literals;

ZenCacheStore::ZenCacheStore(const std::filesystem::path& RootDir) : m_DiskLayer{RootDir}
{
	ZEN_INFO("initializing structured cache at '{}'", RootDir);
	CreateDirectories(RootDir);

	m_DiskLayer.DiscoverBuckets();
}

ZenCacheStore::~ZenCacheStore()
{
}

bool
ZenCacheStore::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
{
	bool Ok = m_MemLayer.Get(InBucket, HashKey, OutValue);

	if (Ok)
	{
		ZEN_ASSERT(OutValue.Value.Size());
	}

	if (!Ok)
	{
		Ok = m_DiskLayer.Get(InBucket, HashKey, OutValue);

		if (Ok)
		{
			ZEN_ASSERT(OutValue.Value.Size());
		}

		if (Ok && (OutValue.Value.Size() <= m_DiskLayerSizeThreshold))
		{
			m_MemLayer.Put(InBucket, HashKey, OutValue);
		}
	}

	return Ok;
}

void
ZenCacheStore::Put(std::string_view InBucket, const IoHash& HashKey, const ZenCacheValue& Value)
{
	// Store value and index

	ZEN_ASSERT(Value.Value.Size());

	m_DiskLayer.Put(InBucket, HashKey, Value);

	if (Value.Value.Size() <= m_DiskLayerSizeThreshold)
	{
		m_MemLayer.Put(InBucket, HashKey, Value);
	}
}

bool
ZenCacheStore::DropBucket(std::string_view Bucket)
{
	ZEN_INFO("dropping bucket '{}'", Bucket);

	// TODO: should ensure this is done atomically across all layers

	const bool MemDropped  = m_MemLayer.DropBucket(Bucket);
	const bool DiskDropped = m_DiskLayer.DropBucket(Bucket);
	const bool AnyDropped  = MemDropped || DiskDropped;

	ZEN_INFO("bucket '{}' was {}", Bucket, AnyDropped ? "dropped" : "not found");

	return AnyDropped;
}

void
ZenCacheStore::Flush()
{
	m_DiskLayer.Flush();
}

void
ZenCacheStore::Scrub(ScrubContext& Ctx)
{
	if (m_LastScrubTime == Ctx.ScrubTimestamp())
	{
		return;
	}

	m_LastScrubTime = Ctx.ScrubTimestamp();

	m_DiskLayer.Scrub(Ctx);
	m_MemLayer.Scrub(Ctx);
}

void
ZenCacheStore::GarbageCollect(GcContext& GcCtx)
{
	m_DiskLayer.GarbageCollect(GcCtx);
	m_MemLayer.GarbageCollect(GcCtx);
}

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

ZenCacheMemoryLayer::ZenCacheMemoryLayer()
{
}

ZenCacheMemoryLayer::~ZenCacheMemoryLayer()
{
}

bool
ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
{
	RwLock::SharedLockScope _(m_Lock);

	auto it = m_Buckets.find(std::string(InBucket));

	if (it == m_Buckets.end())
	{
		return false;
	}

	CacheBucket* Bucket = Bucket = &it->second;

	_.ReleaseNow();

	// There's a race here. Since the lock is released early to allow
	// inserts, the bucket delete path could end up deleting the
	// underlying data structure

	return Bucket->Get(HashKey, OutValue);
}

void
ZenCacheMemoryLayer::Put(std::string_view InBucket, const IoHash& HashKey, const ZenCacheValue& Value)
{
	CacheBucket* Bucket = nullptr;

	{
		RwLock::SharedLockScope _(m_Lock);

		auto it = m_Buckets.find(std::string(InBucket));

		if (it != m_Buckets.end())
		{
			Bucket = &it->second;
		}
	}

	if (Bucket == nullptr)
	{
		// New bucket

		RwLock::ExclusiveLockScope _(m_Lock);

		Bucket = &m_Buckets[std::string(InBucket)];
	}

	// Note that since the underlying IoBuffer is retained, the content type is also

	Bucket->Put(HashKey, Value);
}

bool
ZenCacheMemoryLayer::DropBucket(std::string_view Bucket)
{
	RwLock::ExclusiveLockScope _(m_Lock);

	return !!m_Buckets.erase(std::string(Bucket));
}

void
ZenCacheMemoryLayer::Scrub(ScrubContext& Ctx)
{
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		Kv.second.Scrub(Ctx);
	}
}

void
ZenCacheMemoryLayer::GarbageCollect(GcContext& GcCtx)
{
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		Kv.second.GarbageCollect(GcCtx);
	}
}

void
ZenCacheMemoryLayer::CacheBucket::Scrub(ScrubContext& Ctx)
{
	RwLock::SharedLockScope _(m_bucketLock);

	std::vector<IoHash> BadHashes;

	for (auto& Kv : m_cacheMap)
	{
		if (Kv.first != IoHash::HashBuffer(Kv.second.Payload))
		{
			BadHashes.push_back(Kv.first);
		}
	}

	if (!BadHashes.empty())
	{
		Ctx.ReportBadCasChunks(BadHashes);
	}
}

void
ZenCacheMemoryLayer::CacheBucket::GarbageCollect(GcContext& GcCtx)
{
	// Is it even meaningful to do this? The memory layer shouldn't
	// contain anything which is not already in the disk layer

	RwLock::SharedLockScope _(m_bucketLock);

	std::vector<IoHash> BadHashes;

	for (const auto& Kv : m_cacheMap)
	{
		const IoBuffer& Payload = Kv.second.Payload;

		switch (Payload.GetContentType())
		{
			case ZenContentType::kCbObject:
				{
					CbObject Obj(SharedBuffer{Payload});

					Obj.IterateAttachments([&](CbFieldView Field) { GcCtx.ContributeCids(std::array<IoHash, 1>{Field.AsAttachment()}); });
				}
				break;

			case ZenContentType::kBinary:
				break;
		}
	}
}

bool
ZenCacheMemoryLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
{
	RwLock::SharedLockScope _(m_bucketLock);

	if (auto bucketIt = m_cacheMap.find(HashKey); bucketIt == m_cacheMap.end())
	{
		return false;
	}
	else
	{
		BucketValue& Value = bucketIt.value();
		OutValue.Value	   = Value.Payload;
		Value.LastAccess   = GetCurrentTimeStamp();

		return true;
	}
}

uint64_t
ZenCacheMemoryLayer::CacheBucket::GetCurrentTimeStamp()
{
	return GetLofreqTimerValue();
}

void
ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
{
	RwLock::ExclusiveLockScope _(m_bucketLock);

	m_cacheMap.insert_or_assign(HashKey, BucketValue{.LastAccess = GetCurrentTimeStamp(), .Payload = Value.Value});
}

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

#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;
	static const uint64_t kFlagsMask	  = 0xff00'0000'0000'0000ull;
	static const uint64_t kStandaloneFile = 0x8000'0000'0000'0000ull;
	static const uint64_t kStructured	  = 0x4000'0000'0000'0000ull;
	static const uint64_t kTombStone	  = 0x2000'0000'0000'0000ull;

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

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

struct ZenCacheDiskLayer::CacheBucket
{
	CacheBucket();
	~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 Scrub(ScrubContext& Ctx);
	void GarbageCollect(GcContext& GcCtx);

	inline bool IsOk() const { return m_IsOk; }

private:
	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;

	RwLock												 m_IndexLock;
	tsl::robin_map<IoHash, DiskLocation, IoHash::Hasher> m_Index;
	uint64_t											 m_WriteCursor = 0;

	void BuildPath(WideStringBuilderBase& Path, const IoHash& HashKey);
	void PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value);
	bool GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey, ZenCacheValue& OutValue);
	bool GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue);

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

ZenCacheDiskLayer::CacheBucket::CacheBucket()
{
}

ZenCacheDiskLayer::CacheBucket::~CacheBucket()
{
}

bool
ZenCacheDiskLayer::CacheBucket::Delete(std::filesystem::path BucketDir)
{
	if (std::filesystem::exists(BucketDir))
	{
		DeleteDirectories(BucketDir);

		return true;
	}

	return false;
}

void
ZenCacheDiskLayer::CacheBucket::OpenOrCreate(std::filesystem::path BucketDir, bool AllowCreate)
{
	CreateDirectories(BucketDir);

	m_BucketDir = BucketDir;

	std::filesystem::path ManifestPath{m_BucketDir / "zen_manifest"};
	std::filesystem::path SobsPath{m_BucketDir / "zen.sobs"};
	std::filesystem::path SlogPath{m_BucketDir / "zen.slog"};

	BasicFile ManifestFile;

	// Try opening existing manifest file first

	bool IsNew = false;

	std::error_code Ec;
	ManifestFile.Open(ManifestPath, /* IsCreate */ false, Ec);

	if (!Ec)
	{
		uint64_t FileSize = ManifestFile.FileSize();

		if (FileSize == sizeof(Oid))
		{
			ManifestFile.Read(&m_BucketId, sizeof(Oid), 0);

			m_IsOk = true;
		}

		if (!m_IsOk)
		{
			ManifestFile.Close();
		}
	}

	if (!m_IsOk)
	{
		if (AllowCreate == false)
		{
			// Invalid bucket
			return;
		}

		// No manifest file found, this is a new bucket

		ManifestFile.Open(ManifestPath, /* IsCreate */ true, Ec);

		if (Ec)
		{
			throw std::system_error(Ec, "Failed to create bucket manifest '{}'"_format(ManifestPath));
		}

		m_BucketId.Generate();

		ManifestFile.Write(&m_BucketId, sizeof(Oid), /* FileOffset */ 0);

		IsNew = true;
	}

	// Initialize small object storage related files

	m_SobsFile.Open(SobsPath, IsNew);

	// Open and replay log

	m_SlogFile.Open(SlogPath, IsNew);

	uint64_t MaxFileOffset = 0;

	if (RwLock::ExclusiveLockScope _(m_IndexLock); m_Index.empty())
	{
		m_SlogFile.Replay([&](const DiskIndexEntry& Record) {
			m_Index[Record.Key] = Record.Location;

			MaxFileOffset = std::max<uint64_t>(MaxFileOffset, Record.Location.Offset() + Record.Location.Size());
		});

		m_WriteCursor = (MaxFileOffset + 15) & ~15;
	}

	m_IsOk = true;
}

void
ZenCacheDiskLayer::CacheBucket::BuildPath(WideStringBuilderBase& Path, const IoHash& HashKey)
{
	char HexString[sizeof(HashKey.Hash) * 2];
	ToHexBytes(HashKey.Hash, sizeof HashKey.Hash, HexString);

	Path.Append(m_BucketDir.c_str());
	Path.Append(L"/blob/");
	Path.AppendAsciiRange(HexString, HexString + 3);
	Path.Append(L"/");
	Path.AppendAsciiRange(HexString + 3, HexString + 5);
	Path.Append(L"/");
	Path.AppendAsciiRange(HexString + 5, HexString + sizeof(HexString));
}

bool
ZenCacheDiskLayer::CacheBucket::GetInlineCacheValue(const DiskLocation& Loc, ZenCacheValue& OutValue)
{
	if (Loc.IsFlagSet(DiskLocation::kStandaloneFile))
	{
		return false;
	}

	OutValue.Value = IoBufferBuilder::MakeFromFileHandle(m_SobsFile.Handle(), Loc.Offset(), Loc.Size());
	OutValue.Value.SetContentType(Loc.GetContentType());

	return true;
}

bool
ZenCacheDiskLayer::CacheBucket::GetStandaloneCacheValue(const DiskLocation& Loc, const IoHash& HashKey, ZenCacheValue& OutValue)
{
	WideStringBuilder<128> DataFilePath;
	BuildPath(DataFilePath, HashKey);

	RwLock::SharedLockScope ValueLock(LockForHash(HashKey));

	if (IoBuffer Data = IoBufferBuilder::MakeFromFile(DataFilePath.c_str()))
	{
		OutValue.Value = Data;
		OutValue.Value.SetContentType(Loc.GetContentType());

		return true;
	}

	return false;
}

bool
ZenCacheDiskLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
{
	if (!m_IsOk)
	{
		return false;
	}

	RwLock::SharedLockScope _(m_IndexLock);

	if (auto it = m_Index.find(HashKey); it != m_Index.end())
	{
		const DiskLocation& Loc = it->second;

		if (GetInlineCacheValue(Loc, OutValue))
		{
			return true;
		}

		_.ReleaseNow();

		return GetStandaloneCacheValue(Loc, HashKey, OutValue);
	}

	return false;
}

void
ZenCacheDiskLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
{
	if (!m_IsOk)
	{
		return;
	}

	if (Value.Value.Size() >= m_LargeObjectThreshold)
	{
		return PutStandaloneCacheValue(HashKey, Value);
	}
	else
	{
		// Small object put

		uint64_t EntryFlags = 0;

		if (Value.Value.GetContentType() == ZenContentType::kCbObject)
		{
			EntryFlags |= DiskLocation::kStructured;
		}

		RwLock::ExclusiveLockScope _(m_IndexLock);

		DiskLocation Loc(m_WriteCursor, Value.Value.Size(), 0, EntryFlags);

		m_WriteCursor = RoundUp(m_WriteCursor + Loc.Size(), 16);

		if (auto it = m_Index.find(HashKey); it == m_Index.end())
		{
			// Previously unknown object
			m_Index.insert({HashKey, Loc});
		}
		else
		{
			// TODO: should check if write is idempotent and bail out if it is?
			// this would requiring comparing contents on disk unless we add a
			// content hash to the index entry
			it.value() = Loc;
		}

		m_SlogFile.Append({.Key = HashKey, .Location = Loc});
		m_SobsFile.Write(Value.Value.Data(), Loc.Size(), Loc.Offset());
	}
}

void
ZenCacheDiskLayer::CacheBucket::Drop()
{
	// TODO: add error handling

	m_SobsFile.Close();
	m_SlogFile.Close();
	DeleteDirectories(m_BucketDir);
}

void
ZenCacheDiskLayer::CacheBucket::Flush()
{
	m_SobsFile.Flush();
	m_SlogFile.Flush();
}

void
ZenCacheDiskLayer::CacheBucket::Scrub(ScrubContext& Ctx)
{
	std::vector<IoHash> BadKeys;

	{
		RwLock::SharedLockScope _(m_IndexLock);

		for (auto& Kv : m_Index)
		{
			const IoHash&		HashKey = Kv.first;
			const DiskLocation& Loc		= Kv.second;

			ZenCacheValue Value;

			if (GetInlineCacheValue(Loc, Value))
			{
				// Validate contents
			}
			else if (GetStandaloneCacheValue(Loc, HashKey, Value))
			{
				// Note: we cannot currently validate contents since we don't
				// have a content hash!
			}
			else
			{
				// Value not found
				BadKeys.push_back(HashKey);
			}
		}
	}

	if (BadKeys.empty())
	{
		return;
	}

	if (Ctx.RunRecovery())
	{
		RwLock::ExclusiveLockScope _(m_IndexLock);

		for (const IoHash& BadKey : BadKeys)
		{
			// Log a tombstone and delete the in-memory index for the bad entry

			m_SlogFile.Append(DiskIndexEntry{.Key = BadKey, .Location = {0, 0, 0, DiskLocation::kTombStone}});
			m_Index.erase(BadKey);
		}
	}
}

void
ZenCacheDiskLayer::CacheBucket::GarbageCollect(GcContext& GcCtx)
{
	RwLock::SharedLockScope _(m_IndexLock);

	for (auto& Kv : m_Index)
	{
		const IoHash&		HashKey = Kv.first;
		const DiskLocation& Loc		= Kv.second;

		if (Loc.IsFlagSet(DiskLocation::kStructured) == false)
		{
			continue;
		}

		ZenCacheValue		CacheValue;
		std::vector<IoHash> Attachments;

		auto GatherRefs = [&] {
			ZEN_ASSERT(CacheValue.Value.GetContentType() == ZenContentType::kCbObject);

			CbObject Obj(SharedBuffer{CacheValue.Value});
			Obj.IterateAttachments([&](CbFieldView Field) { Attachments.push_back(Field.AsAttachment()); });
			GcCtx.ContributeCids(Attachments);
		};

		if (GetInlineCacheValue(Loc, /* out */ CacheValue))
		{
			GatherRefs();
		}
		else if (GetStandaloneCacheValue(Loc, HashKey, /* out */ CacheValue))
		{
			GatherRefs();
		}
		else
		{
			// Value not found
		}

		GcCtx.ContributeCids(Attachments);
	}
}

void
ZenCacheDiskLayer::CacheBucket::PutStandaloneCacheValue(const IoHash& HashKey, const ZenCacheValue& Value)
{
	RwLock::ExclusiveLockScope ValueLock(LockForHash(HashKey));

	WideStringBuilder<128> DataFilePath;
	BuildPath(DataFilePath, HashKey);

	TemporaryFile DataFile;

	std::error_code Ec;
	DataFile.CreateTemporary(m_BucketDir.c_str(), Ec);

	if (Ec)
	{
		throw std::system_error(Ec, "Failed to open temporary file for put at '{}'"_format(m_BucketDir));
	}

	DataFile.WriteAll(Value.Value, Ec);

	if (Ec)
	{
		throw std::system_error(Ec, "Failed to write payload ({} bytes) to file"_format(NiceBytes(Value.Value.Size())));
	}

	// Move file into place (atomically)

	std::filesystem::path FsPath{DataFilePath.c_str()};

	DataFile.MoveTemporaryIntoPlace(FsPath, Ec);

	if (Ec)
	{
		int RetryCount = 3;

		do
		{
			std::filesystem::path ParentPath = std::filesystem::path(DataFilePath.c_str()).parent_path();
			CreateDirectories(ParentPath);

			DataFile.MoveTemporaryIntoPlace(FsPath, Ec);

			if (!Ec)
			{
				break;
			}

			std::error_code InnerEc;
			const uint64_t	ExistingFileSize = std::filesystem::file_size(FsPath, InnerEc);

			if (!InnerEc && ExistingFileSize == Value.Value.Size())
			{
				// Concurrent write of same value?
				return;
			}

			// Semi arbitrary back-off
			zen::Sleep(1000 * RetryCount);
		} while (RetryCount--);

		if (Ec)
		{
			throw std::system_error(Ec, "Failed to finalize file '{}'"_format(WideToUtf8(DataFilePath)));
		}
	}

	// Update index

	uint64_t EntryFlags = DiskLocation::kStandaloneFile;

	if (Value.Value.GetContentType() == ZenContentType::kCbObject)
	{
		EntryFlags |= DiskLocation::kStructured;
	}

	RwLock::ExclusiveLockScope _(m_IndexLock);

	DiskLocation Loc(/* Offset */ 0, Value.Value.Size(), 0, EntryFlags);

	if (auto it = m_Index.find(HashKey); it == m_Index.end())
	{
		// Previously unknown object
		m_Index.insert({HashKey, Loc});
	}
	else
	{
		// TODO: should check if write is idempotent and bail out if it is?
		it.value() = Loc;
	}

	m_SlogFile.Append({.Key = HashKey, .Location = Loc});
}

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

ZenCacheDiskLayer::ZenCacheDiskLayer(const std::filesystem::path& RootDir) : m_RootDir(RootDir)
{
}

ZenCacheDiskLayer::~ZenCacheDiskLayer() = default;

bool
ZenCacheDiskLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
{
	CacheBucket* Bucket = nullptr;

	{
		RwLock::SharedLockScope _(m_Lock);

		auto it = m_Buckets.find(std::string(InBucket));

		if (it != m_Buckets.end())
		{
			Bucket = &it->second;
		}
	}

	if (Bucket == nullptr)
	{
		// Bucket needs to be opened/created

		RwLock::ExclusiveLockScope _(m_Lock);

		if (auto it = m_Buckets.find(std::string(InBucket)); it != m_Buckets.end())
		{
			Bucket = &it->second;
		}
		else
		{
			auto It = m_Buckets.try_emplace(std::string(InBucket));
			Bucket	= &It.first->second;

			std::filesystem::path BucketPath = m_RootDir;
			BucketPath /= std::string(InBucket);

			Bucket->OpenOrCreate(BucketPath);
		}
	}

	ZEN_ASSERT(Bucket != nullptr);

	return Bucket->Get(HashKey, OutValue);
}

void
ZenCacheDiskLayer::Put(std::string_view InBucket, const IoHash& HashKey, const ZenCacheValue& Value)
{
	CacheBucket* Bucket = nullptr;

	{
		RwLock::SharedLockScope _(m_Lock);

		auto it = m_Buckets.find(std::string(InBucket));

		if (it != m_Buckets.end())
		{
			Bucket = &it->second;
		}
	}

	if (Bucket == nullptr)
	{
		// New bucket needs to be created

		RwLock::ExclusiveLockScope _(m_Lock);

		if (auto it = m_Buckets.find(std::string(InBucket)); it != m_Buckets.end())
		{
			Bucket = &it->second;
		}
		else
		{
			auto It = m_Buckets.try_emplace(std::string(InBucket));
			Bucket	= &It.first->second;

			std::filesystem::path bucketPath = m_RootDir;
			bucketPath /= std::string(InBucket);

			Bucket->OpenOrCreate(bucketPath);
		}
	}

	ZEN_ASSERT(Bucket != nullptr);

	if (Bucket->IsOk())
	{
		Bucket->Put(HashKey, Value);
	}
}

void
ZenCacheDiskLayer::DiscoverBuckets()
{
	FileSystemTraversal Traversal;
	struct Visitor : public FileSystemTraversal::TreeVisitor
	{
		virtual void VisitFile([[maybe_unused]] const std::filesystem::path& Parent,
							   [[maybe_unused]] const path_view&			 File,
							   [[maybe_unused]] uint64_t					 FileSize) override
		{
		}

		virtual bool VisitDirectory([[maybe_unused]] const std::filesystem::path& Parent, const path_view& DirectoryName) override
		{
			Dirs.push_back(std::wstring(DirectoryName));
			return false;
		}

		std::vector<std::wstring> Dirs;
	} Visit;

	Traversal.TraverseFileSystem(m_RootDir, Visit);

	// Initialize buckets

	RwLock::ExclusiveLockScope _(m_Lock);

	for (const std::wstring& BucketName : Visit.Dirs)
	{
		// New bucket needs to be created

		std::string BucketName8 = WideToUtf8(BucketName);

		if (auto It = m_Buckets.find(BucketName8); It != m_Buckets.end())
		{
		}
		else
		{
			auto InsertResult = m_Buckets.try_emplace(BucketName8);

			std::filesystem::path BucketPath = m_RootDir;
			BucketPath /= BucketName8;

			CacheBucket& Bucket = InsertResult.first->second;

			Bucket.OpenOrCreate(BucketPath, /* AllowCreate */ false);

			if (!Bucket.IsOk())
			{
				ZEN_WARN("Found directory '{}' in our base directory '{}' but it is not a valid bucket", BucketName8, m_RootDir);

				m_Buckets.erase(InsertResult.first);
			}
		}
	}
}

bool
ZenCacheDiskLayer::DropBucket(std::string_view InBucket)
{
	RwLock::ExclusiveLockScope _(m_Lock);

	auto it = m_Buckets.find(std::string(InBucket));

	if (it != m_Buckets.end())
	{
		CacheBucket* Bucket = &it->second;

		Bucket->Drop();

		m_Buckets.erase(it);

		return true;
	}

	std::filesystem::path BucketPath = m_RootDir;
	BucketPath /= std::string(InBucket);

	return CacheBucket::Delete(BucketPath);
}

void
ZenCacheDiskLayer::Flush()
{
	std::vector<CacheBucket*> Buckets;
	Buckets.reserve(m_Buckets.size());

	{
		RwLock::SharedLockScope _(m_Lock);

		for (auto& Kv : m_Buckets)
		{
			Buckets.push_back(&Kv.second);
		}
	}

	for (auto& Bucket : Buckets)
	{
		Bucket->Flush();
	}
}

void
ZenCacheDiskLayer::Scrub(ScrubContext& Ctx)
{
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		Kv.second.Scrub(Ctx);
	}
}

void
ZenCacheDiskLayer::GarbageCollect(GcContext& GcCtx)
{
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		Kv.second.GarbageCollect(GcCtx);
	}
}

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

#if ZEN_WITH_TESTS

TEST_CASE("z$.store")
{
	using namespace fmt::literals;
	using namespace std::literals;

	ScopedTemporaryDirectory TempDir;

	ZenCacheStore Zcs(TempDir.Path() / "cache");

	const int kIterationCount = 100;

	for (int i = 0; i < kIterationCount; ++i)
	{
		const IoHash Key = IoHash::HashBuffer(&i, sizeof i);

		CbObjectWriter Cbo;
		Cbo << "hey" << i;
		CbObject Obj = Cbo.Save();

		ZenCacheValue Value;
		Value.Value = Obj.GetBuffer().AsIoBuffer();
		Value.Value.SetContentType(ZenContentType::kCbObject);

		Zcs.Put("test_bucket"sv, Key, Value);
	}

	for (int i = 0; i < kIterationCount; ++i)
	{
		const IoHash Key = IoHash::HashBuffer(&i, sizeof i);

		ZenCacheValue Value;
		Zcs.Get("test_bucket"sv, Key, /* out */ Value);

		REQUIRE(Value.Value);
		CHECK(Value.Value.GetContentType() == ZenContentType::kCbObject);
		CHECK_EQ(ValidateCompactBinary(Value.Value, CbValidateMode::All), CbValidateError::None);
		CbObject Obj = LoadCompactBinaryObject(Value.Value);
		CHECK_EQ(Obj["hey"].AsInt32(), i);
	}
}

TEST_CASE("z$.gc")
{
	using namespace fmt::literals;
	using namespace std::literals;

	ScopedTemporaryDirectory TempDir;

	CasStoreConfiguration	  Config{.RootDirectory = TempDir.Path()};
	std::unique_ptr<CasStore> Cas{CreateCasStore()};
	Cas->Initialize(Config);

	CidStore Cid(*Cas, TempDir.Path());

	ZenCacheStore Zcs(TempDir.Path() / "cache");

	const int kIterationCount = 100;

	for (int i = 0; i < kIterationCount; ++i)
	{
		const IoHash Key = IoHash::HashBuffer(&i, sizeof i);

		CompressedBuffer CompBuf = CompressedBuffer::Compress(SharedBuffer::MakeView("abcd"sv));
		Cid.AddChunk(CompBuf);

		CbObjectWriter Cbo;
		Cbo << "hey" << i << "ref" << CbAttachment(CompBuf);
		CbObject Obj = Cbo.Save();

		ZenCacheValue Value;
		Value.Value = Obj.GetBuffer().AsIoBuffer();
		Value.Value.SetContentType(ZenContentType::kCbObject);

		Zcs.Put("test_bucket"sv, Key, Value);
	}
}

#endif

void
z$_forcelink()
{
}

}  // namespace zen