#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch
//#pragma enable_d3d11_debug_symbols

#pragma kernel ClearWaterLine
#pragma kernel LineEvaluation1D
#pragma kernel BoundsPropagation

// HDRP generic includes
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Water/Shaders/UnderWaterUtilities.hlsl"

TEXTURE2D_X_UINT2(_StencilTexture);
TEXTURE2D_X(_DepthTexture);

// 1D WaterLine Buffer, contains the vertical (along upVector) height of the water line
// first two elements contains the horizontal min and max of the visible water line
RWStructuredBuffer<uint> _WaterLineRW;

[numthreads(8, 1, 1)]
void ClearWaterLine(uint dispatchThreadId : SV_DispatchThreadID)
{
    _WaterLineRW[dispatchThreadId] = 0;
}

[numthreads(8, 8, 1)]
void LineEvaluation1D(uint3 currentCoord : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(currentCoord.z);

    // Early exit pixels containing no water
    float depthValue = LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy).x;
    if (depthValue == UNITY_RAW_FAR_CLIP_VALUE)
        return;
    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, currentCoord.xy));
    if ((stencilValue & STENCILUSAGE_WATER_SURFACE) == 0)
        return;

    // Output to the water line buffer
    float2 upVector = float2(_UpDirectionX, _UpDirectionY);
    float2 rightVector = float2(_UpDirectionY, -_UpDirectionX);

    uint posX = round(dot((float2)currentCoord.xy, rightVector) - _BoundsSS.x);
    uint posY = round(dot((float2)currentCoord.xy, upVector) - _BoundsSS.z);

    // We use InterlockedMax with depth values on higher bits to find closest pixel to camera
    // We store the pixel height as payload in the lower bits to retrieve the waterline height later
    uint height = posY + 1; // Add one to make sure 0 means missing value
    uint packedValue = PackFloatToUInt(depthValue, 16, 16) | (height & 0xFFFF);

    uint idx = min(posX + 2, _BufferStride) + currentCoord.z * _BufferStride;
    InterlockedMax(_WaterLineRW[idx], packedValue);
}

#define GROUP_SIZE 64
#ifndef PLATFORM_SUPPORTS_WAVE_INTRINSICS
groupshared uint2 gs_bounds[GROUP_SIZE];
#endif

uint2 ParallelReduction(uint threadIdx, uint2 bounds)
{
#ifdef PLATFORM_SUPPORTS_WAVE_INTRINSICS
    return uint2(WaveActiveMin(bounds.x), WaveActiveMax(bounds.y));
#else
    gs_bounds[threadIdx] = bounds;

    GroupMemoryBarrierWithGroupSync();

    UNITY_UNROLL
    for (uint s = GROUP_SIZE / 2u; s > 0u; s >>= 1u)
    {
        if (threadIdx < s)
        {
            gs_bounds[threadIdx] = uint2(
                min(gs_bounds[threadIdx].x, gs_bounds[threadIdx + s].x),
                max(gs_bounds[threadIdx].y, gs_bounds[threadIdx + s].y)
            );
        }

        GroupMemoryBarrierWithGroupSync();
    }

    return gs_bounds[0];
#endif
}

[numthreads(GROUP_SIZE, 1, 1)]
void BoundsPropagation(uint2 currentCoord : SV_DispatchThreadID, uint groupThreadId : SV_GroupThreadID)
{
    const uint2 maxBounds = uint2(0xFFFFFFFF, 0);

    // This kernel finds the leftmost and rightmost pixels containing water
    uint xr = currentCoord.y * _BufferStride;
    uint coord = min(currentCoord.x + 2, _BufferStride) + xr;
    uint packedValue = _WaterLineRW[coord];
    uint2 bounds = packedValue == 0 ? maxBounds : (uint2)currentCoord.x;
    bounds = ParallelReduction(groupThreadId, bounds);

    if (groupThreadId == 0)
    {
        InterlockedMax(_WaterLineRW[0 + xr], 0xFFFFFFFF - bounds.x);
        InterlockedMax(_WaterLineRW[1 + xr], bounds.y);
    }

    uint maxHeight = ceil(_BoundsSS.w - _BoundsSS.z);
    float distanceToSurface = GetWaterCameraHeight();
    float distanceToWaterLine = distanceToSurface > 0 ? 0.0f : maxHeight;
    float nearPlane = _ProjectionParams.y;
    if (abs(distanceToSurface) > nearPlane * 2.0f)
    {
        _WaterLineRW[coord] = ((uint)distanceToWaterLine + 1) & 0xFFFF;
        return;
    }

    if (packedValue != 0)
        return;

    // Patch holes of less than GROUP_SIZE pixels
    if (any(bounds != maxBounds))
    {
        uint bound = bounds.x != maxBounds.x ? bounds.x : bounds.y;
        packedValue = _WaterLineRW[bound + 2 + xr];
    }

    // Last hope to patch a hole
    if (packedValue == 0) packedValue = _WaterLineRW[0 + 2 + xr];
    if (packedValue == 0) packedValue = _WaterLineRW[(_BoundsSS.y-_BoundsSS.x) + 2 + xr];

    if (packedValue != 0)
    {
        distanceToWaterLine = (packedValue & 0xFFFF) - 1;
        distanceToWaterLine = round(distanceToWaterLine / maxHeight) * maxHeight;
    }

    _WaterLineRW[coord] = ((uint)distanceToWaterLine + 1) & 0xFFFF;
}