#import bevy_pbr::{
    pbr_fragment::pbr_input_from_standard_material,
    pbr_functions::{alpha_discard, calculate_tbn_mikktspace},
}

#ifdef PREPASS_PIPELINE
#import bevy_pbr::{
    prepass_io::{VertexOutput, FragmentOutput},
    pbr_deferred_functions::deferred_output,
}
#else
#import bevy_pbr::{
    forward_io::{VertexOutput, FragmentOutput},
    pbr_functions::{apply_pbr_lighting, main_pass_post_lighting_processing},
}
#endif

struct TerrainLayerUniform {
    base_colors: array<vec4<f32>, 4>,
    properties: array<vec4<f32>, 4>,
    uv_scales: vec4<f32>,
    base_texture_enabled: vec4<f32>,
}

@group(#{MATERIAL_BIND_GROUP}) @binding(100) var<uniform> terrain_layers: TerrainLayerUniform;
@group(#{MATERIAL_BIND_GROUP}) @binding(101) var base_color0: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(102) var base_sampler0: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(103) var normal0: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(104) var normal_sampler0: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(105) var base_color1: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(106) var base_sampler1: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(107) var normal1: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(108) var normal_sampler1: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(109) var base_color2: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(110) var base_sampler2: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(111) var normal2: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(112) var normal_sampler2: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(113) var base_color3: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(114) var base_sampler3: sampler;
@group(#{MATERIAL_BIND_GROUP}) @binding(115) var normal3: texture_2d<f32>;
@group(#{MATERIAL_BIND_GROUP}) @binding(116) var normal_sampler3: sampler;

fn layer_base(index: u32, uv: vec2<f32>) -> vec4<f32> {
    if index == 0u { return textureSample(base_color0, base_sampler0, uv); }
    if index == 1u { return textureSample(base_color1, base_sampler1, uv); }
    if index == 2u { return textureSample(base_color2, base_sampler2, uv); }
    return textureSample(base_color3, base_sampler3, uv);
}

fn layer_normal(index: u32, uv: vec2<f32>) -> vec3<f32> {
    var sample = vec3<f32>(0.5, 0.5, 1.0);
    if index == 0u { sample = textureSample(normal0, normal_sampler0, uv).xyz; }
    if index == 1u { sample = textureSample(normal1, normal_sampler1, uv).xyz; }
    if index == 2u { sample = textureSample(normal2, normal_sampler2, uv).xyz; }
    if index == 3u { sample = textureSample(normal3, normal_sampler3, uv).xyz; }
    return normalize(sample * 2.0 - 1.0);
}

@fragment
fn fragment(in: VertexOutput, @builtin(front_facing) is_front: bool) -> FragmentOutput {
    var pbr_input = pbr_input_from_standard_material(in, is_front);
#ifdef VERTEX_UVS_A
    let uv = in.uv;
#else
    let uv = vec2<f32>(0.0);
#endif
#ifdef VERTEX_COLORS
    var weights = max(in.color, vec4<f32>(0.0));
#else
    var weights = vec4<f32>(1.0, 0.0, 0.0, 0.0);
#endif
    let enabled = vec4<f32>(
        terrain_layers.properties[0].w,
        terrain_layers.properties[1].w,
        terrain_layers.properties[2].w,
        terrain_layers.properties[3].w,
    );
    weights *= enabled;
    let total = dot(weights, vec4<f32>(1.0));
    weights = select(vec4<f32>(1.0, 0.0, 0.0, 0.0), weights / total, total > 0.0001);

    var color = vec4<f32>(0.0);
    var normal_ts = vec3<f32>(0.0);
    var metallic = 0.0;
    var roughness = 0.0;
    for (var index = 0u; index < 4u; index += 1u) {
        let weight = weights[index];
        let layer_uv = uv * terrain_layers.uv_scales[index];
        let base_sample = mix(
            vec4<f32>(1.0),
            layer_base(index, layer_uv),
            terrain_layers.base_texture_enabled[index],
        );
        color += base_sample * terrain_layers.base_colors[index] * weight;
        let normal_enabled = terrain_layers.properties[index].z;
        normal_ts += mix(vec3<f32>(0.0, 0.0, 1.0), layer_normal(index, layer_uv), normal_enabled) * weight;
        metallic += terrain_layers.properties[index].x * weight;
        roughness += terrain_layers.properties[index].y * weight;
    }
    pbr_input.material.base_color = alpha_discard(pbr_input.material, color);
    pbr_input.material.metallic = clamp(metallic, 0.0, 1.0);
    pbr_input.material.perceptual_roughness = clamp(roughness, 0.001, 1.0);
#ifdef VERTEX_TANGENTS
    let tbn = calculate_tbn_mikktspace(pbr_input.world_normal, in.world_tangent);
    let n = normalize(normal_ts);
    pbr_input.N = normalize(n.x * tbn[0] + n.y * tbn[1] + n.z * tbn[2]);
#endif
#ifdef PREPASS_PIPELINE
    return deferred_output(in, pbr_input);
#else
    var out: FragmentOutput;
    out.color = apply_pbr_lighting(pbr_input);
    out.color = main_pass_post_lighting_processing(pbr_input, out.color);
    return out;
#endif
}