#version 450

layout(location = 0) in vec2 inTextCoord;
layout(location = 0) out float outAO;
layout(location = 1) out vec4 viewPos;

layout(set = 0, binding = 0) uniform sampler2D posSampler;
layout(set = 0, binding = 1) uniform sampler2D normalSampler;
layout(set = 0, binding = 2) uniform sampler2D noiseSampler;

layout(set = 1, binding = 0) uniform  SSAOKernel {
    vec4 samples[64];
} kernel;

layout(set = 2, binding = 0) uniform SSAOInfo {
    mat4 projection; //64
    mat4 view; //128
    vec2 screenSize; // 136
    float radius; //140
    float bias; //144
    int kernelSize; //148
    float ResolutionScale; //152;
    vec2 Padding; //160;
} ssao;

void main() {

    vec2 ScreenSize = ssao.screenSize * ssao.ResolutionScale;

    int KERNEL_SIZE = ssao.kernelSize;
    vec3 fragPos = texture(posSampler, inTextCoord).xyz;
    vec3 worldNorm = normalize(texture(normalSampler, inTextCoord).xyz);

    if (dot(worldNorm, worldNorm) < 0.001) {
        outAO = 1.0;
        return;
    }

    vec3 normal = normalize(mat3(ssao.view) *  worldNorm);

    viewPos = vec4(fragPos, 1);

    vec2 noiseScale = ScreenSize / 4.0;
    vec3 randomVec = texture(noiseSampler, inTextCoord * noiseScale).xyz;
    vec3 tangent = normalize(randomVec - normal * dot(randomVec, normal));
    vec3 bitangent = cross(normal, tangent);
    mat3 TBN = mat3(tangent, bitangent, normal);
    float fragDepth = -fragPos.z;

    float minOptimalDistance = 15;
    float maxKernelSize = float(ssao.kernelSize);

    // Calculate a dynamic loop limit based on proximity
    // If closer than 1.5 units, smoothly scale down the loop count
    int dynamicKernelSize = int(KERNEL_SIZE);
    if (fragDepth < minOptimalDistance) {
        float proximityFactor = clamp(fragDepth / minOptimalDistance, 0.2, 1.0);
        dynamicKernelSize = int(maxKernelSize * proximityFactor);
    }

    // Force a hard minimum baseline of samples so SSAO doesn't completely disappear
    dynamicKernelSize = max(dynamicKernelSize, 16);

    float occlusion = 0.0;

    float Passes = 0.0f;

    for (int i = 0; i < dynamicKernelSize; i++) {
        vec3 samplePos = fragPos + (TBN * kernel.samples[i].xyz) * ssao.radius;
        vec4 offset = ssao.projection * vec4(samplePos, 1.0);
        offset.xyz /= offset.w;
        offset.xyz = offset.xyz * 0.5 + 0.5;
        vec2 sampleUV = clamp(vec2(offset.x, 1.0 - offset.y), 0.0, 1.0);
        vec3 sampleViewPos = texture(posSampler, sampleUV).xyz;
        float sampleDepth = -sampleViewPos.z;
        float depthDelta = abs(fragDepth - sampleDepth);
        if (depthDelta > ssao.radius * 2.0) {
            continue;
        }
        float rangeCheck = smoothstep(0.0, 1.0, ssao.radius / depthDelta);
        float isOccluded = step(samplePos.z + ssao.bias, sampleViewPos.z);
        occlusion += isOccluded * rangeCheck;
        Passes += 1.0;
    }

    float Subtraction = 0.0;
    if(Passes > 0){
        Subtraction =  (occlusion / Passes);
    }
    occlusion = 1.0 - Subtraction;
    outAO = clamp(pow(occlusion, 3), 0.0, 1.0);
}