import "includes/math";
import "includes/structs";

struct Vertex {
    [[vk::location(0)]] float3 positionOS;
    [[vk::location(1)]] float2 texCoord;
    [[vk::location(2)]] float3 normalOS;
    [[vk::location(3)]] float4 tangentOS;
}

struct Varyings {
    float4 positionCS : SV_Position;
    nointerpolation uint instance;
    float3 positionVS;
    float2 texCoord;
    float3 normalVS;
    float3 tangentVS;
    float3 bitangentVS;
}

[[vk::binding(0, 0)]] ConstantBuffer<GlobalUniforms> _Global;
[[vk::binding(1, 0)]] StructuredBuffer<PointLight> _PointLights;
[[vk::binding(2, 0)]] StructuredBuffer<DirectionalLight> _DirectionalLights;
[[vk::binding(3, 0)]] StructuredBuffer<Material> _Materials;
[[vk::binding(4, 0)]] SamplerState _Sampler;
[[vk::binding(5, 0)]] Texture2D _Textures[];

[[vk::binding(0, 1)]] StructuredBuffer<ObjectUniforms> _Objects;

[shader("vertex")]
Varyings vertexMain(Vertex vert, uint instanceID : SV_VulkanInstanceID) {
    Varyings out;

    let object = _Objects[instanceID];

    let positionWS = mul(object.matrixOStoWS, float4(vert.positionOS, 1.0)).xyz;
    let positionVS = mul(_Global.matrixWStoVS, float4(positionWS, 1.0)).xyz;
    let positionCS = mul(_Global.matrixVStoCS, float4(positionVS, 1.0));

    let normalWS = normalize(mul(object.matrixOStoWSNormal, float4(vert.normalOS, 0.0)).xyz);
    let normalVS = normalize(mul(_Global.matrixWStoVS, float4(normalWS, 0.0)).xyz);

    let tangentWS = normalize(mul(object.matrixOStoWSNormal, float4(vert.tangentOS.xyz, 0.0)).xyz);
    let tangentVS = normalize(mul(_Global.matrixWStoVS, float4(tangentWS, 0.0)).xyz);

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

    out.positionCS = positionCS;
    out.instance = instanceID;
    out.positionVS = positionVS;
    out.texCoord = vert.texCoord;
    out.normalVS = normalVS;
    out.tangentVS = tangentVS;
    out.bitangentVS = bitangentVS;
    return out;
}

[shader("fragment")]
float4 fragmentMain(Varyings frag) : SV_Target {
    let material = _Materials[_Objects[frag.instance].material];

    let baseColorTexel = texture2DAA(_Textures[material.baseColorTexture], _Sampler, frag.texCoord);
    if (baseColorTexel.a < 0.5) {
        discard;
    }

    let occlusionRoughnessMetallicTexel = texture2DAA(_Textures[material.occlusionRoughnessMetallicTexture], _Sampler, frag.texCoord);
    let normalTexel = texture2DAA(_Textures[material.normalTexture], _Sampler, frag.texCoord);
    let emissiveTexel = texture2DAA(_Textures[material.emissiveTexture], _Sampler, frag.texCoord);

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

    let tangentVS = normalize(frag.tangentVS);
    let bitangentVS = normalize(frag.bitangentVS);
    let matrixTStoVS = transpose(float3x3(tangentVS, bitangentVS, frag.normalVS));
    let normalTS = normalTexel.xyz;
    let normalVS = normalize(mul(matrixTStoVS, normalTS));

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

    var f0 = float3(pow((ior - 1.0) / (ior + 1.0), 2.0));
    f0 = lerp(f0, baseColor, metallic);

    let surface = Surface(baseColor, alpha, metallic, f0);

    var outgoingRadiance = float3(0.0);

    let pointLightCount = _Global.pointLightCount;
    for (uint i = 0; i < pointLightCount; i++) {
        let light = _PointLights[i].getIncoming(_Global.matrixWStoVS, positionVS);

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            surface, light,
        );
    }

    let directionalLightCount = _Global.directionalLightCount;
    for (uint i = 0; i < directionalLightCount; i++) {
        let light = _DirectionalLights[i].getIncoming(_Global.matrixWStoVS);

        outgoingRadiance += lightOutgoingRadiance(
            viewDirectionVS, normalVS, dotNV,
            surface, light,
        );
    }

    outgoingRadiance += _Global.ambientLight * baseColor * occlusion;

    let toneMappedLinearColor = toneMapAcesNarkowicz(outgoingRadiance);
    let toneMappedSrgbColor = linearToSrgbColor(toneMappedLinearColor);

    return float4(toneMappedSrgbColor, 1.0);
}