#version 450

struct Light {
    vec4 color;
    vec4 dir;
    vec4 pos;
    vec4 params;
};

struct Material {
    vec4 hasTextures; // normal, color, metallicRoughness
    vec4 baseColorFactor;
    float metallicFactor;
    float roughnessFactor;
};

layout(set = 0, binding = 0) uniform Global {
    mat4 view;
    mat4 proj;
};

layout(std430, set = 1, binding = 0) readonly buffer Materials {
    Material materials[];
};

layout(set = 1, binding = 1) uniform sampler2D normalTexture;
layout(set = 1, binding = 2) uniform sampler2D baseColorTexture;
layout(set = 1, binding = 3) uniform sampler2D metallicRoughnessTexture;

layout(std430, set = 2, binding = 0) readonly buffer Lights {
    Light lights[];
};

layout(set = 4, binding = 0) uniform sampler2D shadowMap;

layout(location = 0) in vec3 iFragPos;
layout(location = 1) in vec3 iNormal;
layout(location = 2) in vec2 iFragTexCoord;
layout(location = 4) in vec4 iLightFragPos;
layout(location = 5) in flat int materialIndex;

layout(location = 0) out vec4 oColor;

Material mat = materials[materialIndex];
vec4 baseColor = mat.hasTextures.y != 0 ? texture(baseColorTexture, iFragTexCoord) : mat.baseColorFactor;
float metallic = mat.hasTextures.z != 0 ? texture(metallicRoughnessTexture, iFragTexCoord).b : mat.metallicFactor;
float roughness = mat.hasTextures.z != 0 ? texture(metallicRoughnessTexture, iFragTexCoord).g : mat.roughnessFactor;
vec3 color = baseColor.rgb;
float alpha = pow(2.0f, metallic * 8.0f);
float kd = roughness;
float ks = 1.0f - kd;
float ka = 0.05f;


vec3 calculateDirLight(int lightIndex) {
    vec3 N = normalize(iNormal);
    vec3 L = normalize(vec3(view * vec4(-lights[lightIndex].dir.xyz, 0.0f))); // transform from world to view space, 4th element is 0 because it represents a direction, not position
    vec3 V = normalize(-iFragPos);
    vec3 H = normalize(L + V);

    float NL = max(dot(N, L), 0.0f);
    float VR = pow(max(dot(N, H), 0.0f), alpha);

    vec3 diffuse  = kd * NL * lights[lightIndex].color.rgb * color.rgb;
    vec3 specular = ks * VR * lights[lightIndex].color.rgb;

    return diffuse + specular;
}

vec3 calculatePointLight(int lightIndex) {
    vec3 lightPos = vec3(view * vec4(lights[lightIndex].pos.xyz, 1.0f)); // transformed to view space

    vec3 N = normalize(iNormal);
    vec3 L = normalize(lightPos - iFragPos);
    vec3 V = normalize(-iFragPos);
    vec3 H = normalize(L + V);

    float NL = max(dot(N, L), 0.0f);
    float VR = pow(max(dot(N, H), 0.0f), alpha);

    float distance = length(lightPos - iFragPos);
    float attenuation = 1.0f / (1.0f + lights[lightIndex].params.x * distance + lights[lightIndex].params.y * (distance * distance));

    vec3 diffuse  = kd * NL * lights[lightIndex].color.rgb * color.rgb;
    vec3 specular = ks * VR * lights[lightIndex].color.rgb;

    diffuse  *= attenuation;
    specular *= attenuation;

    return diffuse + specular;
}

vec3 calculateSpotLight(int lightIndex) {
    vec3 lightPos = vec3(view * vec4(lights[lightIndex].pos.xyz, 1.0f)); // transformed to view space
    vec3 L = normalize(lightPos - iFragPos);
    vec3 spotDir = -normalize(vec3(view * vec4(lights[lightIndex].dir.xyz, 0.0f)));

    vec3 diffuse = vec3(0.0);
    vec3 specular = vec3(0.0);

    float theta = dot(L, spotDir);
    if (theta > lights[lightIndex].params.z)
    {
        vec3 N = normalize(iNormal);
        vec3 V = normalize(-iFragPos);
        vec3 H = normalize(L + V);

        float NL = max(dot(N, L), 0.0f);

        float NH = pow(max(dot(N, H), 0.0f), alpha);

        float distance = length(lightPos - iFragPos);
        float attenuation = 1.0f / (1.0f  + lights[lightIndex].params.x * distance + lights[lightIndex].params.y * (distance * distance));

        float intensity = smoothstep(lights[lightIndex].params.w, lights[lightIndex].params.z, theta);

        diffuse  = kd * NL * lights[lightIndex].color.rgb * color.rgb;
        specular = ks * NH * lights[lightIndex].color.rgb;

        diffuse  *= attenuation * intensity;
        specular *= attenuation * intensity;
    }

    return diffuse + specular;
}

float calculateShadow() {
    // perspective divide (needed in case of perspective projection)
    vec3 projCoords = iLightFragPos.xyz / iLightFragPos.w;
    // for some reason Vulkans NDC are in [-1, 1] for x and y, but in [0, 1] for z, so we convert xy to [0, 1]
    projCoords.xy = projCoords.xy * 0.5f + 0.5f;

    vec2 shadowMapUV = projCoords.xy;
    float currentDepth = projCoords.z;

    // return if fragment is outside shadow map's projection,
    // as in behind the far plane, left/right/top/bottom are taken care of by the sampler
    if (currentDepth > 1.0f) { return 0.0f; }

    // Percentage-Close Filtering (PCF), basically sampling neighboring texels and averaging the result
    float shadow = 0.0f;
    vec2 texelSize = 1.0f / textureSize(shadowMap, 0);

    for (int x = -1; x <= 1; x++) {
        for(int y = -1; y <= 1; y++) {
            vec2 offset = vec2(x, y) * texelSize;
            float pcfDepth = texture(shadowMap, shadowMapUV + offset).r;
            shadow += currentDepth > pcfDepth ? 1.0f : 0.0f;
        }
    }

    shadow /= 9.0f;

    return shadow;
}

void main() {
    vec3 result = vec3(0.0f);
    vec3 lighting = vec3(0.0f);

    vec3 ambient = ka * color.rgb;

    result += ambient;

    for (int i = 0; i < lights.length(); i++) {
        if (lights[i].dir.w == 0.0f) {
            lighting += calculateDirLight(i);
        } else if (lights[i].dir.w == 1.0f) {
            lighting += calculatePointLight(i);
        } else if (lights[i].dir.w == 2.0f) {
            lighting += calculateSpotLight(i);
        }
    }

    float shadow = calculateShadow();
    result += (1.0f - shadow) * lighting;
    oColor = vec4(result, 1.0f);
}
