#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;

layout(push_constant) uniform PushConstants {
    mat4 m;
    vec3 c;
} pushConstants;

struct Light {
    vec4 position, color;
};

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

layout(set = 0, binding = 2) 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.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);

        vec3 add = vec3(diff) * lights[i].color.rgb;

        if(lights[i].position.w == 1) {
            const float radius = 15.0;
            const float dist = distance(inFragPos, lights[i].position.xyz);

            float att = clamp(1.0 - dist / radius, 0.0, 1.0);
            att *= att;

            add *= att;
        } else if(lights[i].position.w == 2) {
            const float shadow = shadowFilterPCF == 1 ? filterPCF(inShadowPos / inShadowPos.w) : textureProj(inShadowPos / inShadowPos.w, vec2(0, 0));
            add *= shadow;
        }

        diffuse += add;
    }

    vec3 matColor = pushConstants.c;
    if(textureSize(albedoSampler, 0).x > 1)
        matColor *= texture(albedoSampler, inUV).rgb;

    outColor = vec4((matColor * vec3(0.1)) + diffuse * matColor, 1.0);
}