#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 = 3, binding = 0) uniform uFog {
    vec4 fogColor;
    float fogEnd;
    float fogDensity;
};

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;

const float PI = 3.14159265359f;
Material mat = materials[materialIndex];
vec4 baseColor = mat.hasTextures.y != 0 ? texture(baseColorTexture, iFragTexCoord) : mat.baseColorFactor;
vec3 albedo = baseColor.rgb;
float metallic = mat.hasTextures.z != 0 ? texture(metallicRoughnessTexture, iFragTexCoord).b : mat.metallicFactor == 1.0f ? mat.metallicFactor - 0.25f : mat.metallicFactor;
float roughness = mat.hasTextures.z != 0 ? texture(metallicRoughnessTexture, iFragTexCoord).g : mat.roughnessFactor;

// fog code inspired by https://youtu.be/oQksg57qsRA?si=iFEi0f9zL1c-w_Gb
float calculateFog() {
    float distanceCameraToFragment = length(-iFragPos);
    float distanceRatio = 4.0f * distanceCameraToFragment / fogEnd;
    float fogFactor = exp(-distanceRatio * distanceRatio * fogDensity * fogDensity);

    return fogFactor;
}

// PBR inspired by
// - https://learnopengl.com/PBR/Theory
// - https://rtarun9.github.io/blogs/physically_based_rendering/

// approximates the reflected light
vec3 fresnelSchlick(float vDotH, vec3 F0) {
    return F0 + (1.0f - F0) * pow(clamp(1.0f - vDotH, 0.0f, 1.0f), 5.0f);
}

// based on roughness approximates alignment of the surface's microfacets
float normalDistribution(vec3 normal, vec3 halfWayVector) {
    float alpha = roughness * roughness;
    float nDotH = max(dot(normal, halfWayVector), 0.0f);

    return (alpha * alpha) / (PI * pow((nDotH * nDotH) * (alpha * alpha - 1.0f) + 1.0f, 2.0f));
}

// approximates the relative surface area where micro surface-details overshadow each other
float geometrySchlick(float x) {
    // k is a remapping of alpha, differs depending on if image based lighting is used or not (we aren't using IBL)
    float r = roughness + 1.0f;
    float k = (r * r) / 8.0f;

    return x / (x * (1.0f - k) + k);
}

// take view (geometry obstruction) and light (geometry shadowing) directions into account
float geometrySmith(vec3 normal, vec3 viewDirection, vec3 lightDirection) {
    float nDotV = max(dot(normal, viewDirection), 0.0f);
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return geometrySchlick(nDotV) * geometrySchlick(nDotL);
}

// PBR lighting calcs (btw calcs stands for calculations, I'm just using slang)
vec3 cookTorrenceBRDF(vec3 normal, vec3 viewDirection, vec3 lightDirection, vec3 halfWayVector) {
    vec3 F0 = mix(vec3(0.04f), albedo, metallic); // F0 = base reflectivity
    vec3 F = fresnelSchlick(max(dot(halfWayVector, viewDirection), 0.0f), F0);
    float NDF = normalDistribution(normal, halfWayVector);
    float G = geometrySmith(normal, viewDirection, lightDirection);

    vec3 specular = (NDF * G * F) / (4.0f * max(dot(normal, viewDirection), 0.0f) * max(dot(normal, lightDirection), 0.0f) + 0.0001f); // last constant to prevent divide by zero

    vec3 kd = (vec3(1.0f) - F) * (1.0f - metallic);
    vec3 diffuse = kd * (albedo / PI);

    return diffuse + specular;
}

vec3 calculateDirLight(int lightIndex) {
    vec3 normal = normalize(iNormal);
    vec3 viewDirection = normalize(-iFragPos);
    vec3 lightDirection = 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 halfWayVector = normalize(lightDirection + viewDirection);

    vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
    vec3 radiance = lights[lightIndex].color.rgb;
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return BDRF * radiance * nDotL;
}

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

    vec3 normal = normalize(iNormal);
    vec3 viewDirection = normalize(-iFragPos);
    vec3 lightDirection = normalize(lightPos - iFragPos);
    vec3 halfWayVector = normalize(lightDirection + viewDirection);

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

    vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
    vec3 radiance = lights[lightIndex].color.rgb * attenuation;
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return BDRF * radiance * nDotL;
}

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

    vec3 result = vec3(0.0);

    float theta = dot(lightDirection, spotDirection);
    if (theta > lights[lightIndex].params.z) {
        vec3 normal = normalize(iNormal);
        vec3 viewDirection = normalize(-iFragPos);
        vec3 halfWayVector = normalize(lightDirection + viewDirection);

        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);
        intensity *= intensity;

        vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
        vec3 radiance = lights[lightIndex].color.rgb * attenuation;
        float nDotL = max(dot(normal, lightDirection), 0.0f);

        result = BDRF * radiance * nDotL * intensity;
    }

    return result;
}

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 lighting = vec3(0.0f);

    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);
        }
    }

    vec3 ambient = vec3(0.03f) * albedo;
    float shadow = calculateShadow();

    vec3 result = ambient + (1.0f - shadow) * lighting;

    result = result / (result + vec3(1.0f)); // HDR tonemapping
    result = pow(result, vec3(1.0f / 2.2f)); // gamma correction

    float fogFactor = calculateFog();
    result = mix(fogColor.rgb, result, fogFactor);

    oColor = vec4(result, 1.0f);
}
