#version 460
#extension GL_EXT_nonuniform_qualifier : require
#extension GL_EXT_scalar_block_layout : require
#extension GL_EXT_shader_16bit_storage : require

in Varyings {
    layout(location = 0) flat uint instance;
    layout(location = 1) vec3 positionVS;
    layout(location = 2) vec2 texCoord;
    layout(location = 3) vec3 normalVS;
    layout(location = 4) vec3 tangentVS;
    layout(location = 5) vec3 bitangentVS;
} var;

#include "main_common.glsl"

layout(location = 0) out vec4 fragColor;

const float INV_PI = 0.31830987;
const float IOR = 1.45;
const vec3 F90 = vec3(1.0);

vec3 fresnelSchlick(float dotVH, vec3 f0) {
    return mix(f0, F90, pow(1.0 - dotVH, 5.0));
}

float visibilityGGX(float dotNL, float dotNV, float alpha) {
    float alphaSquared = alpha * alpha;

    float vGGX = dotNL * sqrt(dotNV * dotNV * (1.0 - alphaSquared) + alphaSquared);
    float lGGX = dotNV * sqrt(dotNL * dotNL * (1.0 - alphaSquared) + alphaSquared);
    float GGX = vGGX + lGGX;
    return mix(0.0, 0.5 / GGX, GGX > 0.0);
}

float distributionGGX(float dotNH, float alpha) {
    float alphaSquared = alpha * alpha;
    float tmp = dotNH * dotNH * (alphaSquared - 1.0) + 1.0;
    return alphaSquared * INV_PI / (tmp * tmp);
}

vec3 toneMapAcesNarkowicz(vec3 color) {
    const float A = 2.51;
    const float B = 0.03;
    const float C = 2.43;
    const float D = 0.59;
    const float E = 0.14;
    return clamp((color * (A * color + B)) / (color * (C * color + D) + E), 0.0, 1.0);
}

vec3 lightOutgoingRadiance(
    vec3 viewDirectionVS, vec3 normalVS, float dotNV,
    vec3 baseColor, float alpha, float metallic, vec3 f0,
    vec3 incomingRadiance, vec3 lightDirectionVS
) {
    vec3 halfVectorVS = normalize(lightDirectionVS + viewDirectionVS);
    float dotVH = clamp(dot(viewDirectionVS, halfVectorVS), 0.0, 1.0);
    float dotNH = clamp(dot(normalVS, halfVectorVS), 0.0, 1.0);
    float dotNL = clamp(dot(normalVS, lightDirectionVS), 0.0, 1.0);

    vec3 fresnel = fresnelSchlick(dotVH, f0);
    float visibility = visibilityGGX(dotNL, dotNV, alpha);
    float distribution = distributionGGX(dotNH, alpha);

    vec3 scatteredFactor = (1.0 - fresnel) * (1.0 - metallic) * baseColor * INV_PI;
    vec3 reflectedFactor = fresnel * visibility * distribution;

    return (scatteredFactor + reflectedFactor) * incomingRadiance * dotNL;
}

vec4 texture2DAA(texture2D tex, vec2 texCoord) {
    vec2 size = vec2(textureSize(sampler2D(tex, _Sampler), 0).xy);
    vec2 texCoordPX = texCoord * size;
    vec2 seam = floor(texCoordPX + vec2(0.5));

    texCoordPX = (texCoordPX - seam) / fwidth(texCoordPX) + seam;
    texCoordPX = clamp(texCoordPX, seam - 0.5, seam + 0.5);

    texCoord = texCoordPX / size;
    return texture(sampler2D(tex, _Sampler), texCoord);
}

#define OBJECT _Object[var.instance]
#define MATERIAL _Materials[uint(OBJECT.material)]

void main() {
    vec4 baseColorTexel = texture2DAA(_Textures[uint(MATERIAL.baseColorTexture)], var.texCoord);
    vec4 occlusionRoughnessMetallicTexel = texture2DAA(_Textures[uint(MATERIAL.occlusionRoughnessMetallicTexture)], var.texCoord);
    vec4 normalTexel = texture2DAA(_Textures[uint(MATERIAL.normalTexture)], var.texCoord);
    vec4 emissiveTexel = texture2DAA(_Textures[uint(MATERIAL.emissiveTexture)], var.texCoord);

    vec3 baseColor = MATERIAL.baseColor * baseColorTexel.rgb;
    float occlusion = 1.0 + MATERIAL.occlusionTextureStrength * (occlusionRoughnessMetallicTexel.r - 1.0);
    float roughness = MATERIAL.roughness * occlusionRoughnessMetallicTexel.g;
    float metallic = MATERIAL.metallic * occlusionRoughnessMetallicTexel.b;
    vec3 emissive = MATERIAL.emissive * emissiveTexel.rgb;
    float ior = MATERIAL.ior;

    vec3 tangentVS = normalize(var.tangentVS);
    vec3 bitangentVS = normalize(var.bitangentVS);
    mat3 matrixTStoVS = mat3(tangentVS, bitangentVS, var.normalVS);
    vec3 normalTS = normalTexel.xyz;
    vec3 normalVS = normalize(matrixTStoVS * normalTS);

    vec3 positionVS = var.positionVS;
    vec3 viewDirectionVS = normalize(-positionVS);
    float dotNV = clamp(dot(normalVS, viewDirectionVS), 0.0, 1.0);
    float alpha = roughness * roughness;

    vec3 f0 = vec3(pow((ior - 1.0) / (ior + 1.0), 2.0));
    f0 = mix(f0, baseColor, metallic);

    vec3 outgoingRadiance = vec3(0.0);

    for (uint i = 0; i < _PointLights.count; i++) {
        PointLight light = _PointLights.lights[i];

        vec3 lightPositionVS = (_Global.matrixWStoVS * vec4(light.positionWS, 1.0)).xyz;
        vec3 lightDirectionVS = normalize(lightPositionVS - positionVS);
        float lightDistance = distance(positionVS, lightPositionVS);
        float lightAttenuation = 1.0 / (lightDistance * lightDistance);
        vec3 incomingRadiance = light.color * lightAttenuation;

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            baseColor, alpha, metallic, f0,
            incomingRadiance, lightDirectionVS
        );
    }

    for (int i = 0; i < _DirectionalLights.count; i++) {
        DirectionalLight light = _DirectionalLights.lights[i];

        vec3 lightDirectionVS = normalize((_Global.matrixWStoVS * vec4(-light.directionWS, 0.0)).xyz);
        vec3 incomingRadiance = light.color;

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            baseColor, alpha, metallic, f0,
            incomingRadiance, lightDirectionVS
        );
    }

    outgoingRadiance += _Global.ambientLight * baseColor * occlusion;

    vec3 toneMappedLinearColor = toneMapAcesNarkowicz(outgoingRadiance);
    vec3 toneMappedSrgbColor = pow(toneMappedLinearColor, vec3(1.0 / 2.2));

    fragColor = vec4(toneMappedSrgbColor, 1.0);
}