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

#include "zenstore/cidstore.h"

#include <zencore/filesystem.h>
#include <zenstore/CAS.h>
#include <zenstore/caslog.h>

#include <spdlog/spdlog.h>
#include <filesystem>

namespace zen {

struct CidStore::CidState
{
	CidState(CasStore& InCasStore) : m_CasStore(InCasStore) {}

	struct IndexEntry
	{
		IoHash Uncompressed;
		IoHash Compressed;
	};

	CasStore&				m_CasStore;
	TCasLogFile<IndexEntry> m_LogFile;

	RwLock						   m_Lock;
	tsl::robin_map<IoHash, IoHash> m_CidMap;

	void AddCompressedCid(const IoHash& DecompressedId, const IoHash& Compressed)
	{
		RwLock::ExclusiveLockScope _(m_Lock);
		m_CidMap.insert_or_assign(DecompressedId, Compressed);
		m_LogFile.Append({.Uncompressed = DecompressedId, .Compressed = Compressed});
	}

	IoBuffer FindChunkByCid(const IoHash& DecompressedId)
	{
		IoHash CompressedHash;
		{
			RwLock::SharedLockScope _(m_Lock);
			if (auto It = m_CidMap.find(DecompressedId); It != m_CidMap.end())
			{
				CompressedHash = It->second;
			}
		}

		if (CompressedHash != IoHash::Zero)
		{
			return m_CasStore.FindChunk(CompressedHash);
		}

		return IoBuffer();
	}

	bool ContainsChunk(const IoHash& DecompressedId)
	{
		RwLock::SharedLockScope _(m_Lock);
		return m_CidMap.find(DecompressedId) != m_CidMap.end();
	}

	void InitializeIndex(const std::filesystem::path& RootDir)
	{
		zen::CreateDirectories(RootDir);
		std::filesystem::path SlogPath{RootDir / "cid.slog"};

		bool IsNew = !std::filesystem::exists(SlogPath);

		m_LogFile.Open(SlogPath, IsNew);

		m_LogFile.Replay([&](const IndexEntry& Ie) { m_CidMap.insert_or_assign(Ie.Uncompressed, Ie.Compressed); });

		spdlog::debug("CID index initialized: {} entries found", m_CidMap.size());
	}

	void Flush() { m_LogFile.Flush(); }
};

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

CidStore::CidStore(CasStore& InCasStore, const std::filesystem::path& RootDir) : m_Impl(std::make_unique<CidState>(InCasStore))
{
	m_Impl->InitializeIndex(RootDir);
}

CidStore::~CidStore()
{
}

void
CidStore::AddCompressedCid(const IoHash& DecompressedId, const IoHash& Compressed)
{
	m_Impl->AddCompressedCid(DecompressedId, Compressed);
}

IoBuffer
CidStore::FindChunkByCid(const IoHash& DecompressedId)
{
	return m_Impl->FindChunkByCid(DecompressedId);
}

bool
CidStore::ContainsChunk(const IoHash& DecompressedId)
{
	return m_Impl->ContainsChunk(DecompressedId);
}

void
CidStore::Flush()
{
	m_Impl->Flush();
}

}  // namespace zen