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

/*=============================================================================
	DistanceFieldAtlas.h
=============================================================================*/

#pragma once

#include "CoreMinimal.h"
#include "Containers/LockFreeList.h"
#include "ProfilingDebugging/ResourceSize.h"
#include "Engine/EngineTypes.h"
#include "UObject/GCObject.h"
#include "Templates/UniquePtr.h"
#include "DerivedMeshDataTaskUtils.h"
#include "Async/AsyncWork.h"

class MESHUTILITIES2_API FSignedDistanceFieldBuildMaterialData2
{
public:
	EBlendMode BlendMode;
	bool bTwoSided;
};

class UStaticMesh;
class UTexture2D;

template <class T>
class TLockFreePointerListLIFO;

// Change DDC key when modifying these (or any DF encoding logic)
namespace DistanceField2
{
	// One voxel border around object for handling gradient
	constexpr int32 MeshDistanceFieldObjectBorder = 1;
	constexpr int32 UniqueDataBrickSize = 7;
	// Half voxel border around brick for trilinear filtering
	constexpr int32 BrickSize = 8;
	// Trade off between SDF memory and number of steps required to find intersection
	constexpr int32 BandSizeInVoxels = 4;
	constexpr int32 NumMips = 3;
	constexpr uint32 InvalidBrickIndex = 0xFFFFFFFF;
	constexpr EPixelFormat DistanceFieldFormat = PF_G8;

	// Must match LoadDFAssetData
	constexpr uint32 MaxIndirectionDimension = 1024;
};

class FSparseDistanceFieldMip2
{
public:
	FSparseDistanceFieldMip2() :
		IndirectionDimensions(FIntVector::ZeroValue),
		NumDistanceFieldBricks(0),
		VolumeToVirtualUVScale(FVector::ZeroVector),
		VolumeToVirtualUVAdd(FVector::ZeroVector),
		DistanceFieldToVolumeScaleBias(FVector2D::ZeroVector),
		BulkOffset(0),
		BulkSize(0)
	{
	}

	FIntVector IndirectionDimensions;
	int32 NumDistanceFieldBricks;
	FVector VolumeToVirtualUVScale;
	FVector VolumeToVirtualUVAdd;
	FVector2D DistanceFieldToVolumeScaleBias;
	uint32 BulkOffset;
	uint32 BulkSize;

	friend FArchive& operator<<(FArchive& Ar, FSparseDistanceFieldMip2& Mip)
	{
		Ar << Mip.IndirectionDimensions << Mip.NumDistanceFieldBricks << Mip.VolumeToVirtualUVScale << Mip.VolumeToVirtualUVAdd << Mip.DistanceFieldToVolumeScaleBias << Mip.BulkOffset << Mip.BulkSize;
		return Ar;
	}

	SIZE_T GetResourceSizeBytes() const
	{
		FResourceSizeEx ResSize;
		GetResourceSizeEx(ResSize);
		return ResSize.GetTotalMemoryBytes();
	}

	void GetResourceSizeEx(FResourceSizeEx& CumulativeResourceSize) const
	{
		CumulativeResourceSize.AddDedicatedSystemMemoryBytes(sizeof(*this));
	}
};

class FAsyncDistanceFieldTask2;

class FAsyncDistanceFieldTaskWorker2 : public FNonAbandonableTask
{
public:
	FAsyncDistanceFieldTaskWorker2(FAsyncDistanceFieldTask2& InTask)
		: Task(InTask)
	{
	}

	FORCEINLINE TStatId GetStatId() const
	{
		RETURN_QUICK_DECLARE_CYCLE_STAT(FAsyncDistanceFieldTaskWorker2, STATGROUP_ThreadPoolAsyncTasks);
	}

	void DoWork();

private:
	FAsyncDistanceFieldTask2& Task;
};

/** A task to build a distance field for a single mesh */
class MESHUTILITIES2_API FAsyncDistanceFieldTask2
{
public:
	FAsyncDistanceFieldTask2();

	// #if WITH_EDITOR
	TArray<FSignedDistanceFieldBuildMaterialData2> MaterialBlendModes;
	// #endif
	FSourceMeshDataForDerivedDataTask SourceMeshData;
	UStaticMesh* StaticMesh;
	UStaticMesh* GenerateSource;
	float DistanceFieldResolutionScale;
	bool bGenerateDistanceFieldAsIfTwoSided;
	const ITargetPlatform* TargetPlatform;
	FString DDCKey;
	FDistanceFieldVolumeData* GeneratedVolumeData;
	TUniquePtr<FAsyncTask<FAsyncDistanceFieldTaskWorker2>> AsyncTask = nullptr;
};

/** Class that manages asynchronous building of mesh distance fields. */
class MESHUTILITIES2_API FDistanceFieldAsyncQueue2 : public FGCObject
{
public:
	FDistanceFieldAsyncQueue2();

	virtual ~FDistanceFieldAsyncQueue2();

	/** Adds a new build task. (Thread-Safe) */
	void AddTask(FAsyncDistanceFieldTask2* Task);

	/** Cancel the build on this specific static mesh or block until it is completed if already started. */
	void CancelBuild(UStaticMesh* StaticMesh);

	/** Blocks the main thread until the async build are either cancelled or completed. */
	void CancelAllOutstandingBuilds();

	/** Blocks the main thread until the async build of the specified mesh is complete. */
	void BlockUntilBuildComplete(UStaticMesh* StaticMesh, bool bWarnIfBlocked);

	/** Blocks the main thread until all async builds complete. */
	void BlockUntilAllBuildsComplete();

	/** Called once per frame, fetches completed tasks and applies them to the scene. */
	void ProcessAsyncTasks(bool bLimitExecutionTime = false);

	/** Exposes UObject references used by the async build. */
	virtual void AddReferencedObjects(FReferenceCollector& Collector) override;

	/** Returns name of class for reference tracking */
	virtual FString GetReferencerName() const override;

	/** Blocks until it is safe to shut down (worker threads are idle). */
	void Shutdown();

	int32 GetNumOutstandingTasks() const
	{
		FScopeLock Lock(&CriticalSection);
		return ReferencedTasks.Num();
	}

private:
	friend FAsyncDistanceFieldTaskWorker2;
	void ProcessPendingTasks();

	TUniquePtr<FQueuedThreadPool> ThreadPool;

	/** Builds a single task with the given threadpool.  Called from the worker thread. */
	void Build(FAsyncDistanceFieldTask2* Task, class FQueuedThreadPool& ThreadPool);

	/** Change the priority of the background task. */
	void RescheduleBackgroundTask(FAsyncDistanceFieldTask2* InTask, EQueuedWorkPriority InPriority);

	/** Task will be sent to a background worker. */
	void StartBackgroundTask(FAsyncDistanceFieldTask2* Task);

	/** Cancel or finish any background work for the given task. */
	void CancelBackgroundTask(TArray<FAsyncDistanceFieldTask2*> Tasks);

	/** Game-thread managed list of tasks in the async system. */
	TArray<FAsyncDistanceFieldTask2*> ReferencedTasks;

	/** Tasks that are waiting on static mesh compilation to proceed */
	TArray<FAsyncDistanceFieldTask2*> PendingTasks;

	/** Tasks that have completed processing. */
	// consider changing this from FIFO to Unordered, which may be faster
	TLockFreePointerListLIFO<FAsyncDistanceFieldTask2> CompletedTasks;

	class IMeshUtilities2* MeshUtilities;

	mutable FCriticalSection CriticalSection;
};

/** Global build queue. */
extern MESHUTILITIES2_API FDistanceFieldAsyncQueue2* GDistanceFieldAsyncQueue2;

extern MESHUTILITIES2_API FString BuildDistanceFieldDerivedDataKey2(const FString& InMeshKey);

extern MESHUTILITIES2_API void BuildMeshDistanceField(UStaticMesh* StaticMesh);
extern MESHUTILITIES2_API void BuildMeshCardRepresentation(UStaticMesh* StaticMeshAsset, class FStaticMeshRenderData& RenderData, FSourceMeshDataForDerivedDataTask* OptionalSourceMeshData);
