#pragma kernel DeferredContactShadow

#pragma multi_compile _ ENABLE_MSAA

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"

#define USE_FPTL_LIGHTLIST // Use light tiles for contact shadows
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightLoop/LightLoopDef.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/Shadow/ContactShadows.hlsl"

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#ifdef SHADER_API_PSSL
#include SHADER_COMPILER_GLOBAL_OPTIMIZE_REGISTER_USAGE
#endif

// #pragma enable_d3d11_debug_symbols

#ifdef ENABLE_MSAA
    TEXTURE2D_X(_CameraDepthValues);
#endif

#define DEFERRED_SHADOW_TILE_SIZE 8

float SampleDepth(float2 UV, bool HalfRes)
{
    float2 pixelCoord = UV.xy * _ScreenSize.xy;

    if (HalfRes)
    {
        pixelCoord.x *= 0.5f;
        pixelCoord.y = pixelCoord.y * 0.5f + _RenderTargetHeight;
    }

    return LoadCameraDepth(pixelCoord);
}

float GetDepthCompareThreshold(float step, float rayStartZ, float rayOrthoZ)
{
    return abs(rayOrthoZ - rayStartZ) * _ContactShadowThickness * max(0.07, step);
}

bool CompareDepth(float depthDiff, float compareThreshold)
{
    return abs(compareThreshold - depthDiff) < compareThreshold;
}


bool ScreenSpaceShadowRayCast(float3 positionWS, float3 rayDirWS, float rayLength, uint2 positionSS, out float fade)
{

    // Dither pattern is shifted by 0.5 because we want to jitter the ray starting position backward and forward (so we need values between -0.5 and 0.5)
    float ditherBias = 0.5f;
    uint taaEnabled = _TaaFrameInfo.w;
    float dither = InterleavedGradientNoise(positionSS, (_FrameCount % 8u) * taaEnabled) - ditherBias;

    float3 rayStartWS = positionWS - positionWS * _ContactShadowBias;
    float3 rayEndWS = rayStartWS + rayDirWS * rayLength;

    float4 rayStartCS = TransformWorldToHClip(rayStartWS);
    float4 rayEndCS = TransformWorldToHClip(rayEndWS);

    // Here we compute a ray perpendicular to view space. This is the ray we use to compute the threshold for rejecting samples.
    // This is done this way so that the threshold is less dependent of ray slope.
    float4 rayOrthoViewSpace = rayStartCS + float4(GetViewToHClipMatrix()[0][2], GetViewToHClipMatrix()[1][2], GetViewToHClipMatrix()[2][2], GetViewToHClipMatrix()[3][2]) * rayLength;
    rayOrthoViewSpace = rayOrthoViewSpace / rayOrthoViewSpace.w;

    rayStartCS.xyz = rayStartCS.xyz / rayStartCS.w;
    rayEndCS.xyz = rayEndCS.xyz / rayEndCS.w;

    // Pixel to light ray in clip space.
    float3 rayDirCS = rayEndCS.xyz - rayStartCS.xyz;

    float step = 1.0f / _SampleCount;
    float compareThreshold = GetDepthCompareThreshold(step, rayStartCS.z, rayOrthoViewSpace.z);

    float occluded = 0.0f;

    // From this point on, all the marching will be done in UV space + Z
    float2 startUV = rayStartCS.xy * 0.5f + 0.5f;
    startUV.y = 1.0f - startUV.y;
    float3 rayStart = float3(startUV, rayStartCS.z);
    float3 rayDir = float3(rayDirCS.x * 0.5f, -rayDirCS.y * 0.5f, rayDirCS.z);

    float t = step * dither + step;

    bool tracingHalfRes = true;
    int i = 0;

    // While instead of for loop to fix HLSLcc compiler error by converting for loop to while loop (case 1158280)
    while (true)
    {
        if (!(i < _SampleCount))
            break;

        float3 sampleAlongRay = rayStart + t * rayDir;


        if (any(sampleAlongRay.xy < 0) || any(sampleAlongRay.xy > 1))
        {
            break;
        }

#ifdef ENABLE_MSAA
        // Depth buffer depth for this sample
        float sampleDepth = SAMPLE_TEXTURE2D_X_LOD(_CameraDepthValues, s_point_clamp_sampler, ClampAndScaleUVForPoint(sampleAlongRay.xy), 0).y;
#else
        // Depth buffer depth for this sample
        float sampleDepth = SampleDepth(saturate(sampleAlongRay.xy), tracingHalfRes);
#endif

        float depthDiff = sampleDepth - sampleAlongRay.z;

        if (depthDiff > 0.0f && CompareDepth(depthDiff, compareThreshold) && sampleAlongRay.z > 0)
        {
            if (tracingHalfRes)
            {
                // Move permanentely to full res. There is a good chance we've hit something, so next full res sample will find the intersection.
                // If not, then allowing each thread to go back and forth between half and full was not beneficial to perf.
                tracingHalfRes = false;
                continue;
            }
            else
            {
                occluded = 1.0f;
                break;
            }
        }

        t += step;
        i++;
    }

    // Off screen masking
    // We remove the occlusion if the ray is occluded and only if direction steps out of the screen
    float2 vignette = max(6.0f * abs(rayStartCS.xy + rayDirCS.xy * t) - 5.0f, 0.0f);
    fade = occluded;
    fade *= saturate(1.0f - dot(vignette, vignette));

    return occluded;
}

bool ComputeContactShadow(PositionInputs posInput, float3 direction, inout float globalFade)
{
    bool    occluded = false;
    float   fade;

    if (_ContactShadowLength > 0.0f)
    {
        //Here LightDirection is not the light direction but the light position
        float rayLength = _ContactShadowLength * max(0.5, posInput.linearDepth * _ContactShadowDistanceScaleFactor);
        occluded = ScreenSpaceShadowRayCast(posInput.positionWS, direction, rayLength, posInput.positionSS, fade);
        // Fade in
        fade *= saturate((posInput.linearDepth - _ContactShadowMinDistance) * rcp(_ContactShadowFadeInEnd));
        // Fade out
        fade *= saturate((_ContactShadowFadeEnd - posInput.linearDepth) * _ContactShadowFadeOneOverRange);

        globalFade = max(globalFade, fade);
    }

    return occluded;
}

[numthreads(DEFERRED_SHADOW_TILE_SIZE, DEFERRED_SHADOW_TILE_SIZE, 1)]
void DeferredContactShadow(uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID, uint3 dispatchThreadId : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    uint2 pixelCoord = groupId * DEFERRED_SHADOW_TILE_SIZE + groupThreadId;
    // There might be a mismatch between the size of a tile and the group size of this shader.
    uint2 tileCoord = (groupId * DEFERRED_SHADOW_TILE_SIZE) / (TILE_SIZE_FPTL);

#ifdef ENABLE_MSAA
    float depth = LOAD_TEXTURE2D_X(_CameraDepthValues, pixelCoord.xy).z;
#else
    float depth = LoadCameraDepth(pixelCoord.xy);
#endif

    PositionInputs posInput = GetPositionInput(pixelCoord.xy, _ScreenSize.zw, depth, UNITY_MATRIX_I_VP, UNITY_MATRIX_V, tileCoord);

    // discard the shadow if we're on the sky or outside of the contact shadow range
    if (depth == UNITY_RAW_FAR_CLIP_VALUE || posInput.linearDepth - _ContactShadowFadeEnd > 1 || posInput.linearDepth < _ContactShadowMinDistance)
    {
        _ContactShadowTextureUAV[COORD_TEXTURE2D_X(pixelCoord)] = 0;

        // TODO: investigate why the following return statement generates invalid glsl code when TEXTURE_2D_X is set to expand to Texture2DArray
        //return;
    }
    else
    {
        // store the 24 bit contact shadow mask mask (1: pixel is in contact shadow, 0: pixel is not shadowed)
        uint    contactShadowMask = 0;
        // the fade is combined for all lights and stored in the 8 remaining bits of the R32 target
        float   globalFade = 0.0;

        uint featureFlags = ~0;

        // Do the contact shadow for the directional light
        if (featureFlags & LIGHTFEATUREFLAGS_DIRECTIONAL)
        {
            if (_DirectionalShadowIndex >= 0)
            {
                DirectionalLightData light = _DirectionalLightDatas[_DirectionalShadowIndex];

                if (light.contactShadowMask != 0 && light.isRayTracedContactShadow == 0.0)
                {
                    // We store the inverse of the contact shadow:
                    bool occluded = ComputeContactShadow(posInput, -light.forward, globalFade);

                    // light.contactShadowMask contains one bit at the position of the contact shadow index that will
                    // be tested in the lightloop, so it insert 1 at the index of the contact shadow if there is a contact shadow
                    // we take full bits at one multiplied by contact shadow and filter the bit at the contact shadow index.
                    contactShadowMask |= light.contactShadowMask * occluded;
                }
            }
        }

        // iterate over all point/spot lights
        if (featureFlags & LIGHTFEATUREFLAGS_PUNCTUAL)
        {
            uint lightCount, lightStart;

#ifndef LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
            GetCountAndStart(posInput, LIGHTCATEGORY_PUNCTUAL, lightStart, lightCount);
#else   // LIGHTLOOP_DISABLE_TILE_AND_CLUSTER
            lightCount = _PunctualLightCount;
            lightStart = 0;
#endif
            uint startFirstLane = 0;
            bool fastPath;

            fastPath = IsFastPath(lightStart, startFirstLane);
            if (fastPath)
            {
                lightStart = startFirstLane;
            }

            uint v_lightIdx = lightStart;
            uint v_lightListOffset = 0;
            while (v_lightListOffset < lightCount)
            {
                v_lightIdx = FetchIndex(lightStart, v_lightListOffset);
                uint s_lightIdx = ScalarizeElementIndex(v_lightIdx, fastPath);
                if (s_lightIdx == -1)
                    break;

                LightData s_lightData = FetchLight(s_lightIdx);

                // If current scalar and vector light index match, we process the light. The v_lightListOffset for current thread is increased.
                // Note that the following should really be ==, however, since helper lanes are not considered by WaveActiveMin, such helper lanes could
                // end up with a unique v_lightIdx value that is smaller than s_lightIdx hence being stuck in a loop. All the active lanes will not have this problem.
                if (s_lightIdx >= v_lightIdx)
                {
                    v_lightListOffset++;
                    if (s_lightData.contactShadowMask != 0 && s_lightData.isRayTracedContactShadow == 0.0)
                    {
                        // Compute light ray direction:
                        float3 direction = normalize(s_lightData.positionRWS.xyz - posInput.positionWS);

                        bool occluded = ComputeContactShadow(posInput, direction, globalFade);

                        // light.contactShadowMask contains one bit at the position of the contact shadow index that will
                        // be tested in the lightloop, so it insert 1 at the index of the contact shadow if there is a contact shadow
                        // we take full bits at one multiplied by contact shadow and filter the bit at the contact shadow index.
                        contactShadowMask |= s_lightData.contactShadowMask * occluded;
                    }
                }
            }
        }

        _ContactShadowTextureUAV[COORD_TEXTURE2D_X(pixelCoord)] = PackContactShadowData(globalFade, contactShadowMask);
    }
}