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

#ifndef THREADGROUP_SIZE
	#define THREADGROUP_SIZE 1
#endif

#include "../Common.ush"
#include "LumenCardCommon.ush"
#include "LumenCardTile.ush"
#define SUPPORT_CONTACT_SHADOWS 0
#include "../DeferredLightingCommon.ush"
#include "../VolumeLightingCommon.ush"
#include "../ForwardShadowingCommon.ush"
#include "../LightGridCommon.ush"
#define DISTANCE_FIELD_IN_VIEW_UB 1
#define DF_SHADOW_QUALITY 2
#include "LumenTracingCommon.ush"
#include "LumenSoftwareRayTracing.ush"
#include "../DistanceFieldShadowingShared.ush"
#include "LumenSceneDirectLightingPerLightShadowCommon.ush"
#include "../LightDataUniforms.ush"

// Bias
float HeightfieldShadowRayBias;
float MeshSDFShadowRayBias;
float GlobalSDFShadowRayBias;

float MaxTraceDistance;
float SDFSurfaceBiasScale;

float ShadowTraceHeightfield(
	float3 WorldRayStart,
	float3 TranslatedWorldRayStart,
	float3 WorldRayDirection,
	float MaxTraceDistance)
{
	float ConeAngle = 0.0f;
	float MinSampleRadius = 0.0f;
	FConeTraceInput ConeTraceInput;
	ConeTraceInput.Setup(
		WorldRayStart,
		TranslatedWorldRayStart,
		WorldRayDirection,
		ConeAngle,
		MinSampleRadius,
		/*MinTraceDistance*/ 0.0f,
		MaxTraceDistance,
		/*StepFactor*/ 1.0f
	);

	uint NumIntersectingObjects = 0;
	uint CulledDataStart = 0;
	GetHeightfieldShadowTileCulledData(TranslatedWorldRayStart, CulledDataStart, NumIntersectingObjects);

	for (uint IndexInCulledList = 0; IndexInCulledList < NumIntersectingObjects; ++IndexInCulledList)
	{
		const uint HeightfieldObjectIndex = HeightfieldShadowTileArrayData.Load(IndexInCulledList + CulledDataStart);
		FLumenHeightfieldData LumenHeightfield = GetLumenHeightfieldData(HeightfieldObjectIndex);
		if (LumenHeightfield.bValid)
		{
			FConeTraceHeightfieldSimpleResult SimpleResult = ConeTraceHeightfieldSimple(ConeTraceInput, HeightfieldObjectIndex);
			if (SimpleResult.bIsHit)
			{
				return 0.0f;
			}
		}
	}

	return 1.0f;
}

float TraceOffscreenShadows(FLumenCardData CardData, uint2 TexelCoord, float3 WorldPosition, float3 TranslatedWorldPosition, float3 L, float3 ToLight, float3 WorldNormal, bool bLocalLight)
{
	float ShadowFactor = 1.0f;

	// Offscreen shadowing, trace to light
	float TraceDistance = MaxTraceDistance;

	#if LIGHT_TYPE != LIGHT_TYPE_DIRECTIONAL
		TraceDistance = min(length(ToLight), MaxTraceDistance);
	#endif

	#if (LIGHT_TYPE == LIGHT_TYPE_DIRECTIONAL) && (OFFSCREEN_SHADOWING_TRACE_MESH_SDF || OFFSCREEN_SHADOWING_TRACE_HEIGHTFIELDS)
	{
		#if OFFSCREEN_SHADOWING_TRACE_MESH_SDF
		{
			float3 TranslatedWorldPositionForShadowing = GetCardWorldPositionForShadowing(TranslatedWorldPosition, L, WorldNormal, MeshSDFShadowRayBias);
			float3 TranslatedWorldRayStart = TranslatedWorldPositionForShadowing;
			float3 TranslatedWorldRayEnd = TranslatedWorldRayStart + L * TraceDistance;

			uint NumIntersectingObjects = GetCulledNumObjects();
			uint CulledDataParameter = 0;
			GetShadowTileCulledData(TranslatedWorldRayStart, CulledDataParameter, NumIntersectingObjects);

			float SubsurfaceDensity = 0;
			bool bUseSubsurfaceTransmission = false;

			ShadowFactor = ShadowRayTraceThroughCulledObjects(
				TranslatedWorldRayStart,
				TranslatedWorldRayEnd,
				TraceDistance,
				0, //@todo - TanLightSourceAngle - causes mismatch with CSM which doesn't support LightSourceAngle
				0,
				100,
				SubsurfaceDensity,
				CulledDataParameter,
				NumIntersectingObjects,
				true,
				true,
				bUseSubsurfaceTransmission,
				/*bExpandSurface*/ true);
		}
		#endif

		#if OFFSCREEN_SHADOWING_TRACE_HEIGHTFIELDS
		{
			float3 WorldPositionForShadowing = GetCardWorldPositionForShadowing(WorldPosition, L, WorldNormal, HeightfieldShadowRayBias * CardData.TexelSize);
			float3 TranslatedWorldPositionForShadowing = GetCardWorldPositionForShadowing(TranslatedWorldPosition, L, WorldNormal, HeightfieldShadowRayBias * CardData.TexelSize);

			ShadowFactor *= ShadowTraceHeightfield(
				WorldPositionForShadowing,
				TranslatedWorldPositionForShadowing,
				L,
				TraceDistance);
		}
		#endif
	}
	#elif OFFSCREEN_SHADOWING_TRACE_GLOBAL_SDF
	{
		FGlobalSDFTraceInput TraceInput = SetupGlobalSDFTraceInput(TranslatedWorldPosition, L, 0.0f, TraceDistance, 1.0f, 1.0f);

		// Bias away from the starting point on surface
		TraceInput.VoxelSizeRelativeBias = GetCardBiasForShadowing(L, WorldNormal, GlobalSDFShadowRayBias);
		TraceInput.DitherScreenCoord = TexelCoord;

		if (bLocalLight)
		{
			// Bias also the end point (local light position) in case if light is near the geometry
			TraceInput.VoxelSizeRelativeRayEndBias = GlobalSDFShadowRayBias;
		}

		FGlobalSDFTraceResult SDFResult = RayTraceGlobalDistanceField(TraceInput);

		//@todo - approximate cone trace from SDF for antialiasing
		ShadowFactor = GlobalSDFTraceResultIsHit(SDFResult) ? 0.0f : 1.0f;
	}
	#endif

	return ShadowFactor;
}

#ifdef LIGHT_TYPE
[numthreads(THREADGROUP_SIZE_X, THREADGROUP_SIZE_Y, 1)]
void LumenSceneDirectLightingTraceDistanceFieldShadowsCS(
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{
	uint LinearGroupThreadId = GroupThreadId.x + GroupThreadId.y * THREADGROUP_SIZE_X;

#if THREADGROUP_SIZE_32
	uint LocalLightTileIndex = GroupId.x / 2;
	uint2 CoordInCardTile = GroupThreadId.xy + uint2(0, GroupId.x % 2 ? 0 : 4);
	uint ShadowMaskSize = SHADOW_MASK_CARD_TILE_DWORDS / 2;
	uint ShadowMaskOffset = GroupId.x % 2 ? 0 : 1;
#else
	uint LocalLightTileIndex = GroupId.x;
	uint2 CoordInCardTile = GroupThreadId.xy;
	uint ShadowMaskSize = SHADOW_MASK_CARD_TILE_DWORDS;
	uint ShadowMaskOffset = 0;
#endif

	uint SlotIndex = bUseLightTilesPerLightType != 0 && LIGHT_TYPE > 0 ? LIGHT_TYPE - 1 : NUM_BATCHABLE_LIGHT_TYPES + LightIndex;
	uint LightTileIndex = LocalLightTileIndex + LightTileOffsetsPerLight[SlotIndex * NumViews + ViewIndex];

	if (LinearGroupThreadId < ShadowMaskSize)
	{
		SharedShadowMask[LinearGroupThreadId] = RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * LightTileIndex + ShadowMaskSize * ShadowMaskOffset + LinearGroupThreadId];
	}

	GroupMemoryBarrierWithGroupSync();

	FLightTileForShadowMaskPass LightTile = UnpackLightTileForShadowMaskPass(LightTiles[LightTileIndex]);
	uint2 TexelInCardPageCoord = LightTile.TileCoord * CARD_TILE_SIZE + CoordInCardTile;

	FLumenCardPageData CardPage = GetLumenCardPageData(LightTile.CardPageIndex + DummyZeroForFixingShaderCompilerBug);
	float2 AtlasUV = CardPage.PhysicalAtlasUVRect.xy + CardPage.PhysicalAtlasUVTexelScale * (TexelInCardPageCoord + 0.5f);
	float2 CardUV = CardPage.CardUVRect.xy + CardPage.CardUVTexelScale * (TexelInCardPageCoord + 0.5f);

	FShadowMaskRay ShadowMaskRay;
	ReadSharedShadowMaskRay(GroupThreadId.xy, ShadowMaskRay);

	if (!ShadowMaskRay.bShadowFactorComplete)
	{
		FDeferredLightData LightData;
		if (bUseLightTilesPerLightType != 0)
		{
			LightData = LoadLumenLight(LightTile.LightIndex, DFHackToFloat(PrimaryView.PreViewTranslation), ViewExposure[ViewIndex]).DeferredLightData;
		}
		else
		{
			LightData = InitDeferredLightFromUniforms(LIGHT_TYPE);
		}

		FLumenCardData CardData = GetLumenCardData(CardPage.CardIndex);
		FLumenSurfaceCacheData SurfaceCacheData = GetSurfaceCacheData(CardData, CardUV, AtlasUV);

		float3 WorldPosition = SurfaceCacheData.WorldPosition;
		float3 TranslatedWorldPosition = WorldPosition + DFHackToFloat(PrimaryView.PreViewTranslation); // LUMEN_LWC_TODO
		float3 WorldNormal = SurfaceCacheData.WorldNormal;

		float3 L = LightData.Direction;
		float3 ToLight = L;
		bool bLocalLight = false;

#if LIGHT_TYPE != LIGHT_TYPE_DIRECTIONAL
		if (LightData.bRadialLight)
		{
			ToLight = LightData.TranslatedWorldPosition - TranslatedWorldPosition;
			L = normalize(ToLight);
			bLocalLight = true;
		}
#endif

		ShadowMaskRay.bShadowFactorComplete = true;
		ShadowMaskRay.ShadowFactor *= TraceOffscreenShadows(CardData, TexelInCardPageCoord, WorldPosition, TranslatedWorldPosition, L, ToLight, WorldNormal, bLocalLight);

		WriteSharedShadowMaskRay(ShadowMaskRay, GroupThreadId.xy, true);
	}

	GroupMemoryBarrierWithGroupSync();

	if (LinearGroupThreadId < ShadowMaskSize)
	{
		RWShadowMaskTiles[SHADOW_MASK_CARD_TILE_DWORDS * LightTileIndex + ShadowMaskSize * ShadowMaskOffset + LinearGroupThreadId] = SharedShadowMask[LinearGroupThreadId];
	}
}
#endif

RWBuffer<uint> RWShadowTraceIndirectArgs;
StructuredBuffer<uint> ShadowTraceAllocator;

[numthreads(THREADGROUP_SIZE, 1, 1)]
void InitShadowTraceIndirectArgsCS(
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{
	if (all(DispatchThreadId.xyz == 0))
	{
		uint NumShadowTraces = ShadowTraceAllocator[0];

		// One thread per card tile
		WriteDispatchIndirectArgs(RWShadowTraceIndirectArgs, 0, DivideAndRoundUp64(NumShadowTraces), 1, 1);
	}
}