#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/Builtin/BuiltinData.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/DynamicScalingClamping.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Packing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"

#include "Packages/com.misaki.ao-volume/Runtime/Shader/Includes/VolumeData.cs.hlsl"

// Each #kernel tells which function to compile; you can have many kernels
#pragma kernel CSMain

StructuredBuffer<VolumeData> _VolumeDatas;
ByteAddressBuffer _VisibleVolumeCount;
ByteAddressBuffer _VisibleVolumeIndices;

RW_TEXTURE2D_X(float, _AOVolumeBuffer);

float3 GetPositionWS(float2 positionSS, float depth)
{
    float4 positionCS = float4(positionSS * 2.0f - 1.0f, depth * 2.0f - 1.0f, 1.0f);
    float4 positionWS = mul(UNITY_MATRIX_I_VP, positionCS);
    positionWS /= positionWS.w;
    
    return positionWS.xyz;
}

float3 GetNormalWS(float2 positionSS)
{
    float4 normalBufferData = LOAD_TEXTURE2D_X(_NormalBufferTexture, positionSS);
    
    NormalData normalData;
    DecodeFromNormalBuffer(normalBufferData, normalData);
    
    return normalData.normalWS.xyz;
}

float3 GetPositionOS(float4x4 inverseWorldMatrix, float3 positionWS)
{
    // To handle camera relative rendering we need to apply translation before converting to object space
    return mul(ApplyCameraTranslationToInverseMatrix(inverseWorldMatrix), float4(positionWS, 1.0)).xyz;
}

float BoxSDF(float4x4 inverseWorldMatrix, float3 size, float3 positionWS)
{
    float3 halfDim = size * 0.5;
    
    float3 positionLS = GetPositionOS(inverseWorldMatrix, positionWS);
    
    float3 d = halfDim - abs(positionLS);
    float distToFace = min(d.x, min(d.y, d.z));
    float maxDist = min(halfDim.x, min(halfDim.y, halfDim.z));
    
    return saturate(distToFace / maxDist);
}

[numthreads(8,8,1)]
void CSMain(uint3 dispatchThreadId : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);
    
    float depth = LoadCameraDepth(dispatchThreadId.xy);
    float3 positionWS = ComputeWorldSpacePosition(dispatchThreadId.xy / _ScreenSize.xy, depth, UNITY_MATRIX_I_VP); // This cloud be a camera-relative position
    float3 normalWS = GetNormalWS(dispatchThreadId.xy);
    
    float gtao = LOAD_TEXTURE2D_X(_AmbientOcclusionTexture, dispatchThreadId.xy).x;
    
    // TODO: Tile/H-z optmization
    uint visibleCount = _VisibleVolumeCount.Load(0);
    uint i = 0;
    while (i < visibleCount)
    {
        uint volumeIndex = _VisibleVolumeIndices.Load(i << 2);
        VolumeData volumeData = _VolumeDatas[volumeIndex];
        
        float3 volumePostionRWS = GetCameraRelativePositionWS(volumeData.worldMatrix._m03_m13_m23);
        float3 volumeDirection = normalize(volumePostionRWS - positionWS);
        float normalFalloff = saturate(dot(normalWS, volumeDirection) + (2.0 * (1.0 - volumeData.normalFalloff)));
        
        float volumeAO = BoxSDF(volumeData.inverseWorldMatrix, volumeData.size, positionWS);
        volumeAO = smoothstep(0.0, volumeData.falloff, volumeAO) * volumeData.intensity * normalFalloff;
        
        gtao += volumeAO;
        
        i++;
    }
    
    // GTAO is output in a mask format
    _AOVolumeBuffer[COORD_TEXTURE2D_X(dispatchThreadId.xy)] = saturate(gtao);
}