#ifndef WGSL_WGPU
diagnostic(off,derivative_uniformity);
#endif

#define USE_MATERIAL_ID

#define FEATURE_ENVIRONMENT_PROBE
#define FEATURE_TEXTURED
#define FEATURE_GRADIENT
#define FEATURE_GRADIENT_4_STOPS

#ifdef GRADIENT_4_STOPS
#define GRADIENT
#endif

#ifndef PINK
#define PINK vec4<f32>(1.0, 0.0, 1.0, 1.0)
#endif

#define USE_MATERIAL_INDEX
#include "lib/Compatibility.wgsl"
#include "lib/Uniforms.wgsl"

#define USE_NDC_COORDINATES

#ifdef ENVIRONMENT_PROBE
#define GLOBAL_ILLUMINATION
@group(0) @binding(17) var specularProbeAtlas: texture_2d_array<f16>;
@group(0) @binding(18) var reflectionProbeSampler: sampler;
#endif

#if defined(TEXTURED) || defined(GRADIENT) || defined(ENVIRONMENT_PROBE)
#define NEED_MATERIAL
#endif

#if defined(TEXTURED) || defined(ENVIRONMENT_PROBE)
#define NEED_UV
#endif

#ifdef NEED_MATERIAL
/**
 * Material definition
 *
 * Contains texture and gradient data.
 */
struct Material {
    #ifdef GRADIENT_4_STOPS
    colorStop3: vec4<f16>,
    colorStop2: vec4<f16>,
    #endif
    #ifdef GRADIENT
    colorStop1: vec4<f16>,
    colorStop0: vec4<f16>,
    #endif
    #ifdef TEXTURED
    texture: u16,
    #endif
    #ifdef ENVIRONMENT_PROBE
    mip: u8,
    #endif
    #ifdef NEED_UV
    rotationY: f8,
    #endif
};

#include "lib/Materials.wgsl"

fn decodeMaterial(matIndex: u32) -> Material {
    {{decoder}}
    return mat;
}

#include "lib/Quaternion.wgsl"
#include "lib/Math.wgsl"
#include "lib/CoordinateSystems.wgsl"
#include "lib/Color.wgsl"

#ifdef TEXTURED
#include "lib/Textures.wgsl"

/**
 * Rescale equirectangular UVs to remove atlas seam.
 *
 * @todo: This function should be removed when:
 *    - The uncompressed atlas supports padding
 *    - The compressed atlas padding isn't mixed with the texture
 */
fn rescaleUV(uv: vec2<f32>, size: vec2<f32>) -> vec2<f32> {
    /* Rescale and intercept uv to reduce the seam from
     * missing bound and compression error. */
    return (uv*(size - vec2<f32>(1.0)) + vec2<f32>(0.5))/size;
}

/**
 * Read texture with seam reduction from texture atlas
 *
 * @param mat Material data
 * @param uv Equirectangular UVs in [0...1]
 */
fn shadeTexture(mat: Material, uv: vec2<f32>) -> vec4<f32> {
    if(mat.texture == 0u) {
        return PINK;
    }

    let b: Bounds = unpackBounds(mat.texture);
    /* Convert UV to [0; 1] and rescale. */
    let scaledUV: vec2<f32> = rescaleUV(uv, vec2<f32>(b.bounds.zw));

    #ifdef COMPRESSED_ATLAS
    if(!bool(b.layer & COMPRESSED)) {
        #ifdef TEXTURE_STREAMING
        let maxMip: u8 = maxLod(b.bounds.zw);
        /* Mip is on purpose computed with un-normalized UVs. */
        let mip: u8 = mipMapLevelEquirectangular(uv - vec2<f32>(0.5), vec2<f32>(b.bounds.zw), maxMip);
        return tileStackLevel(textureAtlasCompressed, b.bounds.xy, b.bounds.zw, uv, mip);
        #else
        /* The mip is fixed, no need to perform extra computation. */
        return textureAtlasFlat(textureAtlasCompressed, uv, b.bounds, b.layer & 0x7Fu);
        #endif
    }
    #endif
    #ifdef UNCOMPRESSED_ATLAS
    return textureAtlasFlat(textureAtlasUncompressed, uv, b.bounds, b.layer);
    #else
    /* Shouldn't happen, return debug pink */
    return vec4<f32>(1.0, 0.0, 1.0, 1.0);
    #endif
}
#endif

#ifdef GRADIENT
fn shadeGradient(mat: Material, value: f32) -> vec4<f32> {
    #ifdef GRADIENT_4_STOPS
    return gradient4(mat.colorStop0, mat.colorStop1, mat.colorStop2, mat.colorStop3, value);
    #else
    return mix(mat.colorStop0, mat.colorStop1, value);
    #endif
}
#endif

#endif

@fragment
fn main(
    @location(0) ndcCoordinates: vec2<f32>
) -> @location(0) vec4<f32> {

    #ifdef NEED_MATERIAL
    let mat: Material = decodeMaterial(drawUniforms.materialIndex);

    let unprojPoint: vec3<f32> = (inverseProjectionMatrix*vec4(ndcCoordinates, 0.0, 1.0)).xyz;
    let direction: vec3<f32> = normalize(quat_transformVector(viewToWorld[0], unprojPoint));
    #if defined(TEXTURED) || defined(ENVIRONMENT_PROBE)
    var uv: vec2<f32> = cartesianToEquirectangular(direction);
    uv.x = fract(uv.x + mat.rotationY);
    #endif

    var outColor = vec4<f32>(1.0);

    #ifdef TEXTURED
    outColor *= shadeTexture(mat, uv);
    #endif

    #ifdef ENVIRONMENT_PROBE
    /* Environment is always encoded as linear */
    let env: vec3<f32> = textureSampleLevel(specularProbeAtlas, reflectionProbeSampler, uv, 0u, f32(mat.mip)).rgb;
    outColor *= linearToSrgb4(vec4<f32>(env, 1.0));
    #endif

    #ifdef GRADIENT
    /* Remap direction to [0, 1] and sample gradient */
    outColor *= shadeGradient(mat, direction.y * 0.5 + 0.5);
    #endif

    #else

    /* Use error color if the shader isn't set for textures
     * nor for gradients */
    outColor = PINK;

    #endif

    return outColor;
}
