// © 2024 Long Ly. All rights reserved. Any unauthorized use, reproduction, or distribution of this trademark is strictly prohibited and may result in legal action.


#include "Systems/OLSAssetManager.h"
#include "Misc/App.h"
#include "Stats/StatsMisc.h"
#include "Misc/ScopedSlowTask.h"
#include "Engine/Engine.h"
#include "Systems/OLSAssetManagerStartupJob.h"
#include "DataAssets/OLSPawnDataAsset.h"

#include UE_INLINE_GENERATED_CPP_BY_NAME(OLSAssetManager)

const FName FOLSBundles::Equipped("Equipped");

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

static FAutoConsoleCommand CVarDumpLoadedAssets(
	TEXT("OLS.DumpLoadedAssets"),
	TEXT("Shows all assets that were loaded via the asset manager and are currently in memory."),
	FConsoleCommandDelegate::CreateStatic(UOLSAssetManager::DumpLoadedAssets)
);

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

#define STARTUP_JOB_WEIGHTED(jobFunc, jobWeight) StartupJobs.Add(FOLSAssetManagerStartupJob(#jobFunc, [this](const FOLSAssetManagerStartupJob& startupJob, TSharedPtr<FStreamableHandle>& loadHandle){jobFunc;}, jobWeight))
#define STARTUP_JOB(jobFunc) STARTUP_JOB_WEIGHTED(jobFunc, 1.f)

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

UOLSAssetManager::UOLSAssetManager()
{
	DefaultPawnData = nullptr;
}

UOLSAssetManager& UOLSAssetManager::Get()
{
	check(GEngine);

	if (UOLSAssetManager* Singleton = Cast<UOLSAssetManager>(GEngine->AssetManager))
	{
		return *Singleton;
	}

	// @Todo replace this our custom log.
	// UE_LOG(LogLyra, Fatal, TEXT("Invalid AssetManagerClassName in DefaultEngine.ini.  It must be set to LyraAssetManager!"));

	// Fatal error above prevents this from being called.
	return *NewObject<UOLSAssetManager>();
}

template <typename AssetType>
AssetType* UOLSAssetManager::GetAsset(const TSoftObjectPtr<AssetType>& assetPointer, bool shouldKeepInMemory)
{
	AssetType* loadedAsset = nullptr;

	const FSoftObjectPath& assetPath = assetPointer.ToSoftObjectPath();

	if (assetPath.IsValid())
	{
		loadedAsset = assetPointer.Get();
		if (!loadedAsset)
		{
			loadedAsset = Cast<AssetType>(SynchronousLoadAsset(assetPath));
			ensureAlwaysMsgf(loadedAsset, TEXT("Failed to load asset [%s]"), *assetPointer.ToString());
		}

		if (loadedAsset && shouldKeepInMemory)
		{
			// Added to loaded asset list.
			Get().AddLoadedAsset(Cast<UObject>(loadedAsset));
		}
	}

	return loadedAsset;
}

template <typename AssetType>
TSubclassOf<AssetType> UOLSAssetManager::GetSubclass(const TSoftClassPtr<AssetType>& assetPointer, bool shouldKeepInMemory)
{
	TSubclassOf<AssetType> loadedSubclass = nullptr;
    
    	const FSoftObjectPath& assetPath = assetPointer.ToSoftObjectPath();
    
    	if (assetPath.IsValid())
    	{
    		loadedSubclass = assetPointer.Get();
    		if (!loadedSubclass)
    		{
    			loadedSubclass = Cast<UClass>(SynchronousLoadAsset(assetPath));
    			ensureAlwaysMsgf(loadedSubclass, TEXT("Failed to load asset class [%s]"), *assetPointer.ToString());
    		}
    
    		if (loadedSubclass && shouldKeepInMemory)
    		{
    			// Added to loaded asset list.
    			Get().AddLoadedAsset(Cast<UObject>(loadedSubclass));
    		}
    	}
    
    	return loadedSubclass;
}

void UOLSAssetManager::DumpLoadedAssets()
{
}

const UOLSPawnDataAsset* UOLSAssetManager::GetDefaultPawnData() const
{
	return GetAsset(DefaultPawnData);
}

template <typename GameDataClass>
const GameDataClass& UOLSAssetManager::GetOrLoadTypedGameData(const TSoftObjectPtr<GameDataClass>& dataPath)
{
	if (const TObjectPtr<UPrimaryDataAsset>* pResult = GameDataMap.Find(GameDataClass::StaticClass()))
	{
		return *CastChecked<GameDataClass>(*pResult);
	}

	// Does a blocking load if needed
	return *CastChecked<const GameDataClass>(
		LoadGameDataOfClass(GameDataClass::StaticClass(), dataPath, GameDataClass::StaticClass()->GetFName()));
}



UObject* UOLSAssetManager::SynchronousLoadAsset(const FSoftObjectPath& assetPath)
{
	if (assetPath.IsValid())
	{
		TUniquePtr<FScopeLogTime> logTimePtr;

		if (ShouldLogAssetLoads())
		{
			logTimePtr = MakeUnique<FScopeLogTime>(*FString::Printf(TEXT("Synchronously loaded asset [%s]"), *assetPath.ToString()), nullptr, FScopeLogTime::ScopeLog_Seconds);
		}

		if (UAssetManager::IsInitialized())
		{
			return UAssetManager::GetStreamableManager().LoadSynchronous(assetPath, false);
		}

		// Use LoadObject if asset manager isn't ready yet.
		return assetPath.TryLoad();
	}

	return nullptr;
}

bool UOLSAssetManager::ShouldLogAssetLoads()
{
	static bool shouldLogAssetLoads = FParse::Param(FCommandLine::Get(), TEXT("LogAssetLoads"));
	return shouldLogAssetLoads;
}

void UOLSAssetManager::AddLoadedAsset(const UObject* Asset)
{
	if (ensureAlways(Asset))
	{
		FScopeLock LoadedAssetsLock(&LoadedAssetsCritical);
		LoadedAssets.Add(Asset);
	}
}

void UOLSAssetManager::StartInitialLoading()
{
	SCOPED_BOOT_TIMING("UOLSAssetManager::StartInitialLoading");

	// This does all of the scanning, need to do this now even if loads are deferred
	Super::StartInitialLoading();

	STARTUP_JOB(InitializeGameplayCueManager());
	
	{
		// Load base game data asset
		// @Todo uncomment this after implementing GetGameData().
		// STARTUP_JOB_WEIGHTED(GetGameData(), 25.f);
	}
	
	// Run all the queued up startup jobs
	DoAllStartupJobs();
}

#if WITH_EDITOR
void UOLSAssetManager::PreBeginPIE(bool shouldStartSimulate)
{
	Super::PreBeginPIE(shouldStartSimulate);

	{
		FScopedSlowTask slowTask(0, NSLOCTEXT("OLSEditor", "BeginLoadingPIEData", "Loading PIE Data"));
		const bool shouldShowCancelButton = false;
		const bool shouldAllowInPIE = true;
		slowTask.MakeDialog(shouldShowCancelButton, shouldAllowInPIE);

		// @Todo unlock this comment.
		// const ULyraGameData& LocalGameDataCommon = GetGameData();

		// Intentionally after GetGameData to avoid counting GameData time in this timer
		SCOPE_LOG_TIME_IN_SECONDS(TEXT("PreBeginPIE asset preloading complete"), nullptr);

		// You could add preloading of anything else needed for the experience we'll be using here
		// (e.g., by grabbing the default experience from the world settings + the experience override in developer settings)
	}
}
#endif

UPrimaryDataAsset* UOLSAssetManager::LoadGameDataOfClass(
	TSubclassOf<UPrimaryDataAsset> dataClass,
	const TSoftObjectPtr<UPrimaryDataAsset>& dataClassPath,
	FPrimaryAssetType primaryAssetType)
{
	UPrimaryDataAsset* asset = nullptr;

	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("Loading GameData Object"), STAT_GameData, STATGROUP_LoadTime);
	if (!dataClassPath.IsNull())
	{
#if WITH_EDITOR
		FScopedSlowTask slowTask(0, FText::Format(NSLOCTEXT("OLSEditor", "BeginLoadingGameDataTask", "Loading GameData {0}"), FText::FromName(dataClass->GetFName())));
		const bool shouldShowCancelButton = false;
		const bool shouldAllowInPIE = true;
		slowTask.MakeDialog(shouldShowCancelButton, shouldAllowInPIE);
#endif
		// @Todo replace this log with our custom.
		// UE_LOG(LogLyra, Log, TEXT("Loading GameData: %s ..."), *dataClassPath.ToString());
		SCOPE_LOG_TIME_IN_SECONDS(TEXT("    ... GameData loaded!"), nullptr);

		// This can be called recursively in the editor because it is called on demand from PostLoad so force a sync load for primary asset and async load the rest in that case
		if (GIsEditor)
		{
			asset = dataClassPath.LoadSynchronous();
			LoadPrimaryAssetsWithType(primaryAssetType);
		}
		else
		{
			TSharedPtr<FStreamableHandle> handle = LoadPrimaryAssetsWithType(primaryAssetType);
			if (handle.IsValid())
			{
				handle->WaitUntilComplete(0.0f, false);

				// This should always work
				asset = Cast<UPrimaryDataAsset>(handle->GetLoadedAsset());
			}
		}
	}

	if (asset)
	{
		GameDataMap.Add(dataClass, asset);
	}
	else
	{
		// It is not acceptable to fail to load any GameData asset. It will result in soft failures that are hard to diagnose.
		// @Todo replace this log with our custom.
		// UE_LOG(LogLyra, Fatal, TEXT("Failed to load GameData asset at %s. Type %s. This is not recoverable and likely means you do not have the correct data to run %s."), *dataClassPath.ToString(), *primaryAssetType.ToString(), FApp::GetProjectName());
	}

	return asset;
}

void UOLSAssetManager::DoAllStartupJobs()
{
	SCOPED_BOOT_TIMING("ULyraAssetManager::DoAllStartupJobs");
	const double AllStartupJobsStartTime = FPlatformTime::Seconds();

	// if (IsRunningDedicatedServer())
	// {
	// 	// No need for periodic progress updates, just run the jobs
	// 	for (const FLyraAssetManagerStartupJob& StartupJob : StartupJobs)
	// 	{
	// 		StartupJob.DoJob();
	// 	}
	// }
	// else
	// {
	// 	if (StartupJobs.Num() > 0)
	// 	{
	// 		float TotalJobValue = 0.0f;
	// 		for (const FLyraAssetManagerStartupJob& StartupJob : StartupJobs)
	// 		{
	// 			TotalJobValue += StartupJob.JobWeight;
	// 		}
	//
	// 		float AccumulatedJobValue = 0.0f;
	// 		for (FLyraAssetManagerStartupJob& StartupJob : StartupJobs)
	// 		{
	// 			const float JobValue = StartupJob.JobWeight;
	// 			StartupJob.SubstepProgressDelegate.BindLambda([This = this, AccumulatedJobValue, JobValue, TotalJobValue](float NewProgress)
	// 				{
	// 					const float SubstepAdjustment = FMath::Clamp(NewProgress, 0.0f, 1.0f) * JobValue;
	// 					const float OverallPercentWithSubstep = (AccumulatedJobValue + SubstepAdjustment) / TotalJobValue;
	//
	// 					This->UpdateInitialGameContentLoadPercent(OverallPercentWithSubstep);
	// 				});
	//
	// 			StartupJob.DoJob();
	//
	// 			StartupJob.SubstepProgressDelegate.Unbind();
	//
	// 			AccumulatedJobValue += JobValue;
	//
	// 			UpdateInitialGameContentLoadPercent(AccumulatedJobValue / TotalJobValue);
	// 		}
	// 	}
	// 	else
	// 	{
	// 		UpdateInitialGameContentLoadPercent(1.0f);
	// 	}
	// }
	//
	// StartupJobs.Empty();

	// @Todo replace this with our custom log.
	// UE_LOG(LogLyra, Display, TEXT("All startup jobs took %.2f seconds to complete"), FPlatformTime::Seconds() - AllStartupJobsStartTime);
}

void UOLSAssetManager::InitializeGameplayCueManager()
{
	SCOPED_BOOT_TIMING("UOLSAssetManager::InitializeGameplayCueManager");

	// ULyraGameplayCueManager* GCM = ULyraGameplayCueManager::Get();
	// check(GCM);
	// GCM->LoadAlwaysLoadedCues();
}

void UOLSAssetManager::UpdateInitialGameContentLoadPercent(float gameContentPercent)
{
	// Could route this to the early startup loading screen
}