// ExampleComputeShaderModule.h
#include "Modules/ModuleManager.h"

class FExampleComputeShaderModule : public IModuleInterface {
    virtual void StartupModule() override;
    // virtual void ShutdownModule() override;
};

// ExampleComputeShader.usf
#include "/Engine/Public/Platform.ush"

float Scale;
float Translate;

RWStructuredBuffer<float> InputBuffer;
RWStructuredBuffer<float> OutputBuffer;

[numthreads(1, 1, 1)]
void FunctionMultiply(
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint GroupIndex : SV_GroupIndex)
{
    OutputBuffer[DispatchThreadId.x] = InputBuffer[DispatchThreadId.x] * Scale + Translate;
}

// ExampleComputeShaderModule.cpp

/**
 * FGlobalShader是定义在全局上的，Unreal也有FMaterialShader，目前我们不做讨论。
 * FGlobalShader的初始化将由Unreal定义的宏进行管理，我们可以使用宏方便地初始化Global Shader
 * 需要注意的是，FGlobalShader不能出现成员变量。我们会用其他的方法设置Shader参数
 */
class FExampleComputeShaderCS : public FGlobalShader
{
public:
	DECLARE_SHADER_TYPE(FExampleComputeShaderCS, Global) // 定义一堆函数，我们不用管是什么
};

IMPLEMENT_SHADER_TYPE(, FExampleComputeShaderCS, TEXT("/MyShaders/ExampleComputeShader.usf"), TEXT("FunctionMultiply"), SF_Compute)
// 使用宏自动实现FExampleComputeShaderCS的函数。第一个参数可为空，第二个参数是类名，第三个是usf的path，第四个是usf的入口函数名，第五个是shader类型(SF_Vertex, SF_Pixel, SF_Compute等)。

// 添加头文件
#include "ShaderParameterStruct.h"

// 使用ShaderParameterStruct.h内的宏定义Shader参数结构体。
// **注意** 这里每个变量名字必须与usf文件中Shader参数名一致。
BEGIN_SHADER_PARAMETER_STRUCT(FExampleComputeShaderParameters,)
SHADER_PARAMETER(float, Scale)     // 基本类型的shader参数
SHADER_PARAMETER(float, Translate) // 基本类型的shader参数
SHADER_PARAMETER_UAV(RWStructuredBuffer<float>, InputBuffer)  // UAV Buffer类型的shader参数，需要我们后续手动管理
SHADER_PARAMETER_UAV(RWStructuredBuffer<float>, OutputBuffer) // UAV Buffer类型的shader参数，需要我们后续手动管理
END_SHADER_PARAMETER_STRUCT()


/**
 * FGlobalShader是定义在全局上的，Unreal也有FMaterialShader，目前我们不做讨论。
 * FGlobalShader的初始化将由Unreal定义的宏进行管理，我们可以使用宏方便地初始化Global Shader
 * 需要注意的是，FGlobalShader不能出现成员变量。我们会用其他的方法设置Shader参数
 */
class FExampleComputeShaderCS : public FGlobalShader
{
public:
	DECLARE_SHADER_TYPE(FExampleComputeShaderCS, Global) // 定义一堆函数，我们不用管是什么
	SHADER_USE_PARAMETER_STRUCT(FExampleComputeShaderCS, FGlobalShader)  // 定义该Shader使用SHADER_PARAMETER_STRUCT来定义Shader参数

	using FParameters = FExampleComputeShaderParameters; // SHADER_USE_PARAMETER_STRUCT需要我们定义FParameters，我们在类外定义的，所以只需要using就可以。
	//当然也可以直接在FGlobalShader内部定义BEGIN_SHADER_PARAMETER_STRUCT(FParameters,) 两种方法使用一个即可
};

IMPLEMENT_SHADER_TYPE(, FExampleComputeShaderCS, TEXT("/MyShaders/ExampleComputeShader.usf"), TEXT("FunctionMultiply"), SF_Compute)
// 使用宏自动实现FExampleComputeShaderCS的函数。第一个参数可为空，第二个参数是类名，第三个是usf的path，第四个是usf的入口函数名，第五个是shader类型(SF_Vertex, SF_Pixel, SF_Compute等)。

// ExampleComputeShaderModule.h

class EXAMPLECOMPUTESHADER_API FExampleComputeShaderResource : public FRenderResource // 使用EXAMPLECOMPUTESHADER_API因为我们要在别的模块访问到这个类
{
	static FExampleComputeShaderResource* GInstance; // 单例模式
public:
	FRWBufferStructured InputBuffer;  // RWStructuredBuffer<float> InputBuffer;
	FRWBufferStructured OutputBuffer; // RWStructuredBuffer<float> OutputBuffer;

	// 初始化所有buffer，override基类
	// 此函数由FRenderResource::InitResource(FRHICommandList&)调用
	virtual void InitRHI(FRHICommandListBase& RHICmdList) override;

	
	// 此函数由FRenderResource::ReleaseResource(FRHICommandList&)调用
	virtual void ReleaseRHI() override; // 释放所有buffer

	static FExampleComputeShaderResource* Get(); // 单例模式
private:
	FExampleComputeShaderResource() {} // 构造函数私有
};

// ExampleComputeShaderModule.cpp

// 初始化单例指针
FExampleComputeShaderResource* FExampleComputeShaderResource::GInstance = nullptr;

FExampleComputeShaderResource* FExampleComputeShaderResource::Get()
{
	if(GInstance==nullptr)
	{
		GInstance = new FExampleComputeShaderResource();
		// 创建执行在RenderThread的任务。第一个括号写全局唯一的标识符，一般为这个Task起个名字。第二个写Lambda表达式，用FRHICommandList& RHICmdList作为函数参数。
		ENQUEUE_RENDER_COMMAND(FInitExampleComputeShaderResource)([](FRHICommandList& RHICmdList)
		{
			GInstance->InitResource(RHICmdList);
		});
	}
	return GInstance;
}

// Buffer初始化函数
void FExampleComputeShaderResource::InitRHI(FRHICommandListBase& RHICmdList)
{
	InputBuffer.Initialize(RHICmdList, TEXT("InputBuffer"), sizeof(float), 1);
	OutputBuffer.Initialize(RHICmdList, TEXT("OutputBuffer"), sizeof(float), 1);
}

// Buffer释放函数
void FExampleComputeShaderResource::ReleaseRHI()
{
	InputBuffer.Release();
	OutputBuffer.Release();
}

// ExampleComputeShaderModule.cpp
void DispatchExampleComputeShader_RenderThread(FRHICommandList& RHICmdList, FExampleComputeShaderResource* Resource, float Scale, float Translate, uint32 ThreadGroupX, uint32 ThreadGroupY, uint32 ThreadGroupZ)
{
	TShaderMapRef<FExampleComputeShaderCS> Shader(GetGlobalShaderMap(GMaxRHIFeatureLevel)); // 声明ShaderMap
	SetComputePipelineState(RHICmdList, Shader.GetComputeShader()); // 设置该RHICmdList的Pipeline为ComputeShader
	{
		typename FExampleComputeShaderCS::FParameters Parameters{}; // 创建Shader参数

		// 设置基本类型
		Parameters.Scale = Scale;
		Parameters.Translate = Translate;

		// 设置buffer类型
		Parameters.InputBuffer = Resource->InputBuffer.UAV;
		Parameters.OutputBuffer = Resource->OutputBuffer.UAV;

		// 传入参数
		SetShaderParameters(RHICmdList, Shader, Shader.GetComputeShader(), Parameters);
	}

	// 调用Compute shader
	DispatchComputeShader(RHICmdList, Shader.GetShader(), ThreadGroupX, ThreadGroupY, ThreadGroupZ);

	// 取消绑定buffer参数
	//UnsetShaderSRVs(RHICmdList, Shader, Shader.GetComputeShader());
	UnsetShaderUAVs(RHICmdList, Shader, Shader.GetComputeShader());
}

void DispatchExampleComputeShader_GameThread(float InputVal, float Scale, float Translate, FExampleComputeShaderResource* Resource)
{
	// 加入RenderThread任务
	ENQUEUE_RENDER_COMMAND(FDispatchExampleComputeShader)([Resource, InputVal, Scale, Translate](FRHICommandListImmediate& RHICmdList)
	{
		// LockBuffer并写入数据
		float* InputGPUBuffer = static_cast<float*>(RHICmdList.LockBuffer(Resource->InputBuffer.Buffer, 0, sizeof(float), RLM_WriteOnly));
		*InputGPUBuffer = InputVal; // 可以使用FMemory::Memcpy
		// UnlockBuffer
		RHICmdList.UnlockBuffer(Resource->InputBuffer.Buffer);
		// 调用RenderThread版本的函数
		DispatchExampleComputeShader_RenderThread(RHICmdList, Resource, Scale, Translate, 1, 1, 1);
	});
}

float GetGPUReadback(FExampleComputeShaderResource* Resource, float& OutputVal)
{
	float* pOutputVal = &OutputVal;
	// Flush所有RenderingCommands, 确保我们的shader已经执行了
	FlushRenderingCommands();
	ENQUEUE_RENDER_COMMAND(FReadbackOutputBuffer)([Resource, &pOutputVal](FRHICommandListImmediate& RHICmdList)
	{
		// LockBuffer并读取数据
		float* OutputGPUBuffer = static_cast<float*>(RHICmdList.LockBuffer(Resource->OutputBuffer.Buffer, 0, sizeof(float), RLM_ReadOnly));
		*pOutputVal = *OutputGPUBuffer;
		// UnlockBuffer
		RHICmdList.UnlockBuffer(Resource->OutputBuffer.Buffer);
	});
	// FlushRenderingCommands, 确保上面的RenderCommand被执行了
	FlushRenderingCommands();
	return OutputVal;

	// 下面是使用FRHIGPUBufferReadback的调用方式, 效果是相同的
	//FRHIGPUBufferReadback ReadbackBuffer(TEXT("ExampleComputeShaderReadback"));
	//FlushRenderingCommands();
	//ENQUEUE_RENDER_COMMAND(FEnqueueGPUReadback)([&ReadbackBuffer, Resource](FRHICommandListImmediate& RHICmdList)
	//{
	//	ReadbackBuffer.EnqueueCopy(RHICmdList, Resource->OutputBuffer.Buffer, sizeof(int32));
	//});
	//FlushRenderingCommands();
	//float OutputVal;
	//ENQUEUE_RENDER_COMMAND(FEnqueueGPUReadbackLock)([&ReadbackBuffer, Resource, &OutputVal](FRHICommandListImmediate& RHICmdList)
	//{
	//	float* pOutputVal = static_cast<float*>(ReadbackBuffer.Lock(sizeof(float)));
	//	checkf(pOutputVal!=nullptr, TEXT("ReadbackBuffer failed to lock"));
	//	OutputVal = *pOutputVal;
	//	ReadbackBuffer.Unlock();
	//});
	//FlushRenderingCommands();
	//return OutputVal;
}

// ShaderFunctionLibraryModule.h
#pragma once

#include "CoreMinimal.h"
#include "Kismet/BlueprintFunctionLibrary.h"
#include "ExampleComputeShaderModule.h"

#include "ShaderFunctionLibraryModule.generated.h"

class FShaderFunctionLibraryModule : public IModuleInterface {
    virtual void StartupModule() override;
    // virtual void ShutdownModule() override;
};

UCLASS(meta=(ScriptName="ShaderFunctionLibrary"), MinimalAPI)
class UShaderFunctionLibrary : public UBlueprintFunctionLibrary
{
	GENERATED_BODY()
public:
	UFUNCTION(BlueprintCallable, meta=(DisplayName="Execute ExampleComputeShader"), Category="My Shader Functions")
	static SHADERFUNCTIONLIBRARY_API float ExecuteExampleComputeShader(float InputVal, float Scale, float Translate) {
		float OutputVal;
		DispatchExampleComputeShader_GameThread(InputVal, Scale, Translate, FExampleComputeShaderResource::Get());
		return GetGPUReadback(FExampleComputeShaderResource::Get(), OutputVal);
	}
};

// ShaderFunctionLibraryModule.cpp
#include "ShaderFunctionLibraryModule.h"

IMPLEMENT_MODULE(FShaderFunctionLibraryModule, ShaderFunctionLibrary)
void FShaderFunctionLibraryModule::StartupModule()
{ }

// ExampleGrapihcsShader.usf
#include "/Engine/Public/Platform.ush"
// HLSL代码的文档见https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-reference


// 顶点着色器的输入
struct VertexAttributes
{
	float4 v_position : ATTRIBUTE0; // 顶点位置。
	//通常情况下我们输入模型坐标系的顶点，经由view matrix和perspective projection matrix变换得到NDC。但为了演示方便起见我们直接输入NDC。
	float4 v_color : ATTRIBUTE1; // 顶点颜色，取值范围为0.0~1.0
};

// 顶点着色器的输出，同时也是pixel shader的interpolated输入
struct Varyings
{
	float4 p_position : SV_POSITION; // 顶点着色器输出，需要在NDC(Normalized device coordinate)坐标空间下; 片元着色器(pixel shader)的坐标位置输入
	float4 p_color : COLOR0; // 顶点着色器输出; 片元颜色输入
};

void MainVS(in VertexAttributes vertex, out Varyings v)
{
	v.p_position = vertex.v_position; // 此处由于输入的顶点坐标即是NDC，所以不作变换
	v.p_color = vertex.v_color;
}

void MainPS(in Varyings v, out float4 f_color : SV_Target0)
{
	f_color = v.p_color; // 每个片元的颜色输出到render target
}

// ExampleGraphicsShaderModule.h
#pragma once

#include "Modules/ModuleManager.h"

class FExampleGraphicsShaderModule : public IModuleInterface {
    virtual void StartupModule() override;
    // virtual void ShutdownModule() override;
};


// 这个struct对应hlsl里面的VertexAttributes
EXAMPLEGRAPHICSSHADER_API struct VertexAttributes
{
	FVector4f Position; // Normally we use homogeneous coordinate so we declared Vec4 but here for demonstration we only use first 2 components to store NDC
	FVector4f Color; // RGBA
};


// 存VS和PS对应的渲染资源
class EXAMPLEGRAPHICSSHADER_API FExampleGraphicsShaderResource : public FRenderResource
{
	static FExampleGraphicsShaderResource* GInstance; // Singleton instance
public:
	virtual void InitRHI(FRHICommandListBase& RHICmdList) override;
	virtual void ReleaseRHI() override;
	
	FVertexDeclarationRHIRef VertexDeclarationRHI; // 定义vertex的数据是如何存储在buffer的。见InitRHI()
	FReadBuffer VertexBuffer; // GPU只读不写，所以定义为ReadBuffer
	static FExampleGraphicsShaderResource* Get(); // Singleton instance
};

EXAMPLEGRAPHICSSHADER_API void RenderExampleGraphicsShader_GameThread(UTextureRenderTarget2D* TextureRenderTarget2D, FExampleGraphicsShaderResource* Resource);

// ExampleGraphicsShaderModule.cpp
#include "ExampleGraphicsShaderModule.h"

#include "ClearQuad.h"
#include "SelectionSet.h"
#include "Engine/TextureRenderTarget2D.h"

IMPLEMENT_MODULE(FExampleGraphicsShaderModule, ExampleGraphicsShader)

void FExampleGraphicsShaderModule::StartupModule()
{
	// 在模块初始化后将项目目录的Shaders文件夹Mount到/MyGraphicsShader
	AddShaderSourceDirectoryMapping(TEXT("/MyGraphicsShader"), FPaths::ProjectDir() / TEXT("Shaders"));
}

class FExampleGraphcisShaderVS: public FGlobalShader
{
public:
	DECLARE_EXPORTED_GLOBAL_SHADER(FExampleGraphcisShaderVS, EXAMPLEGRAPHICSSHADER_API);
	static bool ShouldCache(EShaderPlatform Platform)
	{
		return true;
	}
};

class FExampleGraphcisShaderPS: public FGlobalShader
{
public:
	DECLARE_EXPORTED_GLOBAL_SHADER(FExampleGraphcisShaderPS, EXAMPLEGRAPHICSSHADER_API);
	static bool ShouldCache(EShaderPlatform Platform)
	{
		return true;
	}
};


IMPLEMENT_SHADER_TYPE(, FExampleGraphcisShaderVS, TEXT("/MyGraphicsShader/ExampleGraphicsShader.usf"), TEXT("MainVS"), SF_Vertex);
IMPLEMENT_SHADER_TYPE(, FExampleGraphcisShaderPS, TEXT("/MyGraphicsShader/ExampleGraphicsShader.usf"), TEXT("MainPS"), SF_Pixel);



FExampleGraphicsShaderResource*  FExampleGraphicsShaderResource::GInstance = nullptr;
FExampleGraphicsShaderResource* FExampleGraphicsShaderResource::Get()
{
	if(GInstance==nullptr)
	{
		GInstance = new FExampleGraphicsShaderResource();
		ENQUEUE_RENDER_COMMAND(FInitExampleGraphicsShaderResource)([](FRHICommandList& RHICmdList)
		{
			GInstance->InitResource(RHICmdList);
		});
	}
	return GInstance;
}

void FExampleGraphicsShaderResource::InitRHI(FRHICommandListBase& RHICmdList)
{
	// 我们先hard code顶点数据
	TArray<VertexAttributes> Vertices;
	Vertices.Add({FVector4f(-1, 1, 1, 1), FVector4f(1,0,0,1)});
	Vertices.Add({FVector4f(1, 1,1, 1), FVector4f(0,0,1,1)});
	Vertices.Add({FVector4f(-1,-1,1, 1), FVector4f(1,1,0,1)});

	// 初始化buffer并拷贝
	uint32 NumBytes = sizeof(VertexAttributes)* Vertices.Num();
	FRHIResourceCreateInfo CreateInfo(TEXT("VertexBuffer"));
	VertexBuffer.Buffer = RHICmdList.CreateVertexBuffer(NumBytes, BUF_Static, CreateInfo);
	VertexAttributes* GPUBufferPtr = static_cast<VertexAttributes*>(RHICmdList.LockBuffer(VertexBuffer.Buffer, 0, sizeof(VertexAttributes) * Vertices.Num(), RLM_WriteOnly));
	FMemory::Memcpy(GPUBufferPtr, Vertices.GetData(), NumBytes);
	RHICmdList.UnlockBuffer(VertexBuffer.Buffer);
	// 到此为止，我们的buffer存储了3个vertices，每个vertex有4个float作为position和4个float作为color。所以一个vertex是4*sizeof(float) + 4*sizeof(float)==32字节。当然3个vertices总共96个字节。

	// 下面我们定义vertex buffer的数据排布。

	// 定义VertexDeclaration
	uint16 Stride = sizeof(VertexAttributes);
	FVertexDeclarationElementList Elements;
	Elements.Add(FVertexElement(0, STRUCT_OFFSET(VertexAttributes, Position),VET_Float4,0, Stride));
	Elements.Add(FVertexElement(0, STRUCT_OFFSET(VertexAttributes, Color), VET_Float4, 1, Stride));
	VertexDeclarationRHI = PipelineStateCache::GetOrCreateVertexDeclaration(Elements);
	// 根据vertex的定义，前四个float是position，并且我们在hlsl将其定义为ATTRIBUTE0。所以它的InOffset应该是0（因为是结构体的第一个成员）；stride应该是sizeof(VertexAttributes)因为当前指针到下一个vertex指针的偏移量是sizeof(VertexAttributes)
	// Color同理，但它的InOffset是Position的大小，即4*sizeof(float)
	// 如果读者熟悉opengl API，FVertexElement与glVertexAttribPointer相似
	
}

void FExampleGraphicsShaderResource::ReleaseRHI()
{
	if(VertexBuffer.Buffer)
	{
		VertexBuffer.Release();
	}
	VertexDeclarationRHI.SafeRelease();
}

void RenderExampleGraphicsShader_RenderThread(FRHICommandList& RHICmdList, FExampleGraphicsShaderResource* Resource, FRHITexture* RenderTarget)
{
	// 调用shader进行渲染
	// 与compute shader不同，graphics shader需要我们定义：1）一个或多个Render Target；2）如何处理Rasterizer、Blend、Depth和Stencil的行为
	// 所以需要给到的初始化参数比compute shader多
	// 一个很好的例子是DrawClearQuad函数，里面详细写了从RHICmdList.BeginRenderPass到RHICmdList.EndRenderPass的所有步骤。如果你的shader无法工作，请使用DrawClearQuad函数进行测试。
	FRHIRenderPassInfo RPInfo(RenderTarget, ERenderTargetActions::Clear_Store);
	RHICmdList.BeginRenderPass(RPInfo, TEXT("ExampleGraphicsShaderRenderPass"));
	// 如果需要测试，可以在这里调用DrawClearQuad(RHICmdList, true, FLinearColor(FVector4(1, 0, 1, 1)), true, 1.0, true, 0);
	// 同时注释其他代码。
	auto ShaderMap = GetGlobalShaderMap(GMaxRHIFeatureLevel);

	TShaderMapRef<FExampleGraphcisShaderVS> VertexShader(ShaderMap);
	TShaderMapRef<FExampleGraphcisShaderPS> PixelShader(ShaderMap);

	FGraphicsPipelineStateInitializer GraphicsPSOInit;
	RHICmdList.ApplyCachedRenderTargets(GraphicsPSOInit);
	GraphicsPSOInit.RasterizerState = TStaticRasterizerState<>::GetRHI();
	GraphicsPSOInit.BlendState = TStaticBlendState<>::GetRHI();
	GraphicsPSOInit.DepthStencilState = TStaticDepthStencilState<true, CF_Always>::GetRHI();
	GraphicsPSOInit.PrimitiveType = PT_TriangleStrip;
	GraphicsPSOInit.BoundShaderState.VertexShaderRHI = VertexShader.GetVertexShader();
	GraphicsPSOInit.BoundShaderState.PixelShaderRHI = PixelShader.GetPixelShader();
	GraphicsPSOInit.BoundShaderState.VertexDeclarationRHI = Resource->VertexDeclarationRHI;

	SetGraphicsPipelineState(RHICmdList, GraphicsPSOInit, 0);
	RHICmdList.SetStreamSource(0, Resource->VertexBuffer.Buffer, 0);
	RHICmdList.DrawPrimitive(0, 1, 1);
	RHICmdList.EndRenderPass();
}

void RenderExampleGraphicsShader_GameThread(UTextureRenderTarget2D* TextureRenderTarget2D, FExampleGraphicsShaderResource* Resource)
{
	ENQUEUE_RENDER_COMMAND(FRenderExampleGraphicsShader)([Resource, TextureRenderTarget2D](FRHICommandList& RHICmdList)
	{
		RenderExampleGraphicsShader_RenderThread(RHICmdList, Resource, TextureRenderTarget2D->GetResource()->GetTexture2DRHI());
	});
}

// ShaderFunctionLibraryModule.h

UFUNCTION(BlueprintCallable, meta=(DisplayName="Render ExampleGraphicsShader"),Category="My Shader Functions")
static SHADERFUNCTIONLIBRARY_API void RenderExampleGraphicsShader(UTextureRenderTarget2D* RenderTarget)
{
	FExampleGraphicsShaderResource::Get();
	RenderExampleGraphicsShader_GameThread(RenderTarget, FExampleGraphicsShaderResource::Get());
}