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

#include "IndexedCacheStorageManager.h"
#include "HAL/PlatformMisc.h"
#include "Misc/ConfigCacheIni.h"
#include "Misc/ScopeRWLock.h"
#include "IndexedCacheStorage.h"

DECLARE_LOG_CATEGORY_EXTERN(LogIndexedCacheStorageManager, Display, All);
DEFINE_LOG_CATEGORY(LogIndexedCacheStorageManager);

namespace Experimental
{

FIndexedCacheStorageManager& FIndexedCacheStorageManager::Get()
{
	static FIndexedCacheStorageManager GIndexedCacheStorageManager;
	UE_CALL_ONCE([&]
	{
		GIndexedCacheStorageManager.Initialize();
	}
	);
	
	return GIndexedCacheStorageManager;
}

void FIndexedCacheStorageManager::Initialize()
{
	uint32 NumIndexedCacheEntries = IIndexedCacheStorage::Get().GetNumIndexedCacheStorages();
	IndexedCacheEntriesMetaData.SetNum(NumIndexedCacheEntries);

	// Takes on the pattern
	// (Name="CacheStorageName",Index=CacheStorageIndex,Mount="CacheStorageMount")
	TArray<FString> StorageConfigs;
	GConfig->GetArray(TEXT("IndexedCacheStorage"), TEXT("Storage"), StorageConfigs, GEngineIni);

	for (const FString& Category : StorageConfigs)
	{
		FString TrimmedCategory = Category;
		TrimmedCategory.TrimStartAndEndInline();
		if (TrimmedCategory.Left(1) == TEXT("("))
		{
			TrimmedCategory.RightChopInline(1, EAllowShrinking::No);
		}
		if (TrimmedCategory.Right(1) == TEXT(")"))
		{
			TrimmedCategory.LeftChopInline(1, EAllowShrinking::No);
		}

		// Find all custom chunks and parse
		const TCHAR* PropertyName = TEXT("Name=");
		const TCHAR* PropertyIndex = TEXT("Index=");
		const TCHAR* PropertyMount = TEXT("Mount=");
		FString StorageName;
		int32 StorageIndex;
		FString StorageMount;

		if (FParse::Value(*TrimmedCategory, PropertyName, StorageName) &&
			FParse::Value(*TrimmedCategory, PropertyIndex, StorageIndex) &&
			FParse::Value(*TrimmedCategory, PropertyMount, StorageMount))
		{
			if (StorageName.Len() == 0)
			{
				UE_LOG(LogIndexedCacheStorageManager, Error, TEXT("Found empty Name in [IndexedCacheStorage]:Storage"));
				continue;
			}

			if (StorageIndex <= 0 || StorageIndex >= (int32)NumIndexedCacheEntries)
			{
				UE_LOG(LogIndexedCacheStorageManager, Error, TEXT("Found invalid Index in [IndexedCacheStorage]:Storage: %d, max allowed = %d"), StorageIndex, NumIndexedCacheEntries-1);
				continue;
			}

			if (StorageMount.Len() == 0)
			{
				UE_LOG(LogIndexedCacheStorageManager, Error, TEXT("Found empty Mount in [IndexedCacheStorage]:Storage"));
				continue;
			}

			StorageName.ReplaceInline(TEXT("\""), TEXT(""));
			StorageMount.ReplaceInline(TEXT("\""), TEXT(""));

			int32& ExistingStorageIndex = StorageNameToStorageIndex.FindOrAdd(StorageName, 0);
			if (ExistingStorageIndex > 0)
			{
				UE_LOG(LogIndexedCacheStorageManager, Error, TEXT("Found an existing entry with name %s in [IndexedCacheStorage]:Storage"), *StorageName);
				continue;
			}
			
			FCacheStorageMetaData& CacheStorageMetaData = IndexedCacheEntriesMetaData[StorageIndex];

			if (!CacheStorageMetaData.MountName.IsEmpty())
			{
				UE_LOG(LogIndexedCacheStorageManager, Error, TEXT("Found an existing entry with at index %d for name %s in [IndexedCacheStorage]:Storage"), StorageIndex, *StorageName);
				continue;
			}

			ExistingStorageIndex = StorageIndex;
			CacheStorageMetaData.MountName = StorageMount;
		}
	}

	TArray<int32> CacheStorageIndices;
	IIndexedCacheStorage::Get().EnumerateCacheStorages(CacheStorageIndices);
	TSet<int32> CacheStorageIndicesSet;
	for (int32 CacheStorageIndex : CacheStorageIndices)
	{
		IndexedCacheEntriesMetaData[CacheStorageIndex].bCacheExists = true;
	}
	bHasRegisteredEntries = (CacheStorageIndices.Num() > 0);

}

bool FIndexedCacheStorageManager::SupportsIndexedCacheStorage()
{
	return bHasRegisteredEntries && IIndexedCacheStorage::Get().SupportsIndexedCacheStorage();
}

int32 FIndexedCacheStorageManager::GetStorageIndex(const FString& StorageName)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_ReadOnly);

	int32* StorageIndex = StorageNameToStorageIndex.Find(StorageName);

	return StorageIndex ? *StorageIndex : 0;
}

void FIndexedCacheStorageManager::EnumerateCacheStorages(TArray<FString>& OutCacheStorageNames)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_ReadOnly);

	OutCacheStorageNames.Reserve(StorageNameToStorageIndex.Num());
	for (const TPair<FString, int32>& StorageNameToStorageIndexIter : StorageNameToStorageIndex)
	{
		OutCacheStorageNames.Add(StorageNameToStorageIndexIter.Key);
	}
}


bool FIndexedCacheStorageManager::CreateCacheStorage(uint64 RequestNumberOfBytes, int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_Write);
	if (CacheIndex <= 0 && !ensure(IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex)))
	{
		return false;
	}

	if (IndexedCacheEntriesMetaData[CacheIndex].MountName.IsEmpty())
	{
		return false;
	}

	ensure(IndexedCacheEntriesMetaData[CacheIndex].MountRefCount == 0);
	IndexedCacheEntriesMetaData[CacheIndex].bCacheExists = IIndexedCacheStorage::Get().CreateCacheStorage(RequestNumberOfBytes, CacheIndex);
	return IndexedCacheEntriesMetaData[CacheIndex].bCacheExists;
}

void FIndexedCacheStorageManager::DestroyCacheStorage(int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_Write);

	if (CacheIndex <= 0 && !ensure(IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex)))
	{
		return;
	}

	if (IndexedCacheEntriesMetaData[CacheIndex].MountName.IsEmpty() || !IndexedCacheEntriesMetaData[CacheIndex].bCacheExists)
	{
		return;
	}

	ensure(IndexedCacheEntriesMetaData[CacheIndex].MountRefCount == 0);
	IIndexedCacheStorage::Get().DestroyCacheStorage(CacheIndex);
	IndexedCacheEntriesMetaData[CacheIndex].bCacheExists = false;
}

uint64 FIndexedCacheStorageManager::GetCacheStorageCapacity(int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_ReadOnly);

	if (CacheIndex <= 0 && !ensure(IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex)))
	{
		return 0;
	}

	if (IndexedCacheEntriesMetaData[CacheIndex].MountName.IsEmpty() || !IndexedCacheEntriesMetaData[CacheIndex].bCacheExists)
	{
		return 0;
	}

	uint64 ExistingSpace = 0;
	IIndexedCacheStorage::Get().GetCacheStorageInfo(CacheIndex, ExistingSpace);
	return ExistingSpace;
}

FString FIndexedCacheStorageManager::MountCacheStorage(int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_Write);

	FString ActualMountName;
	if (!IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex) || !IndexedCacheEntriesMetaData[CacheIndex].bCacheExists)
	{
		return ActualMountName;
	}

	check(IndexedCacheEntriesMetaData[CacheIndex].MountRefCount >= 0);

	FString MountedPath;
	if (IndexedCacheEntriesMetaData[CacheIndex].MountRefCount == 0)
	{
		bool bMounted = IIndexedCacheStorage::Get().MountCacheStorage(MountedPath, CacheIndex, IndexedCacheEntriesMetaData[CacheIndex].MountName);
		if (!bMounted)
		{
			return ActualMountName;
		}
	}
	else
	{
		IIndexedCacheStorage::Get().GetCacheStorageMountPath(MountedPath, IndexedCacheEntriesMetaData[CacheIndex].MountName);
	}

	check(!MountedPath.IsEmpty());

	++IndexedCacheEntriesMetaData[CacheIndex].MountRefCount;

	return MountedPath;
}

void FIndexedCacheStorageManager::UnmountCacheStorage(int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_Write);

	FString ActualMountName;
	if (!IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex) || IndexedCacheEntriesMetaData[CacheIndex].MountRefCount == 0)
	{
		return;
	}

	--IndexedCacheEntriesMetaData[CacheIndex].MountRefCount;
	if (IndexedCacheEntriesMetaData[CacheIndex].MountRefCount == 0)
	{
		IIndexedCacheStorage::Get().UnmountCacheStorage(*IndexedCacheEntriesMetaData[CacheIndex].MountName);
	}
}


FString FIndexedCacheStorageManager::GetMountName(int32 CacheIndex)
{
	FRWScopeLock _(IndexedCacheEntriesMetaDataLock, SLT_ReadOnly);
	if (!IndexedCacheEntriesMetaData.IsValidIndex(CacheIndex))
	{
		return FString();
	}

	return IndexedCacheEntriesMetaData[CacheIndex].MountName;

}

FString FIndexedCacheStorageManager::GetMountPath(int32 CacheIndex)
{
	FString MountName = GetMountName(CacheIndex);
	FString BundleMountPath; 
	IIndexedCacheStorage::Get().GetCacheStorageMountPath(BundleMountPath, MountName);
	return BundleMountPath;
}

}