import package::util::{
    VertexInput,
    VertexOutput,
    uniforms,
    all_spotlights_lighting,
    blue_noise_texture,
    blue_noise_sampler,
    snow_light_texture,
    snow_light_sampler,
    TERRAIN_BOUNDS
};

@vertex
fn vs_main(input: VertexInput) -> VertexOutput {
    var output: VertexOutput;

    let instance_model = mat4x4<f32>(
        input.instance_model_0,
        input.instance_model_1,
        input.instance_model_2,
        input.instance_model_3
    );

    let world_pos = instance_model * vec4<f32>(input.position, 1.0);
    output.world_position = world_pos.xyz;
    output.clip_position = uniforms.projection * uniforms.view * world_pos;

    let normal_matrix = mat3x3<f32>(
        instance_model[0].xyz,
        instance_model[1].xyz,
        instance_model[2].xyz
    );
    output.world_normal = normalize(normal_matrix * input.normal);

    output.dissolve_amount = input.instance_dissolve;
    output.uv = input.uv;

    return output;
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    if uniforms.enable_dissolve == 1u && in.dissolve_amount > 0.0 {
        let screen_pos = in.clip_position.xy / in.clip_position.w;
        let noise_uv = screen_pos * 0.5 + 0.5;
        let noise = textureSample(blue_noise_texture, blue_noise_sampler, noise_uv * 10.0).r;
        if in.dissolve_amount > noise {
            discard;
        }
    }

    let tile_scale = uniforms.tile_scale;
    let spotlight_strokes = all_spotlights_lighting(in.world_position, in.clip_position, in.world_normal, tile_scale);

    var brightness = spotlight_strokes;

    if uniforms.enable_snow_light == 1u {
        let terrain_uv = (vec2<f32>(in.world_position.x, in.world_position.z) + TERRAIN_BOUNDS * 0.5) / TERRAIN_BOUNDS;
        let snow_light = textureSample(snow_light_texture, snow_light_sampler, terrain_uv).r;

        if snow_light > 0.05 {
            let screen_pos = in.clip_position.xy / in.clip_position.w;
            let blue_noise_uv = screen_pos * 0.5 + 0.5;
            let blue_noise = textureSample(blue_noise_texture, blue_noise_sampler, blue_noise_uv * 10.0).r;
            let blue_step = step(blue_noise, snow_light / 30.0);

            brightness = max(brightness, blue_step);
        }
    }

    switch uniforms.debug_mode {
        case 1u { return vec4(in.world_normal * 0.5 + 0.5, 1.0); }
        case 2u { return vec4(fract(in.uv), 0.0, 1.0); }
        case 3u {
            let d = in.clip_position.z / in.clip_position.w;
            let lin = (d - 0.95) / 0.05;
            return vec4(vec3(lin), 1.0);
        }
        case 4u { return vec4(1.0); }
        default { return vec4(brightness, brightness, brightness, 1.0); }
    }
}

@group(1) @binding(0)
var heightmap_texture: texture_2d<f32>;

@group(1) @binding(1)
var heightmap_sampler: sampler;

@group(1) @binding(2)
var snow_depth_texture: texture_2d<f32>;

@group(1) @binding(3)
var snow_depth_sampler: sampler;

fn sample_snow_depth_bilinear(uv: vec2<f32>) -> f32 {
    let tex_size = vec2<f32>(textureDimensions(snow_depth_texture, 0));
    let texel_coord = uv * tex_size - 0.5;
    let base = vec2<i32>(floor(texel_coord));
    let f = fract(texel_coord);

    let s00 = textureLoad(snow_depth_texture, clamp(base, vec2<i32>(0), vec2<i32>(tex_size) - 1), 0).r;
    let s10 = textureLoad(snow_depth_texture, clamp(base + vec2<i32>(1, 0), vec2<i32>(0), vec2<i32>(tex_size) - 1), 0).r;
    let s01 = textureLoad(snow_depth_texture, clamp(base + vec2<i32>(0, 1), vec2<i32>(0), vec2<i32>(tex_size) - 1), 0).r;
    let s11 = textureLoad(snow_depth_texture, clamp(base + vec2<i32>(1, 1), vec2<i32>(0), vec2<i32>(tex_size) - 1), 0).r;

    return mix(mix(s00, s10, f.x), mix(s01, s11, f.x), f.y);
}

@vertex
fn vs_snow_main(in: VertexInput) -> VertexOutput {
    let instance_model = mat4x4<f32>(
        in.instance_model_0,
        in.instance_model_1,
        in.instance_model_2,
        in.instance_model_3,
    );

    let instance_world = instance_model * vec4<f32>(in.position, 1.0);
    let world_xz = instance_world.xz;
    let level_bias = instance_world.y;

    let uv = (world_xz + TERRAIN_BOUNDS * 0.5) / TERRAIN_BOUNDS;

    let terrain_height = textureSampleLevel(heightmap_texture, heightmap_sampler, uv, 0.0).r;
    let snow_depth = sample_snow_depth_bilinear(uv) * 5.0;

    let world_y = terrain_height + snow_depth + 0.2 + level_bias;
    let world_position = vec3<f32>(world_xz.x, world_y, world_xz.y);

    let texel_size = 1.0 / vec2<f32>(textureDimensions(heightmap_texture, 0));
    let h_right = textureSampleLevel(heightmap_texture, heightmap_sampler, uv + vec2<f32>(texel_size.x, 0.0), 0.0).r + sample_snow_depth_bilinear(uv + vec2<f32>(texel_size.x, 0.0));
    let h_left = textureSampleLevel(heightmap_texture, heightmap_sampler, uv - vec2<f32>(texel_size.x, 0.0), 0.0).r + sample_snow_depth_bilinear(uv - vec2<f32>(texel_size.x, 0.0));
    let h_up = textureSampleLevel(heightmap_texture, heightmap_sampler, uv + vec2<f32>(0.0, texel_size.y), 0.0).r + sample_snow_depth_bilinear(uv + vec2<f32>(0.0, texel_size.y));
    let h_down = textureSampleLevel(heightmap_texture, heightmap_sampler, uv - vec2<f32>(0.0, texel_size.y), 0.0).r + sample_snow_depth_bilinear(uv - vec2<f32>(0.0, texel_size.y));

    let dx = h_right - h_left;
    let dz = h_up - h_down;
    let normal = normalize(vec3<f32>(-dx, 2.0 * texel_size.x * TERRAIN_BOUNDS.x, -dz));

    var out: VertexOutput;
    out.world_position = world_position;
    out.world_normal = normal;
    out.clip_position = uniforms.projection * uniforms.view * vec4<f32>(world_position, 1.0);
    out.dissolve_amount = in.instance_dissolve;
    out.uv = uv;

    return out;
}
