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

#include <zenstore/cas.h>

#include "compactcas.h"
#include "filecas.h"

#include <zencore/except.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/memory.h>
#include <zencore/string.h>
#include <zencore/testing.h>
#include <zencore/testutils.h>
#include <zencore/thread.h>
#include <zencore/uid.h>

#include <gsl/gsl-lite.hpp>

#include <filesystem>
#include <functional>
#include <unordered_map>

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

namespace zen {

void
ScrubContext::ReportBadChunks(std::span<IoHash> BadChunks)
{
	ZEN_UNUSED(BadChunks);
}

/**
 * CAS store implementation
 *
 * Uses a basic strategy of splitting payloads by size, to improve ability to reclaim space
 * quickly for unused large chunks and to maintain locality for small chunks which are
 * frequently accessed together.
 *
 */
class CasImpl : public CasStore
{
public:
	CasImpl();
	virtual ~CasImpl();

	virtual void				   Initialize(const CasStoreConfiguration& InConfig) override;
	virtual CasStore::InsertResult InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash) override;
	virtual IoBuffer			   FindChunk(const IoHash& ChunkHash) override;
	virtual void				   FilterChunks(CasChunkSet& InOutChunks) override;
	virtual void				   Flush() override;
	virtual void				   Scrub(ScrubContext& Ctx) override;

private:
	CasContainerStrategy m_TinyStrategy;
	CasContainerStrategy m_SmallStrategy;
	FileCasStrategy		 m_LargeStrategy;
};

CasImpl::CasImpl() : m_TinyStrategy(m_Config), m_SmallStrategy(m_Config), m_LargeStrategy(m_Config)
{
}

CasImpl::~CasImpl()
{
}

void
CasImpl::Initialize(const CasStoreConfiguration& InConfig)
{
	m_Config = InConfig;

	ZEN_INFO("initializing CAS pool at '{}'", m_Config.RootDirectory);

	// Ensure root directory exists - create if it doesn't exist already

	std::filesystem::create_directories(m_Config.RootDirectory);

	// Open or create manifest
	//
	//  The manifest is not currently fully implemented. The goal is to
	//  use it for recovery and configuration

	bool IsNewStore = false;

	{
		std::filesystem::path ManifestPath = m_Config.RootDirectory;
		ManifestPath /= ".ucas_root";

		std::error_code Ec;
		BasicFile		Marker;
		Marker.Open(ManifestPath.c_str(), /* IsCreate */ false, Ec);

		if (Ec)
		{
			IsNewStore = true;

			ExtendableStringBuilder<128> manifest;
			manifest.Append("#CAS_ROOT\n");
			manifest.Append("ID=");
			zen::Oid id = zen::Oid::NewOid();
			id.ToString(manifest);

			Marker.Open(ManifestPath.c_str(), /* IsCreate */ true);
			Marker.Write(manifest.c_str(), (DWORD)manifest.Size(), 0);
		}
	}

	// Initialize payload storage

	m_TinyStrategy.Initialize("tobs", 16, IsNewStore);
	m_SmallStrategy.Initialize("sobs", 4096, IsNewStore);

	ScrubContext Ctx;
	Scrub(Ctx);
}

CasStore::InsertResult
CasImpl::InsertChunk(IoBuffer Chunk, const IoHash& ChunkHash)
{
	const uint64_t ChunkSize = Chunk.Size();

	if (ChunkSize < m_Config.TinyValueThreshold)
	{
		ZEN_ASSERT(ChunkSize);

		return m_TinyStrategy.InsertChunk(Chunk, ChunkHash);
	}
	else if (ChunkSize < m_Config.HugeValueThreshold)
	{
		return m_SmallStrategy.InsertChunk(Chunk, ChunkHash);
	}

	return m_LargeStrategy.InsertChunk(Chunk, ChunkHash);
}

IoBuffer
CasImpl::FindChunk(const IoHash& ChunkHash)
{
	if (IoBuffer Found = m_SmallStrategy.FindChunk(ChunkHash))
	{
		return Found;
	}

	if (IoBuffer Found = m_TinyStrategy.FindChunk(ChunkHash))
	{
		return Found;
	}

	if (IoBuffer Found = m_LargeStrategy.FindChunk(ChunkHash))
	{
		return Found;
	}

	// Not found
	return IoBuffer{};
}

void
CasImpl::FilterChunks(CasChunkSet& InOutChunks)
{
	m_SmallStrategy.FilterChunks(InOutChunks);
	m_TinyStrategy.FilterChunks(InOutChunks);
	m_LargeStrategy.FilterChunks(InOutChunks);
}

void
CasImpl::Flush()
{
	m_SmallStrategy.Flush();
	m_TinyStrategy.Flush();
	m_LargeStrategy.Flush();
}

void
CasImpl::Scrub(ScrubContext& Ctx)
{
	m_SmallStrategy.Scrub(Ctx);
	m_TinyStrategy.Scrub(Ctx);
	m_LargeStrategy.Scrub(Ctx);
}

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

CasStore*
CreateCasStore()
{
	return new CasImpl();
}

//////////////////////////////////////////////////////////////////////////
//
// Testing related code follows...
//

#if ZEN_WITH_TESTS

TEST_CASE("CasStore")
{
	ScopedTemporaryDirectory TempDir;

	zen::CasStoreConfiguration config;
	config.RootDirectory = TempDir.Path();

	std::unique_ptr<zen::CasStore> Store{CreateCasStore()};
	Store->Initialize(config);

	ScrubContext Ctx;
	Store->Scrub(Ctx);

	IoBuffer Value1{16};
	memcpy(Value1.MutableData(), "1234567890123456", 16);
	IoHash				   Hash1   = IoHash::HashBuffer(Value1.Data(), Value1.Size());
	CasStore::InsertResult Result1 = Store->InsertChunk(Value1, Hash1);
	CHECK(Result1.New);

	IoBuffer Value2{16};
	memcpy(Value2.MutableData(), "ABCDEFGHIJKLMNOP", 16);
	IoHash				   Hash2   = IoHash::HashBuffer(Value2.Data(), Value2.Size());
	CasStore::InsertResult Result2 = Store->InsertChunk(Value2, Hash2);
	CHECK(Result2.New);

	CasChunkSet ChunkSet;
	ChunkSet.AddChunk(Hash1);
	ChunkSet.AddChunk(Hash2);

	Store->FilterChunks(ChunkSet);
	CHECK(ChunkSet.GetChunkSet().size() == 0);

	IoBuffer Lookup1 = Store->FindChunk(Hash1);
	CHECK(Lookup1);
	IoBuffer Lookup2 = Store->FindChunk(Hash2);
	CHECK(Lookup2);
}

void
CAS_forcelink()
{
}

#endif

}  // namespace zen