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

/*=============================================================================
	MeshCardRepresentation.cpp
=============================================================================*/

#include "MeshCardRepresentation2.h"

#include "MeshCardRepresentation.h"
#include "MeshUtilities2.h"
#include "Modules/ModuleManager.h"
#include "StaticMeshResources.h"
#include "ProfilingDebugging/CookStats.h"
#include "Templates/UniquePtr.h"
#include "Engine/StaticMesh.h"
#include "Misc/AutomationTest.h"
#include "Misc/QueuedThreadPoolWrapper.h"
#include "ObjectCacheContext.h"

#if WITH_EDITOR
#include "AssetCompilingManager.h"
// #include "DerivedDataCacheInterface.h"
// #include "MeshUtilities.h"
// #include "StaticMeshCompiler.h"
#endif

#if WITH_EDITORONLY_DATA
#include "IMeshBuilderModule.h"
#endif

#if ENABLE_COOK_STATS
namespace CardRepresentationCookStats
{
	FCookStats::FDDCResourceUsageStats UsageStats;
	static FCookStatsManager::FAutoRegisterCallback RegisterCookStats([](FCookStatsManager::AddStatFuncRef AddStat)
	{
		UsageStats.LogStats(AddStat, TEXT("CardRepresentation.Usage"), TEXT(""));
	});
}
#endif

// static TAutoConsoleVariable<int32> CVarCardRepresentation(
// 	TEXT("r.MeshCardRepresentation"),
// 	1,
// 	TEXT(""),
// 	ECVF_ReadOnly);

// static TAutoConsoleVariable<float> CVarCardRepresentationMinSurface(
// 	TEXT("r.MeshCardRepresentation.MinSurface"),
// 	0.2f,
// 	TEXT("Min percentage of surface treshold to spawn a new card, [0;1] range."),
// 	ECVF_ReadOnly);

FCardRepresentationAsyncQueue2* GCardRepresentationAsyncQueue2 = NULL;

int32 GUseAsyncCardRepresentationBuildQueue2 = 1;
static FAutoConsoleVariableRef CVarCardRepresentationAsyncBuildQueue2(
	TEXT("r.MeshCardRepresentation.Async2"),
	GUseAsyncCardRepresentationBuildQueue2,
	TEXT("."),
	ECVF_Default | ECVF_ReadOnly
);

FCardRepresentationAsyncQueue2::FCardRepresentationAsyncQueue2()
{
	MeshUtilities = NULL;
#if WITH_EDITOR
	const int32 MaxConcurrency = -1;
	// In Editor, we allow faster compilation by letting the asset compiler's scheduler organize work.
	ThreadPool = MakeUnique<FQueuedThreadPoolWrapper>(FAssetCompilingManager::Get().GetThreadPool(), MaxConcurrency, [](EQueuedWorkPriority) { return EQueuedWorkPriority::Lowest; });
#else
	const int32 MaxConcurrency = -1;
	ThreadPool = MakeUnique<FQueuedThreadPoolWrapper>(GThreadPool, MaxConcurrency, [](EQueuedWorkPriority) { return EQueuedWorkPriority::Lowest; });
#endif
}

FCardRepresentationAsyncQueue2::~FCardRepresentationAsyncQueue2()
{
}

void FAsyncCardRepresentationTaskWorker2::DoWork()
{
	// Put on background thread to avoid interfering with game-thread bound tasks
	FQueuedThreadPoolTaskGraphWrapper TaskGraphWrapper(ENamedThreads::AnyBackgroundThreadNormalTask);
	GCardRepresentationAsyncQueue2->Build(&Task, TaskGraphWrapper);
}

void FCardRepresentationAsyncQueue2::CancelBackgroundTask(TArray<FAsyncCardRepresentationTask2*> Tasks)
{
	// Do all the cancellation first to make sure none of these tasks
	// get scheduled as we're waiting for completion.
	for (FAsyncCardRepresentationTask2* Task : Tasks)
	{
		if (Task->AsyncTask)
		{
			Task->AsyncTask->Cancel();
		}
	}

	for (FAsyncCardRepresentationTask2* Task : Tasks)
	{
		if (Task->AsyncTask)
		{
			Task->AsyncTask->EnsureCompletion();
			Task->AsyncTask.Reset();
		}
	}
}

void FCardRepresentationAsyncQueue2::StartBackgroundTask(FAsyncCardRepresentationTask2* Task)
{
	check(Task->AsyncTask == nullptr);
	Task->AsyncTask = MakeUnique<FAsyncTask<FAsyncCardRepresentationTaskWorker2>>(*Task);
	Task->AsyncTask->StartBackgroundTask(ThreadPool.Get(), EQueuedWorkPriority::Lowest);
}

void FCardRepresentationAsyncQueue2::ProcessPendingTasks()
{
	FScopeLock Lock(&CriticalSection);
	TArray<FAsyncCardRepresentationTask2*> Tasks = MoveTemp(PendingTasks);
	for (FAsyncCardRepresentationTask2* Task : Tasks)
	{
		if (Task->GenerateSource && Task->GenerateSource->IsCompiling())
		{
			PendingTasks.Add(Task);
		}
		else
		{
			StartBackgroundTask(Task);
		}
	}
}

void FCardRepresentationAsyncQueue2::AddTask(FAsyncCardRepresentationTask2* Task)
{
	// #if WITH_EDITOR
	if (!MeshUtilities)
	{
		MeshUtilities = &FModuleManager::Get().LoadModuleChecked<IMeshUtilities2>(TEXT("MeshUtilities2"));
	}

	{
		// Array protection when called from multiple threads
		FScopeLock Lock(&CriticalSection);
		ReferencedTasks.Add(Task);
	}

	// The Source Mesh's RenderData is not yet ready, postpone the build
	if (Task->GenerateSource->IsCompiling())
	{
		// Array protection when called from multiple threads
		FScopeLock Lock(&CriticalSection);
		PendingTasks.Add(Task);
	}
	else
	{
		// If we're already in worker threads there is no need to launch an async task.
		if (GUseAsyncCardRepresentationBuildQueue2 || !IsInGameThread())
		{
			StartBackgroundTask(Task);
		}
		else
		{
			// To avoid deadlocks, we must queue the inner build tasks on another thread pool, so use the task graph.
			// Put on background thread to avoid interfering with game-thread bound tasks
			FQueuedThreadPoolTaskGraphWrapper TaskGraphWrapper(ENamedThreads::AnyBackgroundThreadNormalTask);
			Build(Task, TaskGraphWrapper);
		}
	}
	// #else
	// 	UE_LOG(LogStaticMesh,Fatal,TEXT("Tried to build a card representation without editor support (this should have been done during cooking)"));
	// #endif
}

void FCardRepresentationAsyncQueue2::CancelBuild(UStaticMesh* StaticMesh)
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FCardRepresentationAsyncQueue2::CancelBuild);

	TArray<FAsyncCardRepresentationTask2*> TasksToCancel;
	{
		FScopeLock Lock(&CriticalSection);
		TArray<FAsyncCardRepresentationTask2*> Tasks = MoveTemp(PendingTasks);
		PendingTasks.Reserve(Tasks.Num());
		for (FAsyncCardRepresentationTask2* Task : Tasks)
		{
			if (Task->GenerateSource != StaticMesh && Task->StaticMesh != StaticMesh)
			{
				PendingTasks.Add(Task);
			}
		}

		Tasks = MoveTemp(ReferencedTasks);
		ReferencedTasks.Reserve(Tasks.Num());
		for (FAsyncCardRepresentationTask2* Task : Tasks)
		{
			if (Task->GenerateSource != StaticMesh && Task->StaticMesh != StaticMesh)
			{
				ReferencedTasks.Add(Task);
			}
			else
			{
				TasksToCancel.Add(Task);
			}
		}
	}

	CancelBackgroundTask(TasksToCancel);
	for (FAsyncCardRepresentationTask2* Task : TasksToCancel)
	{
		if (Task->GeneratedCardRepresentation != nullptr)
		{
			// Rendering thread may still be referencing the old one, use the deferred cleanup interface to delete it next frame when it is safe
			BeginCleanup(Task->GeneratedCardRepresentation);
		}

		delete Task;
	}
}

void FCardRepresentationAsyncQueue2::CancelAllOutstandingBuilds()
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FCardRepresentationAsyncQueue2::CancelAllOutstandingBuilds);

	TArray<FAsyncCardRepresentationTask2*> OutstandingTasks;
	{
		FScopeLock Lock(&CriticalSection);
		PendingTasks.Empty();
		OutstandingTasks = MoveTemp(ReferencedTasks);
	}

	CancelBackgroundTask(OutstandingTasks);
	for (FAsyncCardRepresentationTask2* Task : OutstandingTasks)
	{
		delete Task;
	}
}

void FCardRepresentationAsyncQueue2::RescheduleBackgroundTask(FAsyncCardRepresentationTask2* InTask, EQueuedWorkPriority InPriority)
{
	if (InTask->AsyncTask)
	{
		if (InTask->AsyncTask->GetPriority() != InPriority)
		{
			InTask->AsyncTask->Reschedule(GThreadPool, InPriority);
		}
	}
}

void FCardRepresentationAsyncQueue2::BlockUntilBuildComplete(UStaticMesh* StaticMesh, bool bWarnIfBlocked)
{
	// We will track the wait time here, but only the cycles used.
	// This function is called whether or not an async task is pending,
	// so we have to look elsewhere to properly count how many resources have actually finished building.
	COOK_STAT(auto Timer = CardRepresentationCookStats::UsageStats.TimeAsyncWait());
	COOK_STAT(Timer.TrackCyclesOnly());
	bool bReferenced = false;
	bool bHadToBlock = false;
	double StartTime = 0;

	// #if WITH_EDITOR
	// 	FStaticMeshCompilingManager::Get().FinishCompilation({ StaticMesh });
	// #endif

	do
	{
		ProcessAsyncTasks();

		bReferenced = false;

		{
			FScopeLock Lock(&CriticalSection);
			for (int TaskIndex = 0; TaskIndex < ReferencedTasks.Num(); TaskIndex++)
			{
				if (ReferencedTasks[TaskIndex]->StaticMesh == StaticMesh ||
					ReferencedTasks[TaskIndex]->GenerateSource == StaticMesh)
				{
					bReferenced = true;
					RescheduleBackgroundTask(ReferencedTasks[TaskIndex], EQueuedWorkPriority::Highest);
				}
			}
		}

		if (bReferenced)
		{
			if (!bHadToBlock)
			{
				StartTime = FPlatformTime::Seconds();
			}

			bHadToBlock = true;
			FPlatformProcess::Sleep(.01f);
		}
	}
	while (bReferenced);

	if (bHadToBlock &&
		bWarnIfBlocked
#if WITH_EDITOR
		&& !FAutomationTestFramework::Get().GetCurrentTest() // HACK - Don't output this warning during automation test
#endif
	)
	{
		UE_LOG(LogStaticMesh, Display, TEXT("Main thread blocked for %.3fs for async card representation build of %s to complete!  This can happen if the mesh is rebuilt excessively."),
		       (float)(FPlatformTime::Seconds() - StartTime),
		       *StaticMesh->GetName());
	}
}

void FCardRepresentationAsyncQueue2::BlockUntilAllBuildsComplete()
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FCardRepresentationAsyncQueue2::BlockUntilAllBuildsComplete)
	do
	{
		// #if WITH_EDITOR
		// 		FStaticMeshCompilingManager::Get().FinishAllCompilation();
		// #endif
		// Reschedule as highest prio since we're explicitly waiting on them
		{
			FScopeLock Lock(&CriticalSection);
			for (int TaskIndex = 0; TaskIndex < ReferencedTasks.Num(); TaskIndex++)
			{
				RescheduleBackgroundTask(ReferencedTasks[TaskIndex], EQueuedWorkPriority::Highest);
			}
		}

		ProcessAsyncTasks();
		FPlatformProcess::Sleep(.01f);
	}
	while (GetNumOutstandingTasks() > 0);
}

void FCardRepresentationAsyncQueue2::Build(FAsyncCardRepresentationTask2* Task, FQueuedThreadPool& BuildThreadPool)
{
	// #if WITH_EDITOR
	// Editor 'force delete' can null any UObject pointers which are seen by reference collecting (eg UProperty or serialized)
	if (Task->StaticMesh && Task->GenerateSource)
	{
		TRACE_CPUPROFILER_EVENT_SCOPE(FCardRepresentationAsyncQueue2::Build);

		const FStaticMeshLODResources& LODModel = Task->GenerateSource->GetRenderData()->LODResources[0];

		Task->bSuccess = MeshUtilities->GenerateCardRepresentationData(
			Task->StaticMesh->GetName(),
			Task->SourceMeshData,
			LODModel,
			BuildThreadPool,
			Task->MaterialBlendModes,
			Task->GenerateSource->GetRenderData()->Bounds,
			Task->GenerateSource->GetRenderData()->LODResources[0].DistanceFieldData,
			Task->bGenerateDistanceFieldAsIfTwoSided,
			*Task->GeneratedCardRepresentation);
	}

	CompletedTasks.Push(Task);
	// #endif
}

void FCardRepresentationAsyncQueue2::AddReferencedObjects(FReferenceCollector& Collector)
{
	FScopeLock Lock(&CriticalSection);
	for (int TaskIndex = 0; TaskIndex < ReferencedTasks.Num(); TaskIndex++)
	{
		// Make sure none of the UObjects referenced by the async tasks are GC'ed during the task
		Collector.AddReferencedObject(ReferencedTasks[TaskIndex]->StaticMesh);
		Collector.AddReferencedObject(ReferencedTasks[TaskIndex]->GenerateSource);
	}
}

FString FCardRepresentationAsyncQueue2::GetReferencerName() const
{
	return TEXT("FCardRepresentationAsyncQueue2");
}

void FCardRepresentationAsyncQueue2::ProcessAsyncTasks(bool bLimitExecutionTime)
{
	// #if WITH_EDITOR
	TRACE_CPUPROFILER_EVENT_SCOPE(FCardRepresentationAsyncQueue2::ProcessAsyncTasks);

	ProcessPendingTasks();

	FObjectCacheContextScope ObjectCacheScope;
	const double MaxProcessingTime = 0.016f;
	double StartTime = FPlatformTime::Seconds();
	while (!bLimitExecutionTime || (FPlatformTime::Seconds() - StartTime) < MaxProcessingTime)
	{
		FAsyncCardRepresentationTask2* Task = CompletedTasks.Pop();
		if (Task == nullptr)
		{
			break;
		}

		// We want to count each resource built from a DDC miss, so count each iteration of the loop separately.
		COOK_STAT(auto Timer = CardRepresentationCookStats::UsageStats.TimeSyncWork());

		bool bWasCancelled = false;
		{
			FScopeLock Lock(&CriticalSection);
			bWasCancelled = ReferencedTasks.Remove(Task) == 0;
		}

		if (bWasCancelled)
		{
			continue;
		}

		if (Task->AsyncTask)
		{
			Task->AsyncTask->EnsureCompletion();
			Task->AsyncTask.Reset();
		}

		// Editor 'force delete' can null any UObject pointers which are seen by reference collecting (eg UProperty or serialized)
		if (Task->StaticMesh && Task->bSuccess)
		{
			FCardRepresentationData* OldCardData = Task->StaticMesh->GetRenderData()->LODResources[0].CardRepresentationData;

			// Assign the new data, this is safe because the render threads makes a copy of the pointer at scene proxy creation time.
			Task->StaticMesh->GetRenderData()->LODResources[0].CardRepresentationData = Task->GeneratedCardRepresentation;

			// Any already created render state needs to be dirtied
			if (Task->StaticMesh->GetRenderData()->IsInitialized())
			{
				for (UStaticMeshComponent* Component : ObjectCacheScope.GetContext().GetStaticMeshComponents(Task->StaticMesh))
				{
					if (Component->IsRegistered() && Component->IsRenderStateCreated())
					{
						Component->MarkRenderStateDirty();
					}
				}
			}

			// Rendering thread may still be referencing the old one, use the deferred cleanup interface to delete it next frame when it is safe
			BeginCleanup(OldCardData);

			// {
			// 	TArray<uint8> DerivedData;
			// 	// Save built data to DDC
			// 	FMemoryWriter Ar(DerivedData, /*bIsPersistent=*/ true);
			// 	Ar << *(Task->StaticMesh->GetRenderData()->LODResources[0].CardRepresentationData);
			// 	GetDerivedDataCacheRef().Put(*Task->DDCKey, DerivedData, Task->StaticMesh->GetPathName());
			// 	COOK_STAT(Timer.AddMiss(DerivedData.Num()));
			// }
		}

		delete Task;
	}
	// #endif
}

void FCardRepresentationAsyncQueue2::Shutdown()
{
	CancelAllOutstandingBuilds();

	UE_LOG(LogStaticMesh, Log, TEXT("Abandoning remaining async card representation tasks for shutdown"));
	ThreadPool->Destroy();
}
