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

#include "cachememorylayer.h"

#include <zencore/compactbinaryvalidation.h>
#include <zencore/compress.h>
#include <zencore/trace.h>

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

namespace zen {

ZenCacheMemoryLayer::ZenCacheMemoryLayer()
{
}

ZenCacheMemoryLayer::~ZenCacheMemoryLayer()
{
}

bool
ZenCacheMemoryLayer::Get(std::string_view InBucket, const IoHash& HashKey, ZenCacheValue& OutValue)
{
	ZEN_TRACE_CPU("Z$::Mem::Get");

	RwLock::SharedLockScope _(m_Lock);

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

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

	CacheBucket* Bucket = It->second.get();

	_.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)
{
	ZEN_TRACE_CPU("Z$::Mem::Put");

	const auto	 BucketName = std::string(InBucket);
	CacheBucket* Bucket		= nullptr;

	{
		RwLock::SharedLockScope _(m_Lock);

		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
		{
			Bucket = It->second.get();
		}
	}

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

		RwLock::ExclusiveLockScope _(m_Lock);

		if (auto It = m_Buckets.find(std::string(InBucket)); It != m_Buckets.end())
		{
			Bucket = It->second.get();
		}
		else
		{
			auto InsertResult = m_Buckets.emplace(BucketName, std::make_unique<CacheBucket>());
			Bucket			  = InsertResult.first->second.get();
		}
	}

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

	Bucket->Put(HashKey, Value);
}

bool
ZenCacheMemoryLayer::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;
		m_DroppedBuckets.push_back(std::move(It->second));
		m_Buckets.erase(It);
		Bucket.Drop();
		return true;
	}
	return false;
}

void
ZenCacheMemoryLayer::Drop()
{
	RwLock::ExclusiveLockScope				  _(m_Lock);
	std::vector<std::unique_ptr<CacheBucket>> Buckets;
	Buckets.reserve(m_Buckets.size());
	while (!m_Buckets.empty())
	{
		const auto&	 It		= m_Buckets.begin();
		CacheBucket& Bucket = *It->second;
		m_DroppedBuckets.push_back(std::move(It->second));
		m_Buckets.erase(It->first);
		Bucket.Drop();
	}
}

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

	for (auto& Kv : m_Buckets)
	{
		Kv.second->ScrubStorage(Ctx);
	}
}

void
ZenCacheMemoryLayer::GatherAccessTimes(zen::access_tracking::AccessTimes& AccessTimes)
{
	using namespace zen::access_tracking;

	RwLock::SharedLockScope _(m_Lock);

	for (auto& Kv : m_Buckets)
	{
		std::vector<KeyAccessTime>& Bucket = AccessTimes.Buckets[Kv.first];
		Kv.second->GatherAccessTimes(Bucket);
	}
}

void
ZenCacheMemoryLayer::Reset()
{
	RwLock::ExclusiveLockScope _(m_Lock);
	m_Buckets.clear();
}

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

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

	return TotalSize;
}

ZenCacheMemoryLayer::Info
ZenCacheMemoryLayer::GetInfo() const
{
	ZenCacheMemoryLayer::Info Info = {.Config = m_Configuration, .TotalSize = TotalSize()};

	RwLock::SharedLockScope _(m_Lock);
	Info.BucketNames.reserve(m_Buckets.size());
	for (auto& Kv : m_Buckets)
	{
		Info.BucketNames.push_back(Kv.first);
		Info.EntryCount += Kv.second->EntryCount();
	}
	return Info;
}

std::optional<ZenCacheMemoryLayer::BucketInfo>
ZenCacheMemoryLayer::GetBucketInfo(std::string_view Bucket) const
{
	RwLock::SharedLockScope _(m_Lock);

	if (auto It = m_Buckets.find(std::string(Bucket)); It != m_Buckets.end())
	{
		return ZenCacheMemoryLayer::BucketInfo{.EntryCount = It->second->EntryCount(), .TotalSize = It->second->TotalSize()};
	}
	return {};
}

void
ZenCacheMemoryLayer::CacheBucket::ScrubStorage(ScrubContext& Ctx)
{
	RwLock::SharedLockScope _(m_BucketLock);

	std::vector<IoHash> BadHashes;

	auto ValidateEntry = [](const IoHash& Hash, ZenContentType ContentType, IoBuffer Buffer) {
		if (ContentType == ZenContentType::kCbObject)
		{
			CbValidateError Error = ValidateCompactBinary(Buffer, CbValidateMode::All);
			return Error == CbValidateError::None;
		}
		if (ContentType == ZenContentType::kCompressedBinary)
		{
			IoHash	 RawHash;
			uint64_t RawSize;
			if (!CompressedBuffer::ValidateCompressedHeader(Buffer, RawHash, RawSize))
			{
				return false;
			}
			if (Hash != RawHash)
			{
				return false;
			}
		}
		return true;
	};

	for (auto& Kv : m_CacheMap)
	{
		const BucketPayload& Payload = m_Payloads[Kv.second];
		if (!ValidateEntry(Kv.first, Payload.Payload.GetContentType(), Payload.Payload))
		{
			BadHashes.push_back(Kv.first);
		}
	}

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

void
ZenCacheMemoryLayer::CacheBucket::GatherAccessTimes(std::vector<zen::access_tracking::KeyAccessTime>& AccessTimes)
{
	RwLock::SharedLockScope _(m_BucketLock);
	std::transform(m_CacheMap.begin(), m_CacheMap.end(), std::back_inserter(AccessTimes), [this](const auto& Kv) {
		return access_tracking::KeyAccessTime{.Key = Kv.first, .LastAccess = m_AccessTimes[Kv.second]};
	});
}

bool
ZenCacheMemoryLayer::CacheBucket::Get(const IoHash& HashKey, ZenCacheValue& OutValue)
{
	ZEN_TRACE_CPU("Z$::Mem::Bucket::Get");

	RwLock::SharedLockScope _(m_BucketLock);

	if (auto It = m_CacheMap.find(HashKey); It != m_CacheMap.end())
	{
		uint32_t EntryIndex = It.value();
		ZEN_ASSERT_SLOW(EntryIndex < m_Payloads.size());
		ZEN_ASSERT_SLOW(m_AccessTimes.size() == m_Payloads.size());

		const BucketPayload& Payload = m_Payloads[EntryIndex];
		OutValue					 = {.Value = Payload.Payload, .RawSize = Payload.RawSize, .RawHash = Payload.RawHash};
		m_AccessTimes[EntryIndex]	 = GcClock::TickCount();

		return true;
	}

	return false;
}

void
ZenCacheMemoryLayer::CacheBucket::Put(const IoHash& HashKey, const ZenCacheValue& Value)
{
	ZEN_TRACE_CPU("Z$::Mem::Bucket::Put");

	size_t PayloadSize = Value.Value.GetSize();
	{
		GcClock::Tick			   AccessTime = GcClock::TickCount();
		RwLock::ExclusiveLockScope _(m_BucketLock);
		if (m_CacheMap.size() == std::numeric_limits<uint32_t>::max())
		{
			// No more space in our memory cache!
			return;
		}
		if (auto It = m_CacheMap.find(HashKey); It != m_CacheMap.end())
		{
			uint32_t EntryIndex = It.value();
			ZEN_ASSERT_SLOW(EntryIndex < m_Payloads.size());

			m_TotalSize.fetch_sub(PayloadSize, std::memory_order::relaxed);
			BucketPayload& Payload	  = m_Payloads[EntryIndex];
			Payload.Payload			  = Value.Value;
			Payload.RawHash			  = Value.RawHash;
			Payload.RawSize			  = gsl::narrow<uint32_t>(Value.RawSize);
			m_AccessTimes[EntryIndex] = AccessTime;
		}
		else
		{
			uint32_t EntryIndex = gsl::narrow<uint32_t>(m_Payloads.size());
			m_Payloads.emplace_back(
				BucketPayload{.Payload = Value.Value, .RawSize = gsl::narrow<uint32_t>(Value.RawSize), .RawHash = Value.RawHash});
			m_AccessTimes.emplace_back(AccessTime);
			m_CacheMap.insert_or_assign(HashKey, EntryIndex);
		}
		ZEN_ASSERT_SLOW(m_Payloads.size() == m_CacheMap.size());
		ZEN_ASSERT_SLOW(m_AccessTimes.size() == m_Payloads.size());
	}

	m_TotalSize.fetch_add(PayloadSize, std::memory_order::relaxed);
}

void
ZenCacheMemoryLayer::CacheBucket::Drop()
{
	RwLock::ExclusiveLockScope _(m_BucketLock);
	m_CacheMap.clear();
	m_AccessTimes.clear();
	m_Payloads.clear();
	m_TotalSize.store(0);
}

uint64_t
ZenCacheMemoryLayer::CacheBucket::EntryCount() const
{
	RwLock::SharedLockScope _(m_BucketLock);
	return static_cast<uint64_t>(m_CacheMap.size());
}

}  // namespace zen