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

#include "/Engine/Shared/RayTracingTypes.h"
#include "../Common.ush"
#include "../BlueNoise.ush"
#include "../IntersectionUtils.ush"
#include "../RayTracing/RayTracingCommon.ush"
#include "../SceneTexturesCommon.ush"
#include "MegaLights.ush"
#include "MegaLightsRayTracing.ush"
#include "../Lumen/LumenHardwareRayTracingPayloadCommon.ush"
#include "../Lumen/LumenHardwareRayTracingCommon.ush"
#if HAIR_VOXEL_TRACES
	#include "../HairStrands/HairStrandsRaytracing.ush"
#endif

float RayTracingBias;
float RayTracingNormalBias;
float RayTracingPullbackBias;
uint DebugMode;
uint MaxTraversalIterations;
uint MeshSectionVisibilityTest;
RaytracingAccelerationStructure TLAS;
RaytracingAccelerationStructure FarFieldTLAS;

float NearFieldSceneRadius;
float NearFieldMaxTraceDistance;
float NearFieldMaxTraceDistanceDitherScale;
float FarFieldBias;
float FarFieldMaxTraceDistance;

float DistantScreenTraceSlopeCompareTolerance;
float DistantScreenTraceStartDistance;
float DistantScreenTraceLength;

#if LUMEN_HARDWARE_INLINE_RAYTRACING
	StructuredBuffer<FHitGroupRootConstants> HitGroupData;
	StructuredBuffer<FRayTracingSceneMetadataRecord> RayTracingSceneMetadata;
	RWStructuredBuffer<uint> RWInstanceHitCountBuffer;
#endif

RWTexture2D<uint> RWLightSamples;
RWTexture2D<uint> RWLightSampleRays;

Buffer<uint> CompactedTraceTexelData;
Buffer<uint> CompactedTraceTexelAllocator;
Texture2D<float> DownsampledSceneDepth;
Texture2D<UNORM float3> DownsampledSceneWorldNormal;

RAY_TRACING_ENTRY_RAYGEN_OR_INLINE(HardwareRayTraceLightSamples)
{
	uint ThreadIndex = DispatchThreadIndex.x;
	uint GroupIndex = DispatchThreadIndex.y;
	uint TraceTexelIndex = GroupIndex * 64 + ThreadIndex;

	if (TraceTexelIndex < CompactedTraceTexelAllocator[0])
	{
		uint2 SampleCoord = UnpackTraceTexel(CompactedTraceTexelData[TraceTexelIndex]);
		uint2 ScreenCoord = SampleCoordToScreenCoord(SampleCoord);
		uint2 DownsampledScreenCoord = SampleCoordToDownsampledScreenCoord(SampleCoord);

		FShaderPrintContext DebugContext = InitDebugContext(DownsampledScreenCoord, /*bDownsampled*/ true, float2(0.55, 0.45));

		float2 ScreenUV = (ScreenCoord + 0.5f) * View.BufferSizeAndInvSize.zw;
		float SceneDepth = DownsampledSceneDepth[DownsampledScreenCoord];
		float3 SceneWorldNormal = normalize(DecodeNormal(DownsampledSceneWorldNormal[DownsampledScreenCoord]));
		float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, SceneDepth);

		FLightSample LightSample = UnpackLightSample(RWLightSamples[SampleCoord]);
		FLightSampleRay LightSampleRay = UnpackLightSampleRay(RWLightSampleRays[SampleCoord]);
		FLightSampleTrace LightSampleTrace = GetLightSampleTrace(TranslatedWorldPosition, LightSample.LocalLightIndex, LightSampleRay.UV);

		const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);

		const float ClippedNearFieldMaxTraceDistance = ClipAndDitherNearFieldMaxTraceDistance(
			TranslatedWorldPosition,
			LightSampleTrace.Direction,
			ScreenCoord,
			NearFieldSceneRadius,
			NearFieldMaxTraceDistance,
			NearFieldMaxTraceDistanceDitherScale);

		const float MaxTraceDistance = LightSampleTrace.Distance - UnpackLightRayEndBias(ForwardLightData);

		FRayDesc Ray = (FRayDesc)0;
		Ray.Origin = TranslatedWorldPosition;
		Ray.Direction = LightSampleTrace.Direction;
		Ray.TMin = max(LightSampleRay.RayDistance - RayTracingPullbackBias, RayTracingBias);
		Ray.TMax = max(Ray.TMin, min(MaxTraceDistance, ClippedNearFieldMaxTraceDistance));

		bool bFlipNormalBiasDirection = false;
		if (LightSampleRay.bBackfaceDiffuse && dot(LightSampleTrace.Direction, SceneWorldNormal) < 0.0f)
		{
			bFlipNormalBiasDirection = true;
		}

		ApplyPositionBias(Ray.Origin, Ray.Direction, bFlipNormalBiasDirection ? -SceneWorldNormal : SceneWorldNormal, RayTracingNormalBias);

		FRayCone RayCone = (FRayCone)0;
		FRayTracedLightingContext LightingContext = CreateRayTracedLightingContext(
			RayCone,
			0,
			0, // dummy coordinate
			/*CullingMode*/ FORCE_TWO_SIDED ? RAY_FLAG_NONE : RAY_FLAG_CULL_FRONT_FACING_TRIANGLES,
			MaxTraversalIterations,
			MeshSectionVisibilityTest != 0 || MEGA_LIGHTS_LIGHTING_CHANNELS);

		LightingContext.LightingChannelMask = UnpackLightingChannelMask(ForwardLightData);

		// Shadows don't need closest hit distance
		LightingContext.bAcceptFirstHitAndEndSearch = true;

#if LUMEN_HARDWARE_INLINE_RAYTRACING
		LightingContext.HitGroupData = HitGroupData;
		LightingContext.RayTracingSceneMetadata = RayTracingSceneMetadata;
		LightingContext.RWInstanceHitCountBuffer = RWInstanceHitCountBuffer;
#endif // LUMEN_HARDWARE_INLINE_RAYTRACING

		LightingContext.InstanceMask = RAY_TRACING_MASK_OPAQUE_SHADOW;
		LightingContext.bIsShadowRay = true;

		// by default bIsShadowRay causes RAY_FLAG_SKIP_CLOSEST_HIT_SHADER to be used, but for visualizations we need CHS to run to get hit data such as HitT
		if (DebugContext.bIsActive && DebugMode == DEBUG_MODE_VISUALIZE_TRACING)
		{
			LightingContext.bForceClosestHitShader = true;
		}

#if SUPPORT_CONTINUATION
		// Need to evaluate CHS as this pass depends on TraceResult.bAlphaMasked
		LightingContext.bForceClosestHitShader = true;
#endif

#if MEGA_LIGHTS_EVALUATE_MATERIALS
		FLumenMinimalRayResult TraceResult = (FLumenMinimalRayResult)0;
		{
			FPackedMaterialClosestHitPayload Payload = (FPackedMaterialClosestHitPayload)0;
			Payload.SetLumenPayload();
			Payload.SetIgnoreTranslucentMaterials();
			Payload.SetMinimalPayloadMode();
			Payload.SetRandom(0.5f); // Is an actual random number needed?

			uint RayFlags = RAY_FLAG_SKIP_CLOSEST_HIT_SHADER;
			RayFlags |= LightingContext.CullingMode;
			RayFlags |= LightingContext.bAcceptFirstHitAndEndSearch ? RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH : 0;

			TraceLumenShadingRay(TLAS, RayFlags, LightingContext.InstanceMask, RAY_TRACING_SHADER_SLOT_SHADOW, RAY_TRACING_NUM_SHADER_SLOTS, 0, Ray.GetNativeDesc(), Payload);

			TraceResult.bHit = Payload.IsHit();
			if (TraceResult.bHit)
			{
				TraceResult.HitT = Payload.HitT;
			}
			#if ENABLE_FAR_FIELD_TRACING
			else if(Ray.TMax < MaxTraceDistance)
			{
				FRayDesc FarFieldRay = Ray;
				FarFieldRay.TMin = max(Ray.TMax, FarFieldBias);
				FarFieldRay.TMax = max(FarFieldRay.TMin, min(MaxTraceDistance, FarFieldMaxTraceDistance));
				LightingContext.bIsFarFieldRay = true;

				FPackedMaterialClosestHitPayload FarFieldPayload = (FPackedMaterialClosestHitPayload)0;
				FarFieldPayload.SetLumenPayload();
				FarFieldPayload.SetIgnoreTranslucentMaterials();
				FarFieldPayload.SetMinimalPayloadMode();
				FarFieldPayload.SetRandom(0.5f); // Is an actual random number needed?

				TraceLumenShadingRay(FarFieldTLAS, RayFlags, LightingContext.InstanceMask, RAY_TRACING_SHADER_SLOT_SHADOW, RAY_TRACING_NUM_SHADER_SLOTS, 0, FarFieldRay.GetNativeDesc(), FarFieldPayload);

				TraceResult.bHit = FarFieldPayload.IsHit();
				if (TraceResult.bHit)
				{
					TraceResult.HitT = FarFieldPayload.HitT;
				}
			}
			#endif
		}
#else
		FLumenMinimalRayResult TraceResult = TraceLumenMinimalRay(TLAS, Ray, LightingContext);

		#if ENABLE_FAR_FIELD_TRACING
		{
			if (!TraceResult.bHit && Ray.TMax < MaxTraceDistance)
			{
				FRayDesc FarFieldRay = Ray;
				FarFieldRay.TMin = max(Ray.TMax, FarFieldBias);
				FarFieldRay.TMax = max(FarFieldRay.TMin, min(MaxTraceDistance, FarFieldMaxTraceDistance));
				LightingContext.bIsFarFieldRay = true;

				TraceResult = TraceLumenMinimalRay(FarFieldTLAS, FarFieldRay, LightingContext);
			}
		}
		#endif
#endif

		#if HAIR_VOXEL_TRACES
		if (!TraceResult.bHit)
		{
			// #ml_todo: replace with spatiotemporal blue noise
			// #ml_todo: also check if DownsampledScreenCoord should be used instead to benefit from error-diffusion 
			FRandomSequence RandSequence = RandomSequenceCreate(uint3(SampleCoord, MegaLightsStateFrameIndex), 0, 1);
			const float HitT = TraverseHair(ScreenCoord, RandSequence, Ray.Origin, Ray.Direction, Ray.TMax, VirtualVoxel.Raytracing_ShadowOcclusionThreshold);
			if (HitT > 0 && HitT < (TraceResult.bHit ? TraceResult.HitT : Ray.TMax))
			{
				TraceResult.HitT = HitT;
				TraceResult.bHit = true;
			}
		}
		#endif

		if (DebugContext.bIsActive && DebugMode == DEBUG_MODE_VISUALIZE_TRACING)
		{
			const FForwardLightData ForwardLightData = GetForwardLightData(LightSample.LocalLightIndex, 0);
			const FDeferredLightData LightData = ConvertToDeferredLight(ForwardLightData);

			float3 RayStart = Ray.Origin + Ray.Direction * Ray.TMin;
			float3 RayEnd = Ray.Origin + Ray.Direction * (TraceResult.bHit ? TraceResult.HitT : Ray.TMax);
			float4 RayColor = float4(LightData.Color.xyz / Luminance(LightData.Color.xyz), 1.0f);

#if MEGA_LIGHTS_EVALUATE_MATERIALS
			RayColor = ColorRed;
#endif

			if (TraceResult.bHit)
			{
				RayColor.xyz = 0.0f;
			}

			AddLineTWS(DebugContext, RayStart, RayEnd, RayColor);
			AddCrossTWS(DebugContext, RayEnd, 2.0f, RayColor);
		}

#if DISTANT_SCREEN_TRACES
		if (!TraceResult.bHit)
		{
			const float ContactShadowLengthScreenScale = GetScreenRayLengthMultiplierForProjectionType(SceneDepth).y;

			float3 TranslatedDistantRayOrigin = Ray.Origin + Ray.Direction * Ray.TMin;
			float MaxTraceDistanceForDistantScreenTrace = DistantScreenTraceLength * ContactShadowLengthScreenScale;

			// Start the distant linear screen traces where we have no HWRT representation
			float2 RaySphereHit = RayIntersectSphere(TranslatedDistantRayOrigin, Ray.Direction, float4(View.TranslatedWorldCameraOrigin, DistantScreenTraceStartDistance));
			if (RaySphereHit.x < 0 && RaySphereHit.y > 0)
			{
				// Pull the origin inside the sphere slightly to reduce the gap created by jittering ray step offset
				// t + (t - 1) / (n * 2 - 1) where n is the number of distant screen trace steps
				RaySphereHit.y = max((32.f / 31.f) * RaySphereHit.y - (1.f / 31.f) * MaxTraceDistanceForDistantScreenTrace, 0);

				TranslatedDistantRayOrigin += RaySphereHit.y * Ray.Direction;
				MaxTraceDistanceForDistantScreenTrace -= RaySphereHit.y;	
			}

			if (MaxTraceDistanceForDistantScreenTrace > 0)
			{
				const float Dither = InterleavedGradientNoise(SampleCoord.xy + 0.5f, View.StateFrameIndexMod8);
				const bool bHairNoShadowLight = false;
				const uint NumSteps = 8;

				float2 HitUV;
				float HitDistance = CastScreenSpaceShadowRay(TranslatedDistantRayOrigin, Ray.Direction, MaxTraceDistanceForDistantScreenTrace, NumSteps, Dither, DistantScreenTraceSlopeCompareTolerance, bHairNoShadowLight, HitUV);

				if ( HitDistance > 0.0 )
				{
					TraceResult.bHit = true;
				}
			}
		}
		#endif

#if SUPPORT_CONTINUATION
		if (TraceResult.bAlphaMasked)
		{
			// do nothing so sample gets retraced with material evaluation
			// can't shorten the retrace ray since this trace was done with bAcceptFirstHitAndEndSearch
		}
		else
		{
			// Mark sample as complete to skip continuation
			LightSampleRay.bCompleted = true;
			RWLightSampleRays[SampleCoord] = PackLightSampleRay(LightSampleRay);

			if (TraceResult.bHit)
			{
				LightSample.bVisible = false;
			}
			RWLightSamples[SampleCoord] = PackLightSample(LightSample);
		}
#else
		if (TraceResult.bHit) // if not using continuation, only need to write to RWLightSamples on hits
		{
			LightSample.bVisible = false;
			RWLightSamples[SampleCoord] = PackLightSample(LightSample);
		}
#endif
	}
}