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

#include <zenutil/bufferedwritefilecache.h>

#include <zencore/logging.h>
#include <zencore/trace.h>

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

namespace zen {

BufferedWriteFileCache::BufferedWriteFileCache() : m_CacheHitCount(0), m_CacheMissCount(0), m_OpenHandleCount(0), m_DroppedHandleCount(0)
{
}

BufferedWriteFileCache::~BufferedWriteFileCache()
{
	ZEN_TRACE_CPU("~BufferedWriteFileCache()");

	try
	{
		for (TOpenHandles& OpenHandles : m_OpenFiles)
		{
			while (BasicFile* File = OpenHandles.Pop())
			{
				std::unique_ptr<BasicFile> FileToClose(File);
				m_OpenHandleCount--;
			}
		}
		m_OpenFiles.clear();
		m_ChunkWriters.clear();
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("~BufferedWriteFileCache() threw exeption: {}", Ex.what());
	}
}

std::unique_ptr<BasicFile>
BufferedWriteFileCache::Get(uint32_t FileIndex)
{
	ZEN_TRACE_CPU("BufferedWriteFileCache::Get");

	RwLock::ExclusiveLockScope _(m_WriterLock);
	if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
	{
		const uint32_t HandleIndex = It->second;
		TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
		if (BasicFile* File = OpenHandles.Pop(); File != nullptr)
		{
			m_OpenHandleCount--;
			m_CacheHitCount++;
			return std::unique_ptr<BasicFile>(File);
		}
	}
	m_CacheMissCount++;
	return nullptr;
}

void
BufferedWriteFileCache::Put(uint32_t FileIndex, std::unique_ptr<BasicFile>&& Writer)
{
	ZEN_TRACE_CPU("BufferedWriteFileCache::Put");

	if (m_OpenHandleCount.load() >= MaxBufferedCount)
	{
		m_DroppedHandleCount++;
		return;
	}
	RwLock::ExclusiveLockScope _(m_WriterLock);
	if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
	{
		const uint32_t HandleIndex = It->second;
		TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
		if (OpenHandles.Push(Writer.get()))
		{
			Writer.release();
			m_OpenHandleCount++;
		}
		else
		{
			m_DroppedHandleCount++;
		}
	}
	else
	{
		const uint32_t HandleIndex = gsl::narrow<uint32_t>(m_OpenFiles.size());
		m_OpenFiles.push_back(TOpenHandles{});
		m_OpenFiles.back().Push(Writer.release());
		m_ChunkWriters.insert_or_assign(FileIndex, HandleIndex);
		m_OpenHandleCount++;
	}
}

void
BufferedWriteFileCache::Close(std::span<uint32_t> FileIndexes)
{
	ZEN_TRACE_CPU("BufferedWriteFileCache::Close");

	std::vector<std::unique_ptr<BasicFile>> FilesToClose;
	FilesToClose.reserve(FileIndexes.size());
	{
		RwLock::ExclusiveLockScope _(m_WriterLock);
		for (uint32_t FileIndex : FileIndexes)
		{
			if (auto It = m_ChunkWriters.find(FileIndex); It != m_ChunkWriters.end())
			{
				const uint32_t HandleIndex = It->second;
				TOpenHandles&  OpenHandles = m_OpenFiles[HandleIndex];
				while (BasicFile* File = OpenHandles.Pop())
				{
					FilesToClose.emplace_back(std::unique_ptr<BasicFile>(File));
					m_OpenHandleCount--;
				}
				m_ChunkWriters.erase(It);
			}
		}
	}
	FilesToClose.clear();
}

BufferedWriteFileCache::Local::Local(BufferedWriteFileCache& Cache) : m_Cache(Cache)
{
}

BufferedWriteFileCache::Local::Writer*
BufferedWriteFileCache::Local::GetWriter(uint32_t FileIndex)
{
	if (auto It = m_FileIndexToWriterIndex.find(FileIndex); It != m_FileIndexToWriterIndex.end())
	{
		return m_ChunkWriters[It->second].get();
	}
	std::unique_ptr<BasicFile> File = m_Cache.Get(FileIndex);
	if (File)
	{
		const uint32_t WriterIndex = gsl::narrow<uint32_t>(m_ChunkWriters.size());
		m_FileIndexToWriterIndex.insert_or_assign(FileIndex, WriterIndex);
		m_ChunkWriters.emplace_back(std::make_unique<Writer>(Writer{.File = std::move(File)}));
		return m_ChunkWriters.back().get();
	}
	return nullptr;
}

BufferedWriteFileCache::Local::Writer*
BufferedWriteFileCache::Local::PutWriter(uint32_t FileIndex, std::unique_ptr<Writer> Writer)
{
	ZEN_ASSERT(!m_FileIndexToWriterIndex.contains(FileIndex));
	const uint32_t WriterIndex = gsl::narrow<uint32_t>(m_ChunkWriters.size());
	m_FileIndexToWriterIndex.insert_or_assign(FileIndex, WriterIndex);
	m_ChunkWriters.emplace_back(std::move(Writer));
	return m_ChunkWriters.back().get();
}

BufferedWriteFileCache::Local::~Local()
{
	ZEN_TRACE_CPU("BufferedWriteFileCache::~Local()");
	try
	{
		for (auto& It : m_FileIndexToWriterIndex)
		{
			const uint32_t FileIndex   = It.first;
			const uint32_t WriterIndex = It.second;
			m_ChunkWriters[WriterIndex]->Writer.reset();
			std::unique_ptr<BasicFile> File;
			File.swap(m_ChunkWriters[WriterIndex]->File);
			m_Cache.Put(FileIndex, std::move(File));
		}
	}
	catch (const std::exception& Ex)
	{
		ZEN_ERROR("BufferedWriteFileCache::~Local() threw exeption: {}", Ex.what());
	}
}

}  // namespace zen