#type vertex
#version 430

layout(location = 0) in vec3 a_Position;
layout(location = 1) in vec2 a_TexCoord;

out vec2 v_TexCoord;

void main()
{
    vec4 position = vec4(a_Position.xy, 0.0, 1.0);
    v_TexCoord = a_TexCoord;
    gl_Position = position;
}

#type fragment
#version 430

layout(location = 0) out vec4 o_Color;

in vec2 v_TexCoord;

uniform sampler2D u_ColorTexture;
uniform sampler2D u_DepthTexture;
uniform sampler2D u_MaterialInfoTexture;

uniform mat4 u_Projection;
uniform mat4 u_View;
uniform mat4 u_InvViewProjection;

uniform vec3 u_CameraPosition;

uniform vec2 u_ScreenSize;
uniform float u_CameraNear;
uniform float u_CameraFar;

uniform int u_Steps = 32;
uniform float u_Thickness = 0.5;
uniform float u_MaxDistance = 30.0;
uniform float u_StepSize = 0.5;
uniform float u_RayOffset = 0.1;
uniform float u_Intensity = 1.0;

float saturate(float x)
{
    return clamp(x, 0.0, 1.0);
}

float LinearizeDepth(float depth)
{
    float z = depth * 2.0 - 1.0;
    return (2.0 * u_CameraNear * u_CameraFar) / (u_CameraFar + u_CameraNear - z * (u_CameraFar - u_CameraNear));
}

vec3 ComputeWorldSpacePosition(vec2 uv, float depth)
{
    vec4 clipPos = vec4(uv * 2.0 - 1.0, depth * 2.0 - 1.0, 1.0);
    vec4 worldPos = u_InvViewProjection * clipPos;
    return worldPos.xyz / worldPos.w;
}

vec3 ReconstructNormalFromDepth(vec2 tc, float depth)
{
    vec2 texelSize = 1.0 / u_ScreenSize;

    vec2 tcL = tc - vec2(texelSize.x, 0.0);
    vec2 tcR = tc + vec2(texelSize.x, 0.0);
    vec2 tcB = tc - vec2(0.0, texelSize.y);
    vec2 tcT = tc + vec2(0.0, texelSize.y);

    float depthL = texture(u_DepthTexture, tcL).r;
    float depthR = texture(u_DepthTexture, tcR).r;
    float depthB = texture(u_DepthTexture, tcB).r;
    float depthT = texture(u_DepthTexture, tcT).r;

    vec3 pC = ComputeWorldSpacePosition(tc, depth);
    vec3 pR = ComputeWorldSpacePosition(tcR, depthR);
    vec3 pT = ComputeWorldSpacePosition(tcT, depthT);

    vec3 dx = pR - pC;
    vec3 dy = pT - pC;

    return normalize(cross(dx, dy));
}

vec3 ComputeNDCWithZ(vec3 worldPos)
{
    vec4 clipPos = u_Projection * u_View * vec4(worldPos, 1.0);
    vec3 ndc = clipPos.xyz / clipPos.w;
    ndc.xy = ndc.xy * 0.5 + 0.5;
    return ndc;
}

void main()
{
    float depth = texture(u_DepthTexture, v_TexCoord).r;

    if (depth >= 1.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec4 matInfo = texture(u_MaterialInfoTexture, v_TexCoord);
    float metalness = matInfo.r;
    float roughness = matInfo.g;

    if (roughness > 0.9)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 positionWS = ComputeWorldSpacePosition(v_TexCoord, depth);
    vec3 viewWS = normalize(u_CameraPosition - positionWS);
    vec3 normalWS = ReconstructNormalFromDepth(v_TexCoord, depth);
    vec3 reflectionWS = normalize(reflect(-viewWS, normalWS));

    if (dot(reflectionWS, normalWS) < 0.0)
    {
        o_Color = vec4(0.0);
        return;
    }

    vec3 rayWS = positionWS + normalWS * u_RayOffset;

    for (int i = 0; i < u_Steps; i++)
    {
        rayWS += reflectionWS * u_StepSize;

        vec3 rayNDC = ComputeNDCWithZ(rayWS);
        float rayDepth = rayNDC.z;

        vec2 raySS = rayNDC.xy * u_ScreenSize;

        if (rayNDC.x < 0.0 || rayNDC.x > 1.0 || rayNDC.y < 0.0 || rayNDC.y > 1.0)
            break;

        float sceneDepth = texture(u_DepthTexture, rayNDC.xy).r;
        float sceneDepthLinear = LinearizeDepth(sceneDepth);
        float rayDepthLinear = LinearizeDepth(rayDepth * 0.5 + 0.5);

        float delta = rayDepthLinear - sceneDepthLinear;

        if (delta >= 0.0 && delta <= u_Thickness)
        {
            vec3 hitNormalWS = ReconstructNormalFromDepth(rayNDC.xy, sceneDepth);
            if (dot(hitNormalWS, reflectionWS) > 0.0)
                break;

            float ssrStrength = saturate(1.0 - roughness * roughness);

            float dist = distance(positionWS, rayWS);
            float fadeFactor = 1.0 - smoothstep(u_MaxDistance * 0.3, u_MaxDistance, dist);

            float edgeFade = 1.0;
            edgeFade *= smoothstep(0.0, 0.1, rayNDC.x);
            edgeFade *= smoothstep(0.0, 0.1, rayNDC.y);
            edgeFade *= smoothstep(0.0, 0.1, 1.0 - rayNDC.x);
            edgeFade *= smoothstep(0.0, 0.1, 1.0 - rayNDC.y);

            float NdotV = max(dot(normalWS, viewWS), 0.0);
            vec3 F0 = mix(vec3(0.04), vec3(1.0), metalness);
            vec3 Fresnel = F0 + (1.0 - F0) * pow(1.0 - NdotV, 5.0);

            ssrStrength *= fadeFactor * edgeFade * u_Intensity;

            vec4 color = texture(u_ColorTexture, rayNDC.xy);
            color.rgb *= ssrStrength * Fresnel;
            color.a = ssrStrength * dot(Fresnel, vec3(0.333));

            o_Color = color;
            return;
        }
    }

    o_Color = vec4(0.0);
}