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

#include "zenstore/cidstore.h"

#include <zencore/compress.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/string.h>
#include <zenstore/scrubcontext.h>

#include "cas.h"

#include <filesystem>

namespace zen {

struct CidStore::Impl
{
	explicit Impl(CasStore& InCasStore) : m_CasStore(InCasStore) {}

	CasStore& m_CasStore;

	void Initialize(const CidStoreConfiguration& Config) { m_CasStore.Initialize(Config); }

	CidStore::InsertResult AddChunk(const IoBuffer& ChunkData, const IoHash& RawHash, CidStore::InsertMode Mode)
	{
#if ZEN_BUILD_DEBUG
		IoHash	 VerifyRawHash;
		uint64_t _;
		ZEN_ASSERT(CompressedBuffer::ValidateCompressedHeader(ChunkData, VerifyRawHash, _) && RawHash == VerifyRawHash);
#endif

		IoBuffer Payload(ChunkData);
		Payload.SetContentType(ZenContentType::kCompressedBinary);

		CasStore::InsertResult Result = m_CasStore.InsertChunk(Payload, RawHash, static_cast<CasStore::InsertMode>(Mode));
		if (Result.New)
		{
			WriteCount++;
		}
		return {.New = Result.New};
	}

	IoBuffer FindChunkByCid(const IoHash& DecompressedId)
	{
		IoBuffer Result = m_CasStore.FindChunk(DecompressedId);
		if (Result)
		{
			HitCount++;
		}
		else
		{
			MissCount++;
		}
		return Result;
	}

	bool ContainsChunk(const IoHash& DecompressedId) { return m_CasStore.ContainsChunk(DecompressedId); }

	void FilterChunks(HashKeySet& InOutChunks)
	{
		InOutChunks.RemoveHashesIf([&](const IoHash& Hash) { return ContainsChunk(Hash); });
	}

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

	void ScrubStorage(ScrubContext& Ctx)
	{
		if (Ctx.ScrubTimestamp() == m_LastScrubTime)
		{
			return;
		}

		m_LastScrubTime = Ctx.ScrubTimestamp();

		m_CasStore.ScrubStorage(Ctx);
	}

	CidStoreStats Stats() { return CidStoreStats{.HitCount = HitCount, .MissCount = MissCount, .WriteCount = WriteCount}; }

	std::atomic_uint64_t HitCount{};
	std::atomic_uint64_t MissCount{};
	std::atomic_uint64_t WriteCount{};

	uint64_t m_LastScrubTime = 0;
};

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

CidStore::CidStore(GcManager& Gc) : m_CasStore(CreateCasStore(Gc)), m_Impl(std::make_unique<Impl>(*m_CasStore))
{
}

CidStore::~CidStore()
{
}

void
CidStore::Initialize(const CidStoreConfiguration& Config)
{
	m_Impl->Initialize(Config);
}

CidStore::InsertResult
CidStore::AddChunk(const IoBuffer& ChunkData, const IoHash& RawHash, InsertMode Mode)
{
	return m_Impl->AddChunk(ChunkData, RawHash, Mode);
}

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

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

void
CidStore::FilterChunks(HashKeySet& InOutChunks)
{
	return m_Impl->FilterChunks(InOutChunks);
}

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

void
CidStore::ScrubStorage(ScrubContext& Ctx)
{
	m_Impl->ScrubStorage(Ctx);
}

CidStoreSize
CidStore::TotalSize() const
{
	return m_Impl->m_CasStore.TotalSize();
}

CidStoreStats
CidStore::Stats() const
{
	return m_Impl->Stats();
}

}  // namespace zen