#version 460 core

layout(location = 0) in vec3 inFragPos;
layout(location = 1) in vec4 inShadowPos;
layout(location = 2) in vec3 inNormal;
layout(location = 3) in vec2 inUV;

layout(location = 0) out vec4 outColor;

struct Light {
    vec4 position;
    vec3 color;
};

layout(set = 0, binding = 0) uniform Lights {
    Light lights[32];
};

layout(set = 0, binding = 1) uniform sampler2D shadowSampler;
layout(set = 1, binding = 0) uniform sampler2D albedoSampler;

layout(constant_id = 0) const int shadowFilterPCF = 1;

float textureProj(vec4 P, vec2 off) {
    const vec4 shadowCoord = P / P.w; // perspective divide

    float shadow = 1.0;
    if(shadowCoord.z > -1.0 && shadowCoord.z < 1.0) {
        const float dist = texture(shadowSampler, shadowCoord.st + off).r;
        if(shadowCoord.w > 0.0 && dist < shadowCoord.z)
            shadow = 0.1;
    }

    return shadow;
}

float filterPCF(vec4 sc) {
    const ivec2 texDim = textureSize(shadowSampler, 0);
    const float scale = 1.5;
    const float dx = scale * 1.0 / float(texDim.x);
    const float dy = scale * 1.0 / float(texDim.y);
    const int range = 1;

    float shadowFactor = 0.0;
    int count = 0;
    for(int x = -range; x <= range; x++) {
        for(int y = -range; y <= range; y++) {
            shadowFactor += textureProj(sc, vec2(dx * x, dy * y));
            count++;
        }
    }

    return shadowFactor / count;
}

void main() {
    vec3 diffuse = vec3(0);
    for(int i = 0; i < 32; i++) {
        vec3 lightDir = vec3(0);
        if(lights[i].position.w == 0)
            continue;
        else if(lights[i].position.w == 1)
            lightDir = normalize(lights[i].position.xyz - inFragPos);
        else if(lights[i].position.w == 2)
            lightDir = normalize(-lights[i].position.xyz);

        const vec3 norm = normalize(inNormal);

        const float diff = max(dot(norm, lightDir), 0.0);
        const float shadow = shadowFilterPCF == 1 ? filterPCF(inShadowPos / inShadowPos.w) : textureProj(inShadowPos / inShadowPos.w, vec2(0, 0));
        diffuse += vec3(diff) * lights[i].color * shadow;
    }

    outColor = vec4(vec3(0.1) + diffuse * texture(albedoSampler, inUV).rgb, 1.0);
}