#version 450

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

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) readonly buffer SSAOKernel {
    vec4 samples[];
} 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; //160
} ssao;

void main() {
    vec2 uv = inTextCoord;

    vec3 worldPos = texture(posSampler, uv).xyz;
    vec3 worldNormal = normalize(texture(normalSampler, uv).xyz);


    if (length(worldNormal) < 0.001) {
        outAO = 1.0;
        return;
    }

    vec3 fragPos = vec3(ssao.view * vec4(worldPos, 1.0));
    vec3 normal = normalize(mat3(ssao.view) * worldNormal);

    vec2 noiseScale = ssao.screenSize / 4.0;
    vec3 randomVec = normalize(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 occlusion = 0.0;
    float validSamples = 0.0;;

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

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

        vec2 sampleUV = offset.xy;
        sampleUV.y = 1.0 - sampleUV.y;

        if (sampleUV.x < 0.0 || sampleUV.x > 1.0 ||
            sampleUV.y < 0.0 || sampleUV.y > 1.0) {
            continue;
        }

        vec3 sampleWorldPos = texture(posSampler, sampleUV).xyz;
        vec3 sampleWorldNormal = texture(normalSampler, sampleUV).xyz;

        if (length(sampleWorldNormal) < 0.001) {
            continue;
        }

        vec3 sampleViewPos = vec3(ssao.view * vec4(sampleWorldPos, 1.0));

        float depthDelta = abs(fragPos.z - sampleViewPos.z);

        occlusion += abs(fragPos.z)/200.0f;
        validSamples += 1.0;
        continue;

        if (depthDelta > ssao.radius) {
            continue;
        }

        float rangeCheck = 1.0 - depthDelta / ssao.radius;
        rangeCheck = rangeCheck * rangeCheck * (3.0 - 2.0 * rangeCheck);



        if (sampleViewPos.z >= samplePos.z + ssao.bias) {
            occlusion += rangeCheck;
        }

        validSamples += 1.0;
    }

    if (validSamples <= 0.0) {
        outAO = 1.0;
        return;
    }

    occlusion = 1.0 - occlusion / validSamples;
    outAO = clamp(occlusion, 0.0, 1.0);
}