#version 450

layout(location = 0) in vec2 fragTexCoord;
layout(location = 0) out vec4 outColor;

layout(set = 0, binding = 0) uniform sampler2D albedoSampler;
layout(set = 0, binding = 1) uniform sampler2D worldPosSampler;
layout(set = 0, binding = 2) uniform sampler2D normalSampler;
layout(set = 0, binding = 3) uniform sampler2D pbrSampler;

layout(set = 1, binding = 0) uniform CameraProperties {
    mat4 projection;
    mat4 view;
    float ssrStepSize;
    float ssrMaxDistance;
    int maxSteps;
    int padding;
    vec4 cameraWorldPos;
} ubo;

const float THICKNESS = 0.35;
const float MIN_REFLECTION = 0.01;

float edgeFade(vec2 uv) {
    vec2 fade = smoothstep(vec2(0.0), vec2(0.08), uv) *
    smoothstep(vec2(0.0), vec2(0.08), vec2(1.0) - uv);
    return fade.x * fade.y;
}

void main() {
    outColor = vec4(0.0);
    vec4 posSample = texture(worldPosSampler, fragTexCoord);
    vec4 normalSample = texture(normalSampler, fragTexCoord);
    vec4 pbrSample = texture(pbrSampler, fragTexCoord);

    if (posSample.a == 0.0 || length(posSample.xyz) == 0.0 || length(normalSample.xyz) < 0.001) {
        return;
    }

    vec3 worldPos = posSample.xyz;
    vec3 worldNormal = normalize(normalSample.xyz);

    float roughness = clamp(pbrSample.g, 0.0, 1.0);
    float metallic = clamp(pbrSample.b, 0.0, 1.0);
    float reflectiveness = clamp(pbrSample.a, 0.0, 1.0);
    float refractiveness = clamp(normalSample.a, 0.0, 2.0);

    vec3 viewDirToCamera = normalize(ubo.cameraWorldPos.xyz - worldPos);
    float NoV = clamp(dot(worldNormal, viewDirToCamera), 0.0, 1.0);
    float fresnel = pow(1.0 - NoV, 5.0);

    float materialReflection = reflectiveness + (refractiveness - 1.05);
    materialReflection *= mix(0.75, 1.0, metallic);
    materialReflection *= 1.0 - smoothstep(0.05, 0.7, roughness);
    materialReflection *= mix(0.45, 1.0, fresnel);

    if (materialReflection <= MIN_REFLECTION) {
        return;
    }

    vec3 viewDir = normalize(worldPos - ubo.cameraWorldPos.xyz);
    vec3 reflectDir = normalize(reflect(viewDir, worldNormal));

    if (dot(reflectDir, worldNormal) <= 0.0) {
        return;
    }

    int steps = clamp(ubo.maxSteps, 1, 128);
    float stepSize = max(ubo.ssrStepSize, 0.01);
    float maxDistance = max(ubo.ssrMaxDistance, stepSize);

    vec3 rayPos = worldPos;
    vec2 hitUV = vec2(0.0);
    bool hitFound = false;

    for (int i = 0; i < steps; i++) {
        rayPos += reflectDir * stepSize;

        if (distance(rayPos, worldPos) > maxDistance) {
            break;
        }

        vec4 clip = ubo.projection * ubo.view * vec4(rayPos, 1.0);
        if (clip.w <= 0.0) {
            break;
        }

        vec2 uv = clip.xy / clip.w;
        uv = uv * 0.5 + 0.5;
        uv.y = 1.0 - uv.y;

        if (uv.x <= 0.0 || uv.x >= 1.0 || uv.y <= 0.0 || uv.y >= 1.0) {
            break;
        }

        vec4 surfaceWorld = texture(worldPosSampler, uv);
        if (surfaceWorld.a == 0.0 || length(surfaceWorld.xyz) == 0.0) {
            continue;
        }

        float rayViewZ = abs((ubo.view * vec4(rayPos, 1.0)).z);
        float surfaceViewZ = abs((ubo.view * vec4(surfaceWorld.xyz, 1.0)).z);

        float depthDelta = rayViewZ - surfaceViewZ;
        if (depthDelta >= 0.0 && depthDelta < THICKNESS) {
            hitUV = uv;
            hitFound = true;
            break;
        }
    }

    if (!hitFound) {
        return;
    }

    vec3 reflectedColor = texture(albedoSampler, hitUV).rgb;

    float fade = edgeFade(hitUV);
    float weight = clamp(materialReflection * fade, 0.0, 0.85);

    outColor = vec4(reflectedColor * weight, weight);
}