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

#include "/Engine/Shared/MaterialCacheDefinitions.h"
#include "/Engine/Private/Common.ush"
#include "/Engine/Private/Nanite/NaniteShadeCommon.ush"
#include "/Engine/Private/MortonCode.ush"
#include "/Engine/Public/RootConstants.ush"

#define SceneTexturesStruct MaterialCachePass.SceneTextures 
#define EyeAdaptationStruct MaterialCachePass

#include "/Engine/Generated/Material.ush"
#include "/Engine/Generated/VertexFactory.ush"
#include "/Engine/Generated/UniformBuffers/MaterialCachePass.ush"
#include "/Engine/Private/MaterialCache/MaterialCacheShadeCommon.ush"
#include "/Engine/Private/MaterialCache/MaterialCacheCommon.ush"

COMPILER_ALLOW_CS_DERIVATIVES

Buffer<uint> PageIndirections;

FMaterialCacheABufferTag ShadeAttributes(in FPackedNaniteView PackedView, uint2 VisBufferPixelPos, uint ViewIndex)
{
	ResolvedView = ResolveView();

	FVertexFactoryInterpolantsVSToPS Interpolants = (FVertexFactoryInterpolantsVSToPS)0;
	Interpolants.PixelPos  = VisBufferPixelPos;
	Interpolants.ViewIndex = ViewIndex;

	FNaniteView NaniteView = UnpackNaniteView(PackedView);
	PatchViewState(NaniteView, ResolvedView);

	float4 SvPosition = float4(VisBufferPixelPos + 0.5f.xx, 0.0f, 1.0f);

	FMaterialPixelParameters MaterialParameters = GetMaterialPixelParameters(NaniteView, Interpolants, SvPosition);

	const float4 ScreenPosition = SvPositionToResolvedScreenPosition(SvPosition);
	
	FPixelMaterialInputs PixelMaterialInputs;
	CalcMaterialParametersEx(
		MaterialParameters,
		PixelMaterialInputs,
		SvPosition,
		ScreenPosition,
		true,
		MaterialParameters.WorldPosition_CamRelative,
		MaterialParameters.WorldPosition_NoOffsets_CamRelative
	);

	return GetMaterialABuffer(MaterialParameters, PixelMaterialInputs);
}

bool FindNearestInfillVisPixel(inout uint2 VisBufferPixelPos, out uint DepthInt, out uint VisibleClusterIndex, out uint VisTriIndex)
{
	static const uint TriangularIndices[] = { 0, 1, 1, 2, 2, 2, 3, 3, 3, 3 };
	
	const uint BorderWidth = min(MaterialCachePass.TileOrderingParams.w, 4);
	const uint SampleCount = BorderWidth * (BorderWidth + 1) >> 1;

	// A somewhat naive infilling, find the nearest vis-sample and repeat its shading. While not entirely accurate, it does avoid a second pass.
	// TODO[MP]: Likely have to check against the page bounds, otherwise might pull in irrelevant triangles.
	for (uint i = 0; i < SampleCount; i++)
	{
		const uint TriIndex       = TriangularIndices[i];
		const uint TriSampleIndex = i - (TriIndex * (TriIndex + 1) >> 1);
		const int2 TriSample      = int2(TriSampleIndex, TriIndex - TriSampleIndex);

		const int2 OffsetSwizzles[] = {
			int2( TriSample.x + 1,  TriSample.y    ),
			int2( TriSample.x    , -TriSample.y - 1),
			int2(-TriSample.x - 1, -TriSample.y - 1),
			int2(-TriSample.x - 1,  TriSample.y    ),
		};

		UNROLL_N(4)
		for (uint j = 0; j < 4; j++)
		{
			UnpackVisPixel(NaniteShading.VisBuffer64[VisBufferPixelPos + OffsetSwizzles[j]], DepthInt, VisibleClusterIndex, VisTriIndex);
			if (VisibleClusterIndex != 0xFFFFFFFF)
			{
				VisBufferPixelPos += OffsetSwizzles[j];
				return true;
			}
		}
	}

	return false;
}

[numthreads(64, 1, 1)]
void Main(uint3 DTid : SV_DispatchThreadID)
{
	const uint PageIndex     = PageIndirections[GetRootConstant0() + DTid.y];
	const uint ShadingBin	 = GetRootConstant1();
	const uint MeshPassIndex = GetRootConstant2();

	const EMaterialCacheFlag    Flags      = (EMaterialCacheFlag)GetRootConstant3();
	const FPackedNaniteView     PackedView = NaniteShading.InViews[PageIndex];
	const FMaterialCacheBinData BinData    = NaniteShading.ShadingBinData.Load<FMaterialCacheBinData>(sizeof(FMaterialCacheBinData) * PageIndex);
	
	const uint2 PagePixelPos = GetMaterialCacheCSShadePixel(DTid, MaterialCachePass.TileParams, MaterialCachePass.TileOrderingParams);
	uint2 VisBufferPixelPos  = PackedView.ViewRect.xy + PagePixelPos;
	
	uint DepthInt = 0;
	uint VisibleClusterIndex = 0;
	uint TriIndex = 0;
	UnpackVisPixel(NaniteShading.VisBuffer64[VisBufferPixelPos], DepthInt, VisibleClusterIndex, TriIndex);

	// If we're at an empty pixel, try to find the nearest valid pixel for edge in-fill
	if (VisibleClusterIndex == 0xFFFFFFFF)
	{
		if (!FindNearestInfillVisPixel(VisBufferPixelPos, DepthInt, VisibleClusterIndex, TriIndex))
		{
			return;
		}
	}

	FVisibleCluster    VisibleCluster = GetVisibleCluster(VisibleClusterIndex);
	FInstanceSceneData InstanceData   = GetInstanceSceneDataUnchecked(VisibleCluster);
	FCluster           ClusterData    = GetCluster(VisibleCluster.PageIndex, VisibleCluster.ClusterIndex);

	const uint PixelShadingBin = GetMaterialShadingBin(ClusterData, InstanceData.PrimitiveId, MeshPassIndex, TriIndex);
	if (PixelShadingBin != ShadingBin)
	{
		return;
	}
	
	// Shader against the current layer
	FMaterialCacheABufferTag Top = ShadeAttributes(PackedView, VisBufferPixelPos, PageIndex);
	if (Top.Clipped)
	{
		return;
	}

	StoreMaterialABufferPixel(Top, Flags, BinData.ABufferPhysicalPosition, PagePixelPos);
}