import package::shared::{Spotlight, MAX_SPOTLIGHTS, calculate_spotlight_data};

struct AccumulationUniforms {
    terrain_min_xz: vec2<f32>,
    terrain_max_xz: vec2<f32>,
    decay_rate: f32,
    delta_time: f32,
    spotlight_count: u32,
    _padding: u32,
    light_view_projection: mat4x4<f32>,
    shadow_bias: f32,
    terrain_height_scale: f32,
    _padding3: f32,
    _padding4: f32,
    spotlights: array<Spotlight, 4>,
}

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

@group(0) @binding(1)
var light_sampler: sampler;

@group(0) @binding(2)
var<uniform> uniforms: AccumulationUniforms;

@group(0) @binding(3)
var heightmap: texture_2d<f32>;

@group(0) @binding(4)
var heightmap_sampler: sampler;

@group(0) @binding(5)
var shadow_map: texture_depth_2d;

@group(0) @binding(6)
var shadow_sampler: sampler_comparison;

@group(0) @binding(7)
var snow_depth: texture_2d<f32>;

@group(0) @binding(8)
var snow_depth_sampler: sampler;

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
}

@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
    var out: VertexOutput;

    let x = f32((vertex_index << 1u) & 2u);
    let y = f32(vertex_index & 2u);
    out.position = vec4<f32>(x * 2.0 - 1.0, y * 2.0 - 1.0, 0.0, 1.0);
    out.uv = vec2<f32>(x, 1.0 - y);

    return out;
}

fn sample_shadow_map(light_space_pos: vec4<f32>) -> f32 {
    let proj_coords = light_space_pos.xyz / light_space_pos.w;
    let ndc_coords = proj_coords * vec3<f32>(0.5, -0.5, 1.0) + vec3<f32>(0.5, 0.5, 0.0);

    if ndc_coords.x < 0.0 || ndc_coords.x > 1.0 ||
       ndc_coords.y < 0.0 || ndc_coords.y > 1.0 ||
       ndc_coords.z < 0.0 || ndc_coords.z > 1.0 {
        return 1.0;
    }

    let depth = ndc_coords.z - uniforms.shadow_bias;
    let shadow = textureSampleCompare(shadow_map, shadow_sampler, ndc_coords.xy, depth);

    return shadow;
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let prev_light = textureSample(previous_light, light_sampler, in.uv).r;

    let world_xz = mix(uniforms.terrain_min_xz, uniforms.terrain_max_xz, in.uv);
    let terrain_height = textureSampleLevel(heightmap, heightmap_sampler, in.uv, 0.0).r * uniforms.terrain_height_scale;
    let depth = textureSampleLevel(snow_depth, snow_depth_sampler, in.uv, 0.0).r;
    let snow_surface_height = terrain_height + depth;
    let snow_surface_pos = vec3<f32>(world_xz.x, snow_surface_height, world_xz.y);

    let light_space_position = uniforms.light_view_projection * vec4<f32>(snow_surface_pos, 1.0);
    let shadow = sample_shadow_map(light_space_position);

    var current_light = 0.0;
    if shadow > 0.0 {
        let tile_scale = 2.0;
        let surface_normal = vec3<f32>(0.0, 1.0, 0.0);

        for (var i = 0u; i < uniforms.spotlight_count; i++) {
            let spotlight = uniforms.spotlights[i];
            let data = calculate_spotlight_data(snow_surface_pos, surface_normal, spotlight, tile_scale, shadow);
            let light = f32(data.is_lit);
            current_light = max(current_light, light);
        }
    }

    var accumulated: f32;
    if current_light > 0.01 {
        accumulated = current_light;
    } else {
        let decay_factor = exp(-uniforms.decay_rate * uniforms.delta_time * 60.0);
        accumulated = prev_light * decay_factor;

        if accumulated < 0.01 {
            accumulated = 0.0;
        }
    }

    return vec4<f32>(accumulated, 0.0, 0.0, 1.0);
}