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

#pragma once

#ifdef WITH_NNE_RUNTIME_IREE_SHADER

#include "CoreMinimal.h"
#include "NNERuntimeIREEShaderCompileResult.h"
#include "Engine/EngineTypes.h"
#include "RenderResource.h"
#include "RenderingThread.h"
#include "RenderDeferredCleanup.h"
#include "RHI.h"
#include "SceneTypes.h"
#include "Shader.h"
#include "ShaderCompiler.h"
#include "Templates/RefCounting.h"

struct FNNERuntimeIREEShaderParametersMetadataAllocations;
class FNNERuntimeIREEResource;
class FNNERuntimeIREEShader;

/** Stores outputs from the shader compile that need to be saved. */
class FNNERuntimeIREECompilationOutput
{
	DECLARE_TYPE_LAYOUT(FNNERuntimeIREECompilationOutput, NonVirtual);
public:
	FNNERuntimeIREECompilationOutput()
	{}
};

/** Contains all the information needed to uniquely identify a FNNERuntimeIREEShaderMapID. */
class FNNERuntimeIREEShaderMapId
{
	DECLARE_TYPE_LAYOUT(FNNERuntimeIREEShaderMapId, NonVirtual);
public:
	/** Feature level that the shader map is going to be compiled for.  */
	LAYOUT_FIELD(ERHIFeatureLevel::Type, FeatureLevel);

	/**
	 * The hash computed from the shader code generated by the library.
	 */
	LAYOUT_FIELD(uint64, ShaderCodeHash);

	/** Shader types of shaders that are inlined in this shader map in the DDC. */
	LAYOUT_FIELD(TMemoryImageArray<FShaderTypeDependency>, ShaderTypeDependencies);

	/*
	* Type layout parameters of the memory image
	*/
	LAYOUT_FIELD(FPlatformTypeLayoutParameters, LayoutParams);

public:
	FNNERuntimeIREEShaderMapId()
		: FeatureLevel(ERHIFeatureLevel::SM5)
	{ }

	~FNNERuntimeIREEShaderMapId() = default;

#if WITH_EDITOR
	void SetShaderDependencies(const TArray<FShaderType*>& InShaderTypes, EShaderPlatform InShaderPlatform);
#endif

	//void Serialize(FArchive& Ar);

	friend uint32 GetTypeHash(const FNNERuntimeIREEShaderMapId& Ref)
	{
		return GetTypeHash(Ref.ShaderCodeHash);
	}

	SIZE_T GetSizeBytes() const
	{
		return sizeof(*this);
	}

	/** Hashes the kernel specific part of this shader map Id. */
	void GetKernelHash(FSHAHash& OutHash) const;

	/**
	 * Tests this set against another for equality, disregarding override settings.
	 *
	 * @param InReferenceSet	The set to compare against
	 * @return					true if the sets are equal
	 */
	bool operator==(const FNNERuntimeIREEShaderMapId& InReferenceSet) const;

	bool operator!=(const FNNERuntimeIREEShaderMapId& InReferenceSet) const
	{
		return !(*this == InReferenceSet);
	}

#if WITH_EDITOR
	/** Appends string representations of this Id to a key string. */
	void AppendKeyString(FString& OutKeyString) const;
#endif // WITH_EDITOR

	/** Returns true if the requested shader type is a dependency of this shader map Id. */
	bool ContainsShaderType(const FShaderType* ShaderType) const;
};

class FNNERuntimeIREEShaderMapContent : public FShaderMapContent
{
	using Super = FShaderMapContent;
	friend class FNNERuntimeIREEShaderMap;
	DECLARE_TYPE_LAYOUT(FNNERuntimeIREEShaderMapContent, NonVirtual);
private:
	explicit FNNERuntimeIREEShaderMapContent(EShaderPlatform InPlatform) : Super(InPlatform) {}

	/** The kernel's user friendly name */
	LAYOUT_FIELD(FMemoryImageString, FriendlyName);

	/** The static parameter set that this shader map was compiled with */
	LAYOUT_FIELD(FNNERuntimeIREEShaderMapId, ShaderMapId);

	/** Shader compilation output */
	LAYOUT_FIELD(FNNERuntimeIREECompilationOutput, CompilationOutput);
};


/**
 * All the shader permutations for a kernel.
 */
class FNNERuntimeIREEShaderMap : public TShaderMap<FNNERuntimeIREEShaderMapContent, FShaderMapPointerTable>, public FDeferredCleanupInterface
{
public:
	using Super = TShaderMap<FNNERuntimeIREEShaderMapContent, FShaderMapPointerTable>;

	/**
	 * Finds the shader map for a kernel.
	 * @param InShaderMapId - Id to find
	 * @param Platform - The platform to lookup for
	 * @return nullptr if no cached shader map was found.
	 */
	static FNNERuntimeIREEShaderMap* FindId(const FNNERuntimeIREEShaderMapId& InShaderMapId, EShaderPlatform InPlatform);

	// ShaderMap interface
	template<typename ShaderType> TShaderRef<ShaderType> GetShader(int32 PermutationId) const { return TShaderRef<ShaderType>(GetContent()->GetShader<ShaderType>(PermutationId), *this); }
	TShaderRef<FShader> GetShader(FShaderType* ShaderType, int32 PermutationId) const { return TShaderRef<FShader>(GetContent()->GetShader(ShaderType, PermutationId), *this); }

	FNNERuntimeIREEShaderMap();
	~FNNERuntimeIREEShaderMap();

#if WITH_EDITOR
	/**
	 * Compiles the shaders for a kernel and caches them in this shader map.
	 * @param InKernel - The kernel to compile shaders for.
	 * @param InShaderMapId - the set of static parameters to compile for
	 * @param InPlatform - The platform to compile to
	 */
	void Compile(
		FNNERuntimeIREEResource* InKernel,
		const FNNERuntimeIREEShaderMapId& InShaderMapId,
		TRefCountPtr<FSharedShaderCompilerEnvironment> InCompilationEnvironment,
		const FNNERuntimeIREECompilationOutput& InCompilationOutput,
		EShaderPlatform InPlatform,
		bool bSynchronousCompile,
		bool bApplyCompletedShaderMapForRendering
		);

	/** Sorts the incoming compiled jobs into the appropriate shader maps, and finalizes this shader map so that it can be used for rendering. */
	bool ProcessCompilationResults(const TArray<FShaderCommonCompileJobPtr>& InCompilationResults, int32& InOutResultIndex, float& InOutTimeBudget);
#endif // WITH_EDITOR

	/**
	 * Checks whether the shader map is missing any shader types necessary for the given kernel.
	 * @param InKernel - The NNERuntimeIREE shader resource which is checked.
	 * @return True if the shader map has all of the shader types necessary.
	 */
	bool IsComplete(const FNNERuntimeIREEResource* InKernel, bool bSilent);

	/** Builds a list of the shaders in a shader map. */
	void GetShaderList(TMap<FShaderId, TShaderRef<FShader>>& OutShaders) const;
	virtual void GetShaderList(TMap<FHashedName, TShaderRef<FShader>>& OutShaders) const override;
	virtual void GetShaderPipelineList(TArray<FShaderPipelineRef>& OutShaderPipelines) const override;

	/** Registers a NNERuntimeIREE shader map in the global map so it can be used. */
	void Register(EShaderPlatform InShaderPlatform);

	// Reference counting.
	void AddRef();
	void Release();

	/** Serializes the shader map. */
	bool Serialize(FArchive& Ar);

	// Accessors
	const FNNERuntimeIREEShaderMapId& GetShaderMapId() const { return GetContent()->ShaderMapId; }
	EShaderPlatform GetShaderPlatform() const { return GetContent()->GetShaderPlatform(); }
	const FMemoryImageString& GetFriendlyName() const { return GetContent()->FriendlyName; }
	bool IsCompilationFinalized() const { return bCompilationFinalized; }
	bool CompiledSuccessfully() const { return bCompiledSuccessfully; }

#if WITH_EDITOR
	static TMap<TRefCountPtr<FNNERuntimeIREEShaderMap>, TArray<FNNERuntimeIREEResource*> > &GetInFlightShaderMaps() { return ShaderMapsBeingCompiled; }

	void SetCompiledSuccessfully(bool bSuccess) { bCompiledSuccessfully = bSuccess; }
#endif // WITH_EDITOR

	int32 GetNumRefs() const { return NumRefs; }
	uint32 GetCompilingId()  { return CompilingId; }

private:
#if WITH_EDITOR
	FShader* ProcessCompilationResultsForSingleJob(FShaderCompileJob& InSingleJob, const FSHAHash& InShaderMapHash);
#endif

	bool IsIREEShaderComplete(const FNNERuntimeIREEResource* InKernel, const FNNERuntimeIREEShaderType* InShaderType, bool bSilent);

	/**
	 * A global map from a kernel's ID and static switch set to any shader map cached for that kernel.
	 * Note: this does not necessarily contain all kernel shader maps in memory.  Shader maps with the same key can evict each other.
	 * No ref counting needed as these are removed on destruction of the shader map.
	 */
	static TMap<FNNERuntimeIREEShaderMapId, FNNERuntimeIREEShaderMap*> GIdToIREEShaderMap[SP_NumPlatforms];

	/**
	 * All kernel shader maps in memory.
	 * No ref counting needed as these are removed on destruction of the shader map.
	 */
	static TArray<FNNERuntimeIREEShaderMap*> AllKernelShaderMaps;

#if WITH_EDITOR
	/** Tracks resources and their shader maps that need to be compiled but whose compilation is being deferred. */
	static TMap<TRefCountPtr<FNNERuntimeIREEShaderMap>, TArray<FNNERuntimeIREEResource*> > ShaderMapsBeingCompiled;
#endif

	/** Uniquely identifies this shader map during compilation, needed for deferred compilation where shaders from multiple shader maps are compiled together. */
	uint32 CompilingId;

	mutable int32 NumRefs;

	/** Used to catch errors where the shader map is deleted directly. */
	bool bDeletedThroughDeferredCleanup;

	/** Indicates whether this shader map has been registered in GIdToIREEShaderMap */
	uint32 bRegistered : 1;

	/**
	 * Indicates whether this shader map has had ProcessCompilationResults called after Compile.
	 * The shader map must not be used on the rendering thread unless bCompilationFinalized is true.
	 */
	uint32 bCompilationFinalized : 1;

	uint32 bCompiledSuccessfully : 1;

};

class NNERUNTIMEIREESHADER_API FNNERuntimeIREEResource
{
public:
	FNNERuntimeIREEResource();
	virtual ~FNNERuntimeIREEResource();

	/**
	 * Caches the shaders for this kernel with no static parameters on the given platform.
	 * This is used by UKernel
	 */
	bool CacheShaders(EShaderPlatform InPlatform, const ITargetPlatform* TargetPlatform, bool bApplyCompletedShaderMapForRendering, bool bSynchronous);
	bool CacheShaders(const FNNERuntimeIREEShaderMapId& InShaderMapId, EShaderPlatform InPlatform, bool bApplyCompletedShaderMapForRendering, bool bSynchronous);

	/**
	 * Should the shader for this kernel with the given platform and shader type combination be compiled
	 *
	 * @param InPlatform		The platform currently being compiled for
	 * @param InShaderType		Which shader is being compiled
	 *
	 * @return true if the shader should be compiled
	 */
	virtual bool ShouldCache(EShaderPlatform InPlatform, const FShaderType* InShaderType) const;

	bool SerializeShaderMap(FArchive& Ar);

	void GetDependentShaderTypes(EShaderPlatform InPlatform, TArray<FShaderType*>& OutShaderTypes) const;
	virtual void GetShaderMapId(EShaderPlatform InPlatform, const ITargetPlatform* TargetPlatform, FNNERuntimeIREEShaderMapId& OutId) const;

	/**
	 * Called when compilation finishes, after the GameThreadShaderMap is set and the render command to set the RenderThreadShaderMap is queued
	 */
	virtual void NotifyCompilationFinished(FString const& ResultMessage);

#if WITH_EDITOR
	/**
	 * Cancels all outstanding compilation jobs
	 */
	void CancelCompilation();

	/**
	 * Blocks until compilation has completed. Returns immediately if a compilation is not outstanding.
	 */
	void FinishCompilation();
#endif // WITH_EDITOR

	// Accessors.
	FNNERuntimeIREEShaderCompileResults const& GetCompilationResults() const { return CompilationResults; }
	void SetCompilationResults(FNNERuntimeIREEShaderCompileResults const& InCompilationResults) { CompilationResults = InCompilationResults; }

	ERHIFeatureLevel::Type GetFeatureLevel() const { return FeatureLevel; }

	FNNERuntimeIREEShaderMap* GetGameThreadShaderMap() const
	{
		checkSlow(IsInGameThread() || IsInAsyncLoadingThread());
		return GameThreadShaderMap;
	}

	/** Returns owner name for tracking */
	FName GetOwnerFName() const
	{
		return AssetPath;
	}

	/** Note: SetRenderingThreadShaderMap must also be called with the same value, but from the rendering thread. */
	void SetGameThreadShaderMap(FNNERuntimeIREEShaderMap* InShaderMap)
	{
		checkSlow(IsInGameThread() || IsInAsyncLoadingThread());
		GameThreadShaderMap = InShaderMap;
		if (LIKELY(GameThreadShaderMap))
		{
			GameThreadShaderMap->GetResource()->SetOwnerName(GetOwnerFName());
		}
	}

	/** Note: SetGameThreadShaderMap must also be called with the same value, but from the game thread. */
	void SetRenderingThreadShaderMap(FNNERuntimeIREEShaderMap* InShaderMap);

	void AddCompileId(uint32 InIdentifier) 
	{
		OutstandingCompileShaderMapIds.Add(InIdentifier);
	}

	void RemoveOutstandingCompileId(const int32 InOldOutstandingCompileShaderMapId);

	const FString& GetHLSLSource() const
	{
		return ShaderSource;
	}

	const FString& GetEntryPoint() const
	{
		return ShaderEntryPoint;
	}

	const FShaderParametersMetadata* GetShaderParamMetadata() const
	{
		return ShaderParameterMetadata;
	}

	const FString& GetFriendlyName() const { return FriendlyName; }

	int32 GetNumPermutations() const { return 1; }

	void SetupResource(
		ERHIFeatureLevel::Type InFeatureLevel,
		FString const& InFriendlyName,
		FString const& InShaderEntryPoint,
		FString const& InShaderHashKey,
		FString const& InShaderSource,
		TUniquePtr<FNNERuntimeIREEShaderParametersMetadataAllocations> InShaderParameterMetadataAllocations,
		FShaderParametersMetadata* InShaderParameterMetadata,
		FName const& InAssetPath,
		TConstArrayView<uint32> InBufferBindings
	);

	TShaderRef<FNNERuntimeIREEShader> GetShader(int32 PermutationId) const;

	bool IsSame(const FNNERuntimeIREEShaderMapId& InId) const;

	uint32 GetBindingIndex(uint32 BufferIdx) const;

protected:
#if WITH_EDITOR
	/**
	 * Fills the passed array with IDs of shader maps unfinished compilation jobs.
	 */
	void GetShaderMapIDsWithUnfinishedCompilation(TArray<int32>& OutShaderMapIds);
#endif // WITH_EDITOR

private:
	FNNERuntimeIREEShaderCompileResults CompilationResults;

	/** 
	 * Game thread tracked shader map, which is ref counted and manages shader map lifetime. 
	 * The shader map uses deferred deletion so that the rendering thread has a chance to process a release command when the shader map is no longer referenced.
	 * Code that sets this is responsible for updating RenderingThreadShaderMap in a thread safe way.
	 * During an async compile, this will be NULL and will not contain the actual shader map until compilation is complete.
	 */
	TRefCountPtr<FNNERuntimeIREEShaderMap> GameThreadShaderMap = nullptr;

	/** 
	 * Shader map for this FComputeKernelResource which is accessible by the rendering thread. 
	 * This must be updated along with GameThreadShaderMap, but on the rendering thread.
	 */
	FNNERuntimeIREEShaderMap* RenderingThreadShaderMap = nullptr;

	/** Name of shader main function. */
	FString ShaderEntryPoint;

	/** Cached shader source for compilation. */
	FString ShaderSource;

	/** 
	 * Hash computed from shader code. If multiple kernels give
	 * the same output, same shader resource will be shared 
	 */
	uint64 ShaderCodeHash = 0;

	TUniquePtr<FNNERuntimeIREEShaderParametersMetadataAllocations> ShaderParameterMetadataAllocations;

	/** Shader parameter metadata. It is expected that object is owned by ShaderParameterMetadataAllocations.  */
	FShaderParametersMetadata const* ShaderParameterMetadata;

	/** Binding indices for storage buffers, needed for IREE command buffer dispatches (buffers might alias and two of them map to the same RDG buffer) */
	TArray<uint32> BufferBindings;

#if WITH_EDITOR
	/**
	 * Compiles this kernel for InPlatform, storing the result in OutShaderMap if the compile was synchronous
	 */
	bool BeginCompileShaderMap(
		const FNNERuntimeIREEShaderMapId& InShaderMapId,
		EShaderPlatform InPlatform, 
		TRefCountPtr<FNNERuntimeIREEShaderMap>& OutShaderMap, 
		bool bApplyCompletedShaderMapForRendering,
		bool bSynchronous = false);
#endif // WITH_EDITOR

	FString FriendlyName;

	/** Asset using this resource */
	FName AssetPath;

	/** 
	 * Contains the compiling id of this shader map when it is being compiled asynchronously. 
	 * This can be used to access the shader map during async compiling, since GameThreadShaderMap will not have been set yet.
	 */
	TArray<int32, TInlineAllocator<1> > OutstandingCompileShaderMapIds;

	/** Feature level that this kernel is representing. */
	ERHIFeatureLevel::Type FeatureLevel = ERHIFeatureLevel::SM5;

	uint32 bLoadedCookedShaderMapId : 1;
	FNNERuntimeIREEShaderMapId CookedShaderMapId;
};

#endif // WITH_NNE_RUNTIME_IREE_SHADER