#include "DistanceFieldAtlas2.h"
// #include "DistanceFieldAtlas.h"
// #include "HAL/RunnableThread.h"
// #include "HAL/Runnable.h"
// #include "Misc/App.h"
// #include "Serialization/MemoryReader.h"
// #include "Serialization/MemoryWriter.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 "Async/ParallelFor.h"
// #include "DistanceFieldDownsampling.h"
// #include "GlobalShader.h"
#include "RenderGraph.h"
#include "MeshCardRepresentation.h"
#include "Misc/QueuedThreadPoolWrapper.h"
// #include "Async/Async.h"
#include "ObjectCacheContext.h"

// #if WITH_EDITOR
// #include "DerivedDataCacheInterface.h"
#include "Engine/Public/DistanceFieldAtlas.h"

// #include "AssetCompilingManager.h"
#include "MeshCardRepresentation2.h"
#include "StaticMeshCompiler.h"
#include "MeshUtilities2/Public/MeshUtilities2.h"
// #endif

CSV_DEFINE_CATEGORY(DistanceField2, false);

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


FDistanceFieldAsyncQueue2* GDistanceFieldAsyncQueue2 = nullptr;

int32 GUseAsyncDistanceFieldBuildQueue2 = 1;
static FAutoConsoleVariableRef CVarAOAsyncBuildQueue(
	TEXT("r.AOAsyncBuildQueue2"),
	GUseAsyncDistanceFieldBuildQueue2,
	TEXT("是否从网格异步构建距离场体数据。"),
	ECVF_Default | ECVF_ReadOnly
);

//构建距离场
void BuildMeshDistanceField(UStaticMesh* StaticMesh)
{
	GDistanceFieldAsyncQueue2->CancelBuild(StaticMesh);
	GCardRepresentationAsyncQueue2->CancelBuild(StaticMesh);

	//ref:StaticMesh.cpp 2782
	auto RenderData = StaticMesh->GetRenderData();
	if (RenderData->LODResources.IsValidIndex(0))
	{
		auto& LODResource = RenderData->LODResources[0];
		if (!LODResource.DistanceFieldData)
		{
			LODResource.DistanceFieldData = new FDistanceFieldVolumeData();
			LODResource.DistanceFieldData->AssetName = StaticMesh->GetFName();
		}

		// We don't actually build the resource until later, so only track the cycles used here.
		// COOK_STAT(Timer.TrackCyclesOnly());
		FAsyncDistanceFieldTask2* NewTask = new FAsyncDistanceFieldTask2;
		NewTask->DDCKey = "";
		// check(Mesh && GenerateSource);
		// NewTask->TargetPlatform = RunningPlatform;
		NewTask->StaticMesh = StaticMesh;
		NewTask->GenerateSource = StaticMesh; //GenerateSource;
		NewTask->DistanceFieldResolutionScale = 2; //DistanceFieldResolutionScale;
		NewTask->bGenerateDistanceFieldAsIfTwoSided = false; // bGenerateDistanceFieldAsIfTwoSided;
		NewTask->GeneratedVolumeData = new FDistanceFieldVolumeData(); //;给一个新的用于存放
		NewTask->GeneratedVolumeData->AssetName = StaticMesh->GetFName();
		NewTask->GeneratedVolumeData->bAsyncBuilding = true;

		for (int32 MaterialIndex = 0; MaterialIndex < StaticMesh->GetStaticMaterials().Num(); MaterialIndex++)
		{
			FSignedDistanceFieldBuildMaterialData2 MaterialData;
			// Default material blend mode
			MaterialData.BlendMode = BLEND_Opaque;
			MaterialData.bTwoSided = false;

			if (StaticMesh->GetStaticMaterials()[MaterialIndex].MaterialInterface)
			{
				MaterialData.BlendMode = StaticMesh->GetStaticMaterials()[MaterialIndex].MaterialInterface->GetBlendMode();
				MaterialData.bTwoSided = StaticMesh->GetStaticMaterials()[MaterialIndex].MaterialInterface->IsTwoSided();
			}

			NewTask->MaterialBlendModes.Add(MaterialData);
		}

		// // Nanite overrides source static mesh with a coarse representation. Need to load original data before we build the mesh SDF.
		// if (StaticMesh->NaniteSettings.bEnabled)
		// {
		// 	IMeshBuilderModule& MeshBuilderModule = IMeshBuilderModule::GetForPlatform(TargetPlatform);
		// 	if (!MeshBuilderModule.BuildMeshVertexPositions(Mesh, NewTask->SourceMeshData.TriangleIndices, NewTask->SourceMeshData.VertexPositions))
		// 	{
		// 		UE_LOG(LogStaticMesh, Error, TEXT("Failed to build static mesh. See previous line(s) for details."));
		// 	}
		// }

		GDistanceFieldAsyncQueue2->AddTask(NewTask);
	}
}

void BuildMeshCardRepresentation(UStaticMesh* StaticMeshAsset, FStaticMeshRenderData& RenderData, FSourceMeshDataForDerivedDataTask* OptionalSourceMeshData)
{
	if (RenderData.LODResources.IsValidIndex(0))
	{
		if (!RenderData.LODResources[0].CardRepresentationData)
		{
			RenderData.LODResources[0].CardRepresentationData = new FCardRepresentationData();
		}

		// const FMeshBuildSettings& BuildSettings = StaticMeshAsset->GetSourceModel(0).BuildSettings;
		// UStaticMesh* MeshToGenerateFrom = StaticMeshAsset;

		// We don't actually build the resource until later, so only track the cycles used here.
		// COOK_STAT(Timer.TrackCyclesOnly());
		FAsyncCardRepresentationTask2* NewTask = new FAsyncCardRepresentationTask2;
		// NewTask->DDCKey = InDDCKey;
		check(StaticMeshAsset);
		NewTask->StaticMesh = StaticMeshAsset;
		NewTask->GenerateSource = StaticMeshAsset;
		NewTask->GeneratedCardRepresentation = new FCardRepresentationData();
		NewTask->bGenerateDistanceFieldAsIfTwoSided = false;

		for (int32 MaterialIndex = 0; MaterialIndex < StaticMeshAsset->GetStaticMaterials().Num(); MaterialIndex++)
		{
			FSignedDistanceFieldBuildMaterialData2 MaterialData;
			// Default material blend mode
			MaterialData.BlendMode = BLEND_Opaque;
			MaterialData.bTwoSided = false;

			if (StaticMeshAsset->GetStaticMaterials()[MaterialIndex].MaterialInterface)
			{
				MaterialData.BlendMode = StaticMeshAsset->GetStaticMaterials()[MaterialIndex].MaterialInterface->GetBlendMode();
				MaterialData.bTwoSided = StaticMeshAsset->GetStaticMaterials()[MaterialIndex].MaterialInterface->IsTwoSided();
			}

			NewTask->MaterialBlendModes.Add(MaterialData);
		}

		// Nanite overrides source static mesh with a coarse representation. Need to load original data before we build the mesh SDF.
		if (OptionalSourceMeshData)
		{
			NewTask->SourceMeshData = *OptionalSourceMeshData;
		}
		// else if (StaticMeshAsset->NaniteSettings.bEnabled)
		// {
		// 	IMeshBuilderModule& MeshBuilderModule = IMeshBuilderModule::GetForPlatform(TargetPlatform);
		// 	if (!MeshBuilderModule.BuildMeshVertexPositions(Mesh, NewTask->SourceMeshData.TriangleIndices, NewTask->SourceMeshData.VertexPositions))
		// 	{
		// 		UE_LOG(LogStaticMesh, Error, TEXT("Failed to build static mesh. See previous line(s) for details."));
		// 	}
		// }

		GCardRepresentationAsyncQueue2->AddTask(NewTask);
	}
}


FAsyncDistanceFieldTask2::FAsyncDistanceFieldTask2()
	: StaticMesh(nullptr)
	  , GenerateSource(nullptr)
	  , DistanceFieldResolutionScale(0.0f)
	  , bGenerateDistanceFieldAsIfTwoSided(false)
	  , GeneratedVolumeData(nullptr)
{
}

FDistanceFieldAsyncQueue2::FDistanceFieldAsyncQueue2()
{
	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
}

FDistanceFieldAsyncQueue2::~FDistanceFieldAsyncQueue2()
{
}

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

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

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

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

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

void FDistanceFieldAsyncQueue2::AddTask(FAsyncDistanceFieldTask2* 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 (GUseAsyncDistanceFieldBuildQueue2 || !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 distance field without editor support (this should have been done during cooking)"));
	// #endif
}

void FDistanceFieldAsyncQueue2::CancelBuild(UStaticMesh* StaticMesh)
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FDistanceFieldAsyncQueue2::CancelBuild)

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

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

	CancelBackgroundTask(TasksToCancel);
	for (FAsyncDistanceFieldTask2* Task : TasksToCancel)
	{
		if (Task->GeneratedVolumeData != 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->GeneratedVolumeData);
		}

		delete Task;
	}
}

void FDistanceFieldAsyncQueue2::CancelAllOutstandingBuilds()
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FDistanceFieldAsyncQueue2::CancelAllOutstandingBuilds)

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

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

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

void FDistanceFieldAsyncQueue2::BlockUntilBuildComplete(UStaticMesh* StaticMesh, bool bWarnIfBlocked)
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FDistanceFieldAsyncQueue2::BlockUntilBuildComplete)

	// 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 = DistanceFieldCookStats::UsageStats.TimeAsyncWait());
	COOK_STAT(Timer.TrackCyclesOnly());
	bool bReferenced = false;
	bool bHadToBlock = false;
	double StartTime = 0;

	TSet<UStaticMesh*> RequiredFinishCompilation;
	do
	{
		ProcessAsyncTasks();

		bReferenced = false;
		RequiredFinishCompilation.Reset();

		// Reschedule the tasks we're waiting on as highest prio
		{
			FScopeLock Lock(&CriticalSection);
			for (int TaskIndex = 0; TaskIndex < ReferencedTasks.Num(); TaskIndex++)
			{
				if (ReferencedTasks[TaskIndex]->StaticMesh == StaticMesh ||
					ReferencedTasks[TaskIndex]->GenerateSource == StaticMesh)
				{
					bReferenced = true;

					// If the task we are waiting on depends on other static meshes
					// we need to force finish them too.
					// #if WITH_EDITOR

					if (ReferencedTasks[TaskIndex]->GenerateSource != nullptr &&
						ReferencedTasks[TaskIndex]->GenerateSource->IsCompiling())
					{
						RequiredFinishCompilation.Add(ReferencedTasks[TaskIndex]->GenerateSource);
					}

					if (ReferencedTasks[TaskIndex]->StaticMesh != nullptr &&
						ReferencedTasks[TaskIndex]->StaticMesh->IsCompiling())
					{
						RequiredFinishCompilation.Add(ReferencedTasks[TaskIndex]->StaticMesh);
					}
					// #endif

					RescheduleBackgroundTask(ReferencedTasks[TaskIndex], EQueuedWorkPriority::Highest);
				}
			}
		}

#if WITH_EDITOR
		// Call the finish compilation outside of the critical section since those compilations
		// might need to register new distance field tasks which also uses the critical section.
		if (RequiredFinishCompilation.Num())
		{
			FStaticMeshCompilingManager::Get().FinishCompilation(RequiredFinishCompilation.Array());
		}
#endif

		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 distance field build of %s to complete!  This can happen if the mesh is rebuilt excessively."),
		       (float)(FPlatformTime::Seconds() - StartTime),
		       *StaticMesh->GetName());
	}
}

void FDistanceFieldAsyncQueue2::BlockUntilAllBuildsComplete()
{
	TRACE_CPUPROFILER_EVENT_SCOPE(FDistanceFieldAsyncQueue2::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 FDistanceFieldAsyncQueue2::Build(FAsyncDistanceFieldTask2* Task, FQueuedThreadPool& BuildThreadPool)
{
	//#if WITH_EDITOR
	// Editor 'force delete' can null any UObject pointers which are seen by reference collecting (eg FProperty or serialized)
	if (Task->StaticMesh && Task->GenerateSource)
	{
		TRACE_CPUPROFILER_EVENT_SCOPE(FDistanceFieldAsyncQueue2::Build);

		const FStaticMeshLODResources& LODModel = Task->GenerateSource->GetRenderData()->LODResources[0];
		MeshUtilities->GenerateSignedDistanceFieldVolumeData(
			Task->StaticMesh->GetName(),
			Task->SourceMeshData,
			LODModel,
			BuildThreadPool,
			Task->MaterialBlendModes,
			Task->GenerateSource->GetRenderData()->Bounds,
			Task->DistanceFieldResolutionScale,
			Task->bGenerateDistanceFieldAsIfTwoSided,
			*Task->GeneratedVolumeData);
	}

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

void FDistanceFieldAsyncQueue2::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 FDistanceFieldAsyncQueue2::GetReferencerName() const
{
	return TEXT("FDistanceFieldAsyncQueue2");
}

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

	ProcessPendingTasks();

	FObjectCacheContextScope ObjectCacheScope;
	const double MaxProcessingTime = 0.016f;
	double StartTime = FPlatformTime::Seconds();
	while (!bLimitExecutionTime || (FPlatformTime::Seconds() - StartTime) < MaxProcessingTime)
	{
		FAsyncDistanceFieldTask2* 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 = DistanceFieldCookStats::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 FProperty or serialized)
		if (Task->StaticMesh)
		{
			Task->GeneratedVolumeData->bAsyncBuilding = false;
			FDistanceFieldVolumeData* OldVolumeData = Task->StaticMesh->GetRenderData()->LODResources[0].DistanceFieldData;

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

			// Renderstates are not initialized between UStaticMesh::PreEditChange() and UStaticMesh::PostEditChange()
			if (Task->StaticMesh->GetRenderData()->IsInitialized())
			{
				for (UStaticMeshComponent* Component : ObjectCacheScope.GetContext().GetStaticMeshComponents(Task->StaticMesh))
				{
					if (Component->IsRegistered() && Component->IsRenderStateCreated())
					{
						Component->MarkRenderStateDirty();
					}
				}
			}

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

			//第二阶段
			BuildMeshCardRepresentation(Task->StaticMesh, *Task->StaticMesh->GetRenderData(), &Task->SourceMeshData);
		}

		delete Task;
	}
	// #endif
}

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

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