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

#include "../Common.ush"
#include "../SceneData.ush"
#include "../DeferredShadingCommon.ush"
#include "../SplineMeshCommon.ush"

#include "NaniteDataDecode.ush"
#include "NaniteAttributeDecode.ush"
#include "NaniteVertexFetch.ush"
#include "NaniteEditorCommon.ush"

Texture2D<UlongType> VisBuffer64;
#if MSAA_SAMPLE_COUNT > 1
Texture2DMS<float, MSAA_SAMPLE_COUNT> SceneDepth;
#else
Texture2D<float> SceneDepth;
#endif
Texture2D<uint> ShadingMask;
ByteAddressBuffer DebugViewData;

ByteAddressBuffer ShadingBinData;

// .xy = min, .zw = max
uint4 ViewRect;

float InvShaderBudget;
float3 SelectionColor;
float3 OverlayIntensityColor;

uint DebugViewMode;

#define SOBEL_WIREFRAME_FILTER_ONLY 0

float3 ApplyWireframeFilter(uint2 PixelPosXY, uint DepthInt, float3 WireColor)
{
	// Sobel edge detect depth
	static int SobelX[] =
	{
		1,  0, -1,
		2,  0, -2,
		1,  0, -1
	};

	static int SobelY[] =
	{
		 1,  2,  1,
		 0,  0,  0,
		-1, -2, -1
	};

	static uint2 UVSample[] =
	{
		{-1,  1},  {0,  1},  {1,  1},
		{-1,  0},  {0,  0},  {1,  0},
		{-1, -1},  {0, -1},  {1, -1}
	};

	float2 DepthGrad = 0.0f;
	uint2 BitGrad = 0x88888888;

	uint DepthIntCurrent;
	uint VisibleClusterIndexCurrent;
	uint TriIndexCurrent;

	for (uint Tap = 0; Tap < 9u; ++Tap)
	{
		const UlongType VisPixelCurrent = VisBuffer64[PixelPosXY + UVSample[Tap]];
		UnpackVisPixel(VisPixelCurrent, DepthIntCurrent, VisibleClusterIndexCurrent, TriIndexCurrent);

		float SampleDensityDepth = log2(ConvertFromDeviceZ(asfloat(DepthIntCurrent)) + 1.0f) * 10.0f;

		DepthGrad += float2(SobelX[Tap], SobelY[Tap]) * SampleDensityDepth;

		uint Bits = 0;
		for (uint BitIndex = 0; BitIndex < 8; ++BitIndex)
		{
			Bits |= ((TriIndexCurrent >> BitIndex) & 1u) << (BitIndex * 4u);
		}

		BitGrad.x += SobelX[Tap] * Bits;
		BitGrad.y += SobelY[Tap] * Bits;
	}

	float Wireframe = 0;
	for (uint BitIndex = 0; BitIndex < 8; ++BitIndex)
	{
		float2 Grad = (BitGrad >> (BitIndex * 4u)) & 0xF;
		Wireframe = max(Wireframe, length(Grad - 8u));
	}
	Wireframe *= 0.25f;

	BRANCH
	if (Wireframe == 0.0f)
	{
		discard;
	}

	return saturate(WireColor * Wireframe);
}

struct FNaniteMaterialDebugViewInfo
{
	uint InstructionCountVS;
	uint InstructionCountPS;
	uint InstructionCountCS;
	uint LWCComplexityVS;
	uint LWCComplexityPS;
	uint LWCComplexityCS;
};

FNaniteMaterialDebugViewInfo ReadMaterialDebugViewInfo(uint ShadingBin)
{
	FNaniteMaterialDebugViewInfo Result;

	const uint Stride = MATERIAL_DEBUG_VIEW_INFO_STRIDE;
	uint3 Data = DebugViewData.Load3(ShadingBin * Stride);

	Result.InstructionCountVS	= BitFieldExtractU32(Data.x, 16, 0);
	Result.InstructionCountPS	= BitFieldExtractU32(Data.x, 16, 16);

	Result.InstructionCountCS	= BitFieldExtractU32(Data.y, 16, 0);
	Result.LWCComplexityVS		= BitFieldExtractU32(Data.y, 16, 16);

	Result.LWCComplexityPS		= BitFieldExtractU32(Data.z, 16, 0);
	Result.LWCComplexityCS		= BitFieldExtractU32(Data.z, 16, 16);

	return Result;
}

void ExportDebugViewPS(
	in float4 SvPosition : SV_Position,
	out float4 OutColor  : SV_Target0,
	out float OutDepth   : SV_Depth
)
{
	const uint2 PixelPos = (uint2)SvPosition.xy;
	const UlongType VisPixel = VisBuffer64[PixelPos];

	OutDepth = 0;

	uint DepthInt = 0;
	uint VisibleClusterIndex = 0;
	uint TriIndex = 0;
	UnpackVisPixel(VisPixel, DepthInt, VisibleClusterIndex, TriIndex);

	if (VisibleClusterIndex != 0xFFFFFFFFu)
	{
		const float NaniteDepth = asfloat(DepthInt);
		OutDepth = NaniteDepth;

#if MSAA_SAMPLE_COUNT > 1
		// only the first MSAA depth sample is tested, as this shader runs per-pixel rather than per-sample
		float SceneDepthValue = SceneDepth.Load(PixelPos.xy, 0);
#else
		float SceneDepthValue = SceneDepth[PixelPos.xy];
#endif

		if (NaniteDepth >= SceneDepthValue)
		{
			FNaniteView NaniteView = GetNaniteView(0);

			FVisibleCluster VisibleCluster				= GetVisibleCluster(VisibleClusterIndex);
			FInstanceSceneData InstanceData				= GetInstanceSceneDataUnchecked(VisibleCluster);
			FPrimitiveSceneData PrimitiveData			= GetPrimitiveData(InstanceData.PrimitiveId);
			FInstanceDynamicData InstanceDynamicData	= CalculateInstanceDynamicData(NaniteView, InstanceData);
			FCluster Cluster							= GetCluster(VisibleCluster.PageIndex, VisibleCluster.ClusterIndex);

			FShadingMask UnpackedMask = UnpackShadingMask(ShadingMask[PixelPos]);
			FNaniteMaterialDebugViewInfo MaterialDebugViewInfo = ReadMaterialDebugViewInfo(UnpackedMask.ShadingBin);

			BRANCH
			if (DebugViewMode == DEBUG_VIEW_WIREFRAME)
			{
				const bool bIsSelected = IsInstanceSelected(InstanceData, VisibleCluster, TriIndex);
				const float3 WireColor = CondMask(bIsSelected, SelectionColor, PrimitiveData.WireframeColor);
				//const float3 WireColor = PrimitiveData.PrimitiveColor;

			#if SOBEL_WIREFRAME_FILTER_ONLY
				const bool bUseSobelFilter = true;
			#else
				// If the visible cluster has WPO, we can't accurately calculate the screen positions of the triangles
				// and therefore cannot compute the barycentrics. So in this case, fall back to the Sobel filter
				const bool bUseSobelFilter = (VisibleCluster.Flags & NANITE_CULLING_FLAG_ENABLE_WPO) != 0;
			#endif

				BRANCH
				if (bUseSobelFilter)
				{
					OutColor = float4(ApplyWireframeFilter(PixelPos, DepthInt, WireColor), 1.0f);
				}
				else
				{
					const uint3 TriIndices = DecodeTriangleIndices(Cluster, TriIndex);

					float3 PointLocal[3];
					FetchLocalNaniteTrianglePositions(InstanceData, Cluster, VisibleCluster, TriIndices, PointLocal);

					// We can handle spline mesh deformation here because its deformation is fixed function
					const uint PackedFlags = ShadingBinData.Load<FNaniteShadingBinMeta>(UnpackedMask.ShadingBin * NANITE_SHADING_BIN_META_BYTES).MaterialFlags;
					const FNaniteMaterialFlags MaterialFlags = UnpackNaniteMaterialFlags(PackedFlags);

					BRANCH
					if (MaterialFlags.bSplineMesh)
					{
						FSplineMeshShaderParams SplineMeshParams = SplineMeshLoadParamsFromInstancePayload(InstanceData);
						PointLocal[0] = SplineMeshDeformLocalPos(SplineMeshParams, PointLocal[0]);
						PointLocal[1] = SplineMeshDeformLocalPos(SplineMeshParams, PointLocal[1]);
						PointLocal[2] = SplineMeshDeformLocalPos(SplineMeshParams, PointLocal[2]);
					}

					const float3 PointWorld0 = mul(float4(PointLocal[0], 1), InstanceDynamicData.LocalToTranslatedWorld).xyz;
					const float3 PointWorld1 = mul(float4(PointLocal[1], 1), InstanceDynamicData.LocalToTranslatedWorld).xyz;
					const float3 PointWorld2 = mul(float4(PointLocal[2], 1), InstanceDynamicData.LocalToTranslatedWorld).xyz;

					const float4 PointClip0 = mul(float4(PointWorld0, 1), NaniteView.TranslatedWorldToClip);
					const float4 PointClip1 = mul(float4(PointWorld1, 1), NaniteView.TranslatedWorldToClip);
					const float4 PointClip2 = mul(float4(PointWorld2, 1), NaniteView.TranslatedWorldToClip);

					const float2 PixelClip = (PixelPos + 0.5 - View.ViewRectMin.xy) * View.ViewSizeAndInvSize.zw * float2(2, -2) + float2(-1, 1);

					// Calculate perspective correct barycentric coordinates with screen derivatives
					const FBarycentrics Barycentrics = CalculateTriangleBarycentrics(PixelClip, PointClip0, PointClip1, PointClip2, View.ViewSizeAndInvSize.zw);
					
					float3 Deltas = abs(Barycentrics.Value_dx) + abs(Barycentrics.Value_dy);

					float SmoothingVal = 1.0f; // 0-10
					float ThicknessVal = 0.05f; // 0-10

					float3 Smoothing = Deltas * SmoothingVal;
					float3 Thickness = Deltas * ThicknessVal;

					const float3 SmoothBarycentrics = smoothstep(Thickness, Thickness + Smoothing, Barycentrics.Value);

					const float MinBarycentrics = min3(SmoothBarycentrics.x, SmoothBarycentrics.y, SmoothBarycentrics.z);
					if (MinBarycentrics == 1.0f)
					{
						discard;
					}

					OutColor = float4(WireColor, 1.0f - MinBarycentrics);
				}
			}
		#if USE_EDITOR_SHADERS
			else if (DebugViewMode == DEBUG_VIEW_SHADER_COMPLEXITY)
			{
				// If mode is DVSM_QuadComplexity, different calculations used
				// See: FComplexityAccumulateInterface::GetDebugViewModeShaderBindings()

				// float3(PS/ShaderBudget, VS/ShaderBudget, overdraw)
				float3 NormalizedComplexity;
				if (MaterialDebugViewInfo.InstructionCountCS != 0u)
				{
					NormalizedComplexity = float3(
						MaterialDebugViewInfo.InstructionCountCS * InvShaderBudget,
						MaterialDebugViewInfo.InstructionCountCS * InvShaderBudget,
						1.0 / 32.0f
					);
				}
				else
				{
					NormalizedComplexity = float3(
						MaterialDebugViewInfo.InstructionCountPS * InvShaderBudget,
						MaterialDebugViewInfo.InstructionCountVS * InvShaderBudget,
						1.0 / 32.0f
					);
				}

				float3 FinalComplexity = NormalizedComplexity.xyz;

				//BRANCH
				//if (bShowQuadOverdraw && NormalizedComplexity.x > 0.0f)
				//{
				//	uint Coverage = ComputeQuadCoverage(Input.SvPosition.xy, Input.PrimitiveID, 24, false, false, 0);
				//	// The extra cost from quad overdraw is assumed to be linear.
				//	FinalComplexity.x *= 4.f / (float)(Coverage);
				//}

				// Use the maximum range allowed for scene color
				// Alpha channel needs to be 1 to have decals working where the framebuffer blend is setup to multiply with alpha
				OutColor = float4(FinalComplexity.xyz, 1.0f);
			}
			else if (DebugViewMode == DEBUG_VIEW_LWC_COMPLEXITY)
			{
				float3 NormalizedComplexity;
				if (MaterialDebugViewInfo.InstructionCountCS != 0u)
				{
					NormalizedComplexity = float3(MaterialDebugViewInfo.InstructionCountCS, MaterialDebugViewInfo.InstructionCountCS, 0) * InvShaderBudget;
				}
				else
				{
					NormalizedComplexity = float3(MaterialDebugViewInfo.LWCComplexityPS, MaterialDebugViewInfo.LWCComplexityVS, 0) * InvShaderBudget;
				}

				// Use the maximum range allowed for scene color
				// Alpha channel needs to be 1 to have decals working where the framebuffer blend is setup to multiply with alpha
				OutColor = float4(NormalizedComplexity.xyz, 1.0f);
			}
			else if (DebugViewMode == DEBUG_VIEW_PRIMITIVE_COLOR)
			{
				const bool bIsSelected = IsInstanceSelected(InstanceData, VisibleCluster, TriIndex);
				const float3 SolidColor = CondMask(bIsSelected, SelectionColor, OverlayIntensityColor * PrimitiveData.PrimitiveColor);
				OutColor = float4(SolidColor, 1.0f);
			}
		#endif
			else
			{
				// Invalid debug mode mode
				discard;
			}
		}
		else
		{
			// Nanite is behind scene depth
			discard;
		}
	}
	else
	{
		// Not a Nanite pixel
		discard;
	}
}