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

#include "structuredcachestore.h"

#include "cachetracking.h"

#include <zencore/compactbinarybuilder.h>
#include <zencore/compactbinarypackage.h>
#include <zencore/compactbinaryvalidation.h>
#include <zencore/compress.h>
#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/scopeguard.h>
#include <zencore/string.h>
#include <zencore/testing.h>
#include <zencore/testutils.h>
#include <zencore/thread.h>
#include <zencore/windows.h>
#include <zenstore/basicfile.h>
#include <zenstore/caslog.h>
#include <zenstore/cidstore.h>

#include <chrono>
#include <concepts>
#include <memory_resource>
#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;

static CbObject
LoadCompactBinaryObject(const std::filesystem::path& Path)
{
	FileContents Result = ReadFile(Path);

	if (!Result.ErrorCode)
	{
		IoBuffer Buffer = Result.Flatten();
		if (CbValidateError Error = ValidateCompactBinary(Buffer, CbValidateMode::All); Error == CbValidateError::None)
		{
			return LoadCompactBinaryObject(Buffer);
		}
	}

	return CbObject();
}

static void
SaveCompactBinaryObject(const std::filesystem::path& Path, const CbObject& Object)
{
	WriteFile(Path, Object.GetBuffer().AsIoBuffer());
}

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

	m_DiskLayer.DiscoverBuckets();

	m_AccessTracker.reset(new ZenCacheTracker(RootDir));
}

ZenCacheStore::~ZenCacheStore()
{
}

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

	auto _ = MakeGuard([&] {
		if (!Ok)
			return;

		m_AccessTracker->TrackAccess(InBucket, HashKey);
	});

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

		return true;
	}

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

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

		if (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 ZEN_USE_REF_TRACKING
	if (Value.Value.GetContentType() == ZenContentType::kCbObject)
	{
		if (ValidateCompactBinary(Value.Value, CbValidateMode::All) == CbValidateError::None)
		{
			CbObject Object{SharedBuffer(Value.Value)};

			uint8_t								TempBuffer[8 * sizeof(IoHash)];
			std::pmr::monotonic_buffer_resource Linear{TempBuffer, sizeof TempBuffer};
			std::pmr::polymorphic_allocator		Allocator{&Linear};
			std::pmr::vector<IoHash>			CidReferences{Allocator};

			Object.IterateAttachments([&](CbFieldView Field) { CidReferences.push_back(Field.AsAttachment()); });

			m_Gc.OnNewCidReferences(CidReferences);
		}
	}
#endif

	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::GatherReferences(GcContext& GcCtx)
{
	m_MemLayer.GatherReferences(GcCtx);
	m_DiskLayer.GatherReferences(GcCtx);
}

ZenCacheSize
ZenCacheStore::TotalSize() const
{
	return {.MemorySize = m_MemLayer.TotalSize(), .DiskSize = m_DiskLayer.TotalSize()};
}

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

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 = &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::GatherReferences(GcContext& GcCtx)
{
	RwLock::SharedLockScope _(m_Lock);

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

uint64_t
ZenCacheMemoryLayer::TotalSize() const
{
	uint64_t				TotalSize{};
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		TotalSize += Kv.second.TotalSize();
	}

	return TotalSize;
}

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::GatherReferences(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> Cids;

	for (const auto& Kv : m_cacheMap)
	{
		const IoBuffer& Payload = Kv.second.Payload;
		if (Payload.GetContentType() != ZenContentType::kCbObject || GcCtx.Expired(Kv.second.LastAccess))
		{
			continue;
		}

		Cids.clear();
		CbObject Obj(SharedBuffer{Payload});
		Obj.IterateAttachments([&Cids](CbFieldView Field) { Cids.push_back(Field.AsAttachment()); });
		GcCtx.ContributeCids(Cids);
	}
}

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   = GcClock::TickCount();

		return true;
	}
}

void
ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
{
	{
		RwLock::ExclusiveLockScope _(m_bucketLock);
		m_cacheMap.insert_or_assign(HashKey, BucketValue{.LastAccess = GcClock::TickCount(), .Payload = Value.Value});
	}

	m_TotalSize.fetch_add(Value.Value.GetSize());
}

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

#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		GatherReferences(GcContext& GcCtx);

	inline bool		IsOk() const { return m_IsOk; }
	inline uint64_t TotalSize() const { return m_TotalSize; }

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;

	struct IndexEntry
	{
		DiskLocation  Location;
		GcClock::Tick LastAccess{};
	};

	RwLock											   m_IndexLock;
	tsl::robin_map<IoHash, IndexEntry, IoHash::Hasher> m_Index;
	uint64_t										   m_WriteCursor = 0;
	std::atomic_uint64_t							   m_TotalSize{};

	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)
{
	using namespace std::literals;

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

	bool IsNew = false;

	CbObject Manifest = LoadCompactBinaryObject(ManifestPath);

	if (Manifest)
	{
		m_BucketId = Manifest["BucketId"].AsObjectId();
		m_IsOk	   = m_BucketId != Oid::Zero;
	}
	else if (AllowCreate)
	{
		m_BucketId.Generate();

		CbObjectWriter Writer;
		Writer << "BucketId"sv << m_BucketId;
		Manifest = Writer.Save();
		SaveCompactBinaryObject(ManifestPath, Manifest);
		IsNew = true;
	}
	else
	{
		return;
	}

	// Initialize small object storage related files

	m_SobsFile.Open(SobsPath, IsNew);

	// Open and replay log

	m_SlogFile.Open(SlogPath, IsNew);

	uint64_t MaxFileOffset	   = 0;
	uint64_t InvalidEntryCount = 0;

	if (RwLock::ExclusiveLockScope _(m_IndexLock); m_Index.empty())
	{
		m_SlogFile.Replay([&](const DiskIndexEntry& Entry) {
			if (Entry.Key == IoHash::Zero)
			{
				++InvalidEntryCount;
			}
			else if (Entry.Location.IsFlagSet(DiskLocation::kTombStone))
			{
				m_TotalSize.fetch_sub(Entry.Location.Size());
			}
			else
			{
				m_Index[Entry.Key] = {.Location = Entry.Location, .LastAccess = GcClock::TickCount()};
				m_TotalSize.fetch_add(Entry.Location.Size());
			}
			MaxFileOffset = std::max<uint64_t>(MaxFileOffset, Entry.Location.Offset() + Entry.Location.Size());
		});

		if (InvalidEntryCount)
		{
			ZEN_WARN("found {} invalid entries in '{}'", InvalidEntryCount, SlogPath);
		}

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

	for (CbFieldView Entry : Manifest["Timestamps"])
	{
		const CbObjectView Obj = Entry.AsObjectView();
		const IoHash	   Key = Obj["Key"sv].AsHash();

		if (auto It = m_Index.find(Key); It != m_Index.end())
		{
			It.value().LastAccess = Obj["LastAccess"sv].AsInt64();
		}
	}

	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())
	{
		IndexEntry& Entry = It.value();
		Entry.LastAccess  = GcClock::TickCount();

		if (GetInlineCacheValue(Entry.Location, OutValue))
		{
			return true;
		}

		_.ReleaseNow();

		return GetStandaloneCacheValue(Entry.Location, 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, GcClock::TickCount()}});
		}
		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
			IndexEntry& Entry = It.value();
			Entry.Location	  = Loc;
			Entry.LastAccess  = GcClock::TickCount();
		}

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

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

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

void
ZenCacheDiskLayer::CacheBucket::Flush()
{
	using namespace std::literals;

	RwLock::SharedLockScope _(m_IndexLock);

	m_SobsFile.Flush();
	m_SlogFile.Flush();

	// Update manifest
	{
		CbObjectWriter Writer;
		Writer << "BucketId"sv << m_BucketId;

		if (!m_Index.empty())
		{
			Writer.BeginArray("Timestamps"sv);
			for (auto& Kv : m_Index)
			{
				const IoHash&	  Key	= Kv.first;
				const IndexEntry& Entry = Kv.second;
				Writer << "Key"sv << Key << "LastAccess"sv << Entry.LastAccess;
			}
			Writer.EndArray();
		}

		SaveCompactBinaryObject(m_BucketDir / "zen_manifest", Writer.Save());
	}
}

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.Location;

			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

			const auto			It		 = m_Index.find(BadKey);
			const DiskLocation& Location = It->second.Location;
			m_SlogFile.Append(DiskIndexEntry{.Key = BadKey, .Location = {Location.Offset(), Location.Size(), 0, DiskLocation::kTombStone}});
			m_Index.erase(BadKey);
		}
	}
}

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

	std::vector<IoHash> Cids;

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

		if (Loc.IsFlagSet(DiskLocation::kStructured) == false || Loc.IsFlagSet(DiskLocation::kTombStone) ||
			GcCtx.Expired(Kv.second.LastAccess))
		{
			continue;
		}

		ZenCacheValue CacheValue;
		if (!GetInlineCacheValue(Loc, CacheValue))
		{
			GetStandaloneCacheValue(Loc, HashKey, CacheValue);
		}

		if (CacheValue.Value)
		{
			ZEN_ASSERT(CacheValue.Value.GetContentType() == ZenContentType::kCbObject);

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

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);
	IndexEntry	 Entry = IndexEntry{.Location = Loc, .LastAccess = GcClock::TickCount()};

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

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

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

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

		const std::string BucketName8 = ToUtf8(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_INFO("Discovered bucket '{}'", BucketName8);
			}
			else
			{
				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::GatherReferences(GcContext& GcCtx)
{
	RwLock::SharedLockScope _(m_Lock);

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

uint64_t
ZenCacheDiskLayer::TotalSize() const
{
	uint64_t				TotalSize{};
	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		TotalSize += Kv.second.TotalSize();
	}

	return TotalSize;
}

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

#if ZEN_WITH_TESTS

using namespace std::literals;
using namespace fmt::literals;

namespace testutils {

	IoHash CreateKey(size_t KeyValue) { return IoHash::HashBuffer(&KeyValue, sizeof(size_t)); }

}  // namespace testutils

TEST_CASE("zcache.store")
{
	ScopedTemporaryDirectory TempDir;

	CasGc Gc;

	ZenCacheStore Zcs(Gc, 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("zcache.size")
{
	const auto CreateCacheValue = [](size_t Size) -> CbObject {
		std::vector<uint8_t> Buf;
		Buf.resize(Size);

		CbObjectWriter Writer;
		Writer.AddBinary("Binary"sv, Buf.data(), Buf.size());
		return Writer.Save();
	};

	SUBCASE("mem/disklayer")
	{
		const size_t			 Count = 16;
		ScopedTemporaryDirectory TempDir;

		ZenCacheSize CacheSize;

		{
			CasGc		  Gc;
			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");

			CbObject CacheValue = CreateCacheValue(Zcs.DiskLayerThreshold() - 256);

			IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
			Buffer.SetContentType(ZenContentType::kCbObject);

			for (size_t Key = 0; Key < Count; ++Key)
			{
				const size_t Bucket = Key % 4;
				Zcs.Put("test_bucket-{}"_format(Bucket), IoHash::HashBuffer(&Key, sizeof(uint32_t)), {.Value = Buffer});
			}

			CacheSize = Zcs.TotalSize();
			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.DiskSize);
			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.MemorySize);
		}

		{
			CasGc		  Gc;
			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");

			const ZenCacheSize SerializedSize = Zcs.TotalSize();
			CHECK_EQ(SerializedSize.MemorySize, 0);
			CHECK_EQ(SerializedSize.DiskSize, CacheSize.DiskSize);

			for (size_t Bucket = 0; Bucket < 4; ++Bucket)
			{
				Zcs.DropBucket("test_bucket-{}"_format(Bucket));
			}
			CHECK_EQ(0, Zcs.TotalSize().DiskSize);
		}
	}

	SUBCASE("disklayer")
	{
		const size_t			 Count = 16;
		ScopedTemporaryDirectory TempDir;

		ZenCacheSize CacheSize;

		{
			CasGc		  Gc;
			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");

			CbObject CacheValue = CreateCacheValue(Zcs.DiskLayerThreshold() + 64);

			IoBuffer Buffer = CacheValue.GetBuffer().AsIoBuffer();
			Buffer.SetContentType(ZenContentType::kCbObject);

			for (size_t Key = 0; Key < Count; ++Key)
			{
				const size_t Bucket = Key % 4;
				Zcs.Put("test_bucket-{}"_format(Bucket), IoHash::HashBuffer(&Key, sizeof(uint32_t)), {.Value = Buffer});
			}

			CacheSize = Zcs.TotalSize();
			CHECK_EQ(CacheValue.GetSize() * Count, CacheSize.DiskSize);
			CHECK_EQ(0, CacheSize.MemorySize);
		}

		{
			CasGc		  Gc;
			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");

			const ZenCacheSize SerializedSize = Zcs.TotalSize();
			CHECK_EQ(SerializedSize.MemorySize, 0);
			CHECK_EQ(SerializedSize.DiskSize, CacheSize.DiskSize);

			for (size_t Bucket = 0; Bucket < 4; ++Bucket)
			{
				Zcs.DropBucket("test_bucket-{}"_format(Bucket));
			}
			CHECK_EQ(0, Zcs.TotalSize().DiskSize);
		}
	}
}

TEST_CASE("zcache.gc")
{
	using namespace testutils;

	SUBCASE("gather references does NOT add references for expired cache entries")
	{
		ScopedTemporaryDirectory TempDir;
		std::vector<IoHash>		 Cids{CreateKey(1), CreateKey(2), CreateKey(3)};

		const auto CollectAndFilter = [](CasGc&					 Gc,
										 GcClock::TimePoint		 Time,
										 GcClock::Duration		 MaxDuration,
										 std::span<const IoHash> Cids,
										 std::vector<IoHash>&	 OutKeep) {
			GcContext GcCtx(Time);
			GcCtx.MaxCacheDuration(MaxDuration);
			Gc.CollectGarbage(GcCtx);
			OutKeep.clear();
			GcCtx.FilterCids(Cids, [&OutKeep](const IoHash& Hash) { OutKeep.push_back(Hash); });
		};

		{
			CasGc		  Gc;
			ZenCacheStore Zcs(Gc, TempDir.Path() / "cache");
			const auto	  Bucket = "teardrinker"sv;

			// Create a cache record
			const IoHash   Key = CreateKey(42);
			CbObjectWriter Record;
			Record << "Key"sv
				   << "SomeRecord"sv;

			for (size_t Idx = 0; auto& Cid : Cids)
			{
				Record.AddBinaryAttachment("attachment-{}"_format(Idx++), Cid);
			}

			IoBuffer Buffer = Record.Save().GetBuffer().AsIoBuffer();
			Buffer.SetContentType(ZenContentType::kCbObject);

			Zcs.Put(Bucket, Key, {.Value = Buffer});

			std::vector<IoHash> Keep;

			// Collect garbage with 1 hour max cache duration
			{
				CollectAndFilter(Gc, GcClock::Now(), std::chrono::hours(1), Cids, Keep);
				CHECK_EQ(Cids.size(), Keep.size());
			}

			// Move forward in time
			{
				CollectAndFilter(Gc, GcClock::Now() + std::chrono::hours(2), std::chrono::hours(1), Cids, Keep);
				CHECK_EQ(0, Keep.size());
			}
		}

		// Expect timestamps to be serialized
		{
			CasGc				Gc;
			ZenCacheStore		Zcs(Gc, TempDir.Path() / "cache");
			std::vector<IoHash> Keep;

			// Collect garbage with 1 hour max cache duration
			{
				CollectAndFilter(Gc, GcClock::Now(), std::chrono::hours(1), Cids, Keep);
				CHECK_EQ(3, Keep.size());
			}

			// Move forward in time
			{
				CollectAndFilter(Gc, GcClock::Now() + std::chrono::hours(2), std::chrono::hours(1), Cids, Keep);
				CHECK_EQ(0, Keep.size());
			}
		}
	}
}

#endif

void
z$_forcelink()
{
}

}  // namespace zen