struct Vertex {
    @location(0) positionOS: vec3<f32>,
    @location(1) texCoord: vec2<f32>,
    @location(2) normalOS: vec3<f32>,
    @location(3) tangentOS: vec4<f32>,
}

struct Varyings {
    @builtin(position) positionCS: vec4<f32>,
    @location(0) positionVS: vec3<f32>,
    @location(1) texCoord: vec2<f32>,
    @location(2) normalVS: vec3<f32>,
    @location(3) tangentVS: vec3<f32>,
    @location(4) bitangentVS: vec3<f32>,
}

struct PointLight {
    positionWS: vec3<f32>,
    color: vec3<f32>,
}

struct DirectionalLight {
    directionWS: vec3<f32>,
    color: vec3<f32>,
}

struct GlobalUniforms {
    matrixWStoVS: mat4x4<f32>,
    matrixVStoCS: mat4x4<f32>,
    ambientLight: vec3<f32>,
    pointLightCount: u32,
    directionalLightCount: u32,
}

struct ObjectUniforms {
    matrixOStoWS: mat4x4<f32>,
    matrixOStoWSNormal: mat4x4<f32>,
}

// --- GROUP 0 --- GLOBAL ------------------------------------------------------

@group(0) @binding(0) var<uniform> _Global: GlobalUniforms;

@group(0) @binding(1) var<storage> _PointLights: array<PointLight>;
@group(0) @binding(2) var<storage> _DirectionalLights: array<DirectionalLight>;

@group(0) @binding(3) var _Sampler: sampler;
@group(0) @binding(4) var _BaseColorTexture: texture_2d_array<f32>;
@group(0) @binding(5) var _OcclusionRoughnessMetallicTexture: texture_2d_array<f32>;
@group(0) @binding(6) var _NormalTexture: texture_2d_array<f32>;

// --- GROUP 1 --- PER MATERIAL ------------------------------------------------

@group(1) @binding(0) var<uniform> _TextureIndex: u32;

// --- GROUP 2 --- PER OBJECT --------------------------------------------------

@group(2) @binding(0) var<uniform> _Object: ObjectUniforms;

// -----------------------------------------------------------------------------

const INV_PI: f32 = 0.31830987;
const IOR: f32 = 1.45;
const DIELECTRIC_F0: vec3<f32> = vec3(pow((IOR - 1.0) / (IOR + 1.0), 2.0));
const F90 = vec3(1.0);

fn fresnelSchlick(dotVH: f32, f0: vec3<f32>) -> vec3<f32> {
    return mix(f0, F90, pow(1.0 - dotVH, 5.0));
}

fn visibilityGGX(dotNL: f32, dotNV: f32, alpha: f32) -> f32 {
    let alphaSquared = alpha * alpha;

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

fn distributionGGX(dotNH: f32, alpha: f32) -> f32 {
    let alphaSquared = alpha * alpha;
    let tmp = dotNH * dotNH * (alphaSquared - 1.0) + 1.0;
    return alphaSquared * INV_PI / (tmp * tmp);
}

fn toneMapAcesNarkowicz(color: vec3<f32>) -> vec3<f32> {
    const A: f32 = 2.51;
    const B: f32 = 0.03;
    const C: f32 = 2.43;
    const D: f32 = 0.59;
    const E: f32 = 0.14;
    return saturate((color * (A * color + B)) / (color * (C * color + D) + E));
}

fn lightOutgoingRadiance(
    viewDirectionVS: vec3<f32>, normalVS: vec3<f32>, dotNV: f32,
    baseColor: vec3<f32>, alpha: f32, metallic: f32, f0: vec3<f32>,
    incomingRadiance: vec3<f32>, lightDirectionVS: vec3<f32>,
) -> vec3<f32> {
    let halfVectorVS = normalize(lightDirectionVS + viewDirectionVS);
    let dotVH = saturate(dot(viewDirectionVS, halfVectorVS));
    let dotNH = saturate(dot(normalVS, halfVectorVS));
    let dotNL = saturate(dot(normalVS, lightDirectionVS));

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

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

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

@vertex
fn vert(vertex: Vertex) -> Varyings {
    let positionWS = (_Object.matrixOStoWS * vec4(vertex.positionOS, 1.0)).xyz;
    let positionVS = (_Global.matrixWStoVS * vec4(positionWS, 1.0)).xyz;
    let positionCS = _Global.matrixVStoCS * vec4(positionVS, 1.0);

    let normalWS = normalize((_Object.matrixOStoWSNormal * vec4(vertex.normalOS, 0.0)).xyz);
    let normalVS = normalize((_Global.matrixWStoVS * vec4(normalWS, 0.0)).xyz);

    let tangentWS = normalize((_Object.matrixOStoWS * vec4(vertex.tangentOS.xyz, 0.0)).xyz);
    let tangentVS = normalize((_Global.matrixWStoVS * vec4(tangentWS, 0.0)).xyz);

    let bitangentVS = vertex.tangentOS.w * normalize(cross(normalVS, tangentVS));

    var output: Varyings;
    output.positionCS = positionCS;
    output.positionVS = positionVS;
    output.texCoord = vertex.texCoord;
    output.normalVS = normalVS;
    output.tangentVS = tangentVS;
    output.bitangentVS = bitangentVS;

    return output;
}

@fragment
fn frag(fragment: Varyings) -> @location(0) vec4<f32> {
    let baseColorTexel = textureSample(_BaseColorTexture, _Sampler, fragment.texCoord, _TextureIndex);
    let occlusionRoughnessMetallicTexel = textureSample(_OcclusionRoughnessMetallicTexture, _Sampler, fragment.texCoord, _TextureIndex);
    let normalTextureTexel = textureSample(_NormalTexture, _Sampler, fragment.texCoord, _TextureIndex);

    let baseColor = baseColorTexel.rgb;
    let occlusion = occlusionRoughnessMetallicTexel.r;
    let roughness = occlusionRoughnessMetallicTexel.g;
    let metallic = occlusionRoughnessMetallicTexel.b;

    let tangentVS = normalize(fragment.tangentVS);
    let bitangentVS = normalize(fragment.bitangentVS);
    let matrixTStoVS = mat3x3(tangentVS, bitangentVS, fragment.normalVS);
    let normalTS = normalTextureTexel.xyz * 2.0 - 1.0;
    let normalVS = normalize(matrixTStoVS * normalTS);

    let positionVS = fragment.positionVS;
    let viewDirectionVS = normalize(-positionVS);
    let dotNV = saturate(dot(normalVS, viewDirectionVS));
    let alpha = roughness * roughness;

    let f0 = mix(DIELECTRIC_F0, baseColor, metallic);

    var outgoingRadiance = vec3(0.0);

    for (var i: u32 = 0; i < _Global.pointLightCount; i++) {
        let light = _PointLights[i];

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

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

    for (var i: u32 = 0; i < _Global.directionalLightCount; i++) {
        let light = _DirectionalLights[i];

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

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

    outgoingRadiance += _Global.ambientLight * baseColor * occlusion;

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

    return vec4(toneMappedSrgbColor, 1.0);
}