#version 450

layout(location = 0) in vec2 inUV;

layout(location = 0) out vec4 outColor;

layout(set = 0, binding = 0) uniform sampler2D uPosition;
layout(set = 0, binding = 1) uniform sampler2D uNormal;
layout(set = 0, binding = 2) uniform sampler2D uAlbedo;
layout(set = 0, binding = 3) uniform sampler2D uMaterialProperties;

const int kernelSize = 64;
layout(set = 1, binding = 0) uniform Samples {
    vec4[kernelSize] samples;
};

layout(set = 1, binding = 1) uniform sampler2D uNoiseTexture;

layout (set = 2, binding = 0) uniform UniformBuffer {
    mat4  view;
    mat4  proj;
    mat4  viewInverse;
    mat4  projInverse;
    vec4  position;
    float nearPlane;
    float farPlane;
    int isDebugCam;
    int disableSSAO;
} ub_data;

const float radius = 0.5;
const float bias = 0.025;

void main() {
    if (ub_data.disableSSAO == 1) {
        outColor = vec4(1.0);
        return;
    }
    
    vec2 noiseScale = vec2(textureSize(uPosition, 0)) / vec2(textureSize(uNoiseTexture, 0));
    
    vec3 fragPos = texture(uPosition, inUV).xyz;
    vec3 normal = normalize(texture(uNormal, inUV).xyz);
    vec3 randomVec = normalize(texture(uNoiseTexture, inUV * noiseScale).xyz);

    vec3 tangent = normalize(randomVec - normal * dot(randomVec, normal));
    vec3 bitangent = cross(normal, tangent);
    mat3 TBN = mat3(tangent, bitangent, normal);  

    float occlusion = 0.0;
    for (int i = 0; i < kernelSize; i++) {
        vec3 samplePos = TBN * samples[i].xyz;
        samplePos = fragPos + samplePos * radius;

        vec4 offset = vec4(samplePos, 1.0);
        offset = ub_data.proj * offset;
        offset.xyz /= offset.w;
        offset.xyz = offset.xyz * 0.5 + 0.5;

        float sampleDepth = texture(uPosition, offset.xy).z;

        float rangeCheck = smoothstep(0.0, 1.0, radius / abs(fragPos.z - sampleDepth));
        occlusion += (sampleDepth >= samplePos.z + bias ? 1.0 : 0.0) * rangeCheck;
    }

    occlusion = 1.0 - (occlusion / kernelSize);
    outColor = vec4(occlusion);
}
