struct TerrainRenderUniforms {
    view_projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    sun_direction: vec4<f32>,
    sun_color: vec4<f32>,
    ambient: vec4<f32>,
    params: vec4<f32>,
    splat: vec4<f32>,
    rock_color: vec4<f32>,
    snow_color: vec4<f32>,
    wind_time: vec4<f32>,
    frustum_planes: array<vec4<f32>, 6>,
    level_regions: array<vec4<f32>, 12>,
    cascade_view_projections: array<mat4x4<f32>, 4>,
    cascade_splits: vec4<f32>,
    cascade_atlas_offsets: array<vec4<f32>, 4>,
    cascade_atlas_scale: vec4<f32>,
};

struct TerrainInstance {
    data: vec4<f32>,
};

@group(0) @binding(0) var<uniform> uniforms: TerrainRenderUniforms;
@group(0) @binding(1) var<storage, read> instances: array<TerrainInstance>;
@group(0) @binding(2) var cache: texture_2d_array<f32>;
@group(0) @binding(3) var<storage, read> types: array<GrassType>;
@group(0) @binding(4) var ambient_cache: texture_2d_array<f32>;
@group(0) @binding(5) var shadow_texture: texture_depth_2d;
@group(0) @binding(6) var shadow_sampler: sampler;

struct TerrainVarying {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) world_position: vec3<f32>,
    @location(1) @interpolate(flat) level: u32,
};

const TERRAIN_GRID_CELLS: f32 = 32.0;

fn level_texel_size(level: u32) -> f32 {
    return uniforms.params.x * f32(1u << level);
}

fn level_height(level: u32, world_xz: vec2<f32>) -> f32 {
    let clamped = min(level, u32(uniforms.params.y) - 1u);
    return terrain_cache_sample(cache, clamped, level_texel_size(clamped), world_xz);
}

@vertex
fn vs_terrain(
    @builtin(vertex_index) vertex_index: u32,
    @builtin(instance_index) instance_index: u32,
) -> TerrainVarying {
    let instance = instances[instance_index].data;
    let level = u32(instance.w);
    let texel_size = level_texel_size(level);
    var grid = vec2<f32>(f32(vertex_index % 33u), f32(vertex_index / 33u));

    let region = uniforms.level_regions[level];
    let world_pre = instance.xy + grid * texel_size;
    let center_distance = max(
        abs(world_pre.x - region.x),
        abs(world_pre.y - region.y),
    );
    let morph = clamp((center_distance / region.z - 0.7) / 0.25, 0.0, 1.0);

    grid -= fract(grid * 0.5) * 2.0 * morph;
    let world_xz = instance.xy + grid * texel_size;

    let height_fine = level_height(level, world_xz);
    let height_coarse = level_height(level + 1u, world_xz);
    let height = clamp(
        mix(height_fine, height_coarse, morph),
        uniforms.params.z,
        uniforms.params.w,
    );

    let world = vec3<f32>(world_xz.x, height, world_xz.y);
    var output: TerrainVarying;
    output.clip_position = uniforms.view_projection * vec4<f32>(world, 1.0);
    output.world_position = world;
    output.level = level;
    return output;
}

const TERRAIN_PICK_ID: u32 = 0x7F7FFFF0u;

struct TerrainFragmentOutput {
    @location(0) color: vec4<f32>,
    @location(1) entity_id: f32,
};

fn terrain_shadow(world_pos: vec3<f32>, normal: vec3<f32>, view_distance: f32) -> f32 {
    if uniforms.cascade_atlas_scale.z < 0.5 {
        return 1.0;
    }
    var cascade = 3;
    for (var cascade_index = 0; cascade_index < 4; cascade_index++) {
        if view_distance < uniforms.cascade_splits[cascade_index] {
            cascade = cascade_index;
            break;
        }
    }
    let texel_world = uniforms.cascade_atlas_offsets[cascade].z;
    let offset_position = world_pos
        + normal * texel_world * 1.5
        + uniforms.sun_direction.xyz * uniforms.cascade_atlas_scale.w;
    let light_space = uniforms.cascade_view_projections[cascade] * vec4<f32>(offset_position, 1.0);
    let ndc = light_space.xyz / light_space.w;
    if ndc.z < 0.0 || ndc.z > 1.0 {
        return 1.0;
    }
    var shadow_uv = ndc.xy * vec2<f32>(0.5, -0.5) + vec2<f32>(0.5, 0.5);
    let atlas_offset = uniforms.cascade_atlas_offsets[cascade].xy;
    let atlas_scale = uniforms.cascade_atlas_scale.xy;
    shadow_uv = shadow_uv * atlas_scale + atlas_offset;
    let slot_min = atlas_offset;
    let slot_max = atlas_offset + atlas_scale;
    if shadow_uv.x < slot_min.x || shadow_uv.x > slot_max.x
        || shadow_uv.y < slot_min.y || shadow_uv.y > slot_max.y {
        return 1.0;
    }
    let atlas_size = f32(textureDimensions(shadow_texture).x);
    let texel = 1.0 / atlas_size;
    var visibility = 0.0;
    for (var offset_y = -1; offset_y <= 1; offset_y++) {
        for (var offset_x = -1; offset_x <= 1; offset_x++) {
            let sample_uv = clamp(
                shadow_uv + vec2<f32>(f32(offset_x), f32(offset_y)) * texel,
                slot_min,
                slot_max,
            );
            let sampled_depth = textureSampleLevel(shadow_texture, shadow_sampler, sample_uv, 0);
            visibility += select(0.0, 1.0, ndc.z >= sampled_depth);
        }
    }
    return visibility / 9.0;
}

@fragment
fn fs_terrain(input: TerrainVarying) -> TerrainFragmentOutput {
    let level = input.level;
    let texel_size = level_texel_size(level);
    let world_xz = input.world_position.xz;
    let offset = texel_size;
    let height_x0 = level_height(level, world_xz - vec2<f32>(offset, 0.0));
    let height_x1 = level_height(level, world_xz + vec2<f32>(offset, 0.0));
    let height_z0 = level_height(level, world_xz - vec2<f32>(0.0, offset));
    let height_z1 = level_height(level, world_xz + vec2<f32>(0.0, offset));
    let normal = normalize(vec3<f32>(
        height_x0 - height_x1,
        2.0 * offset,
        height_z0 - height_z1,
    ));

    let slope = 1.0 - normal.y;
    let height = input.world_position.y;
    let detail_fade = 1.0 - smoothstep(
        300.0,
        1500.0,
        distance(uniforms.camera_position.xz, input.world_position.xz),
    );
    let splat_noise = terrain_fbm(world_xz * 0.05, 2u, 977u) * detail_fade;

    let view_distance = distance(uniforms.camera_position.xz, world_xz);
    let type_count = max(u32(uniforms.splat.z), 1u);
    let edge_data = grass_type_edge(world_xz, uniforms.splat.w);
    let type_a = types[min(u32(edge_data.x * f32(type_count)), type_count - 1u)];
    let type_b = types[min(u32(edge_data.y * f32(type_count)), type_count - 1u)];
    let type_blend = 0.5 * (1.0 - edge_data.z);
    let color_base = mix(type_a.color_base.rgb, type_b.color_base.rgb, type_blend);
    let color_tip = mix(type_a.color_tip.rgb, type_b.color_tip.rgb, type_blend);
    var albedo = mix(color_base, color_tip, 0.18) * 0.9;
    let clump = grass_voronoi(world_xz, uniforms.rock_color.w * 2.0, 23u);
    albedo *= 1.0 + (0.2 * clump.cell_value - 0.1);

    let rock_amount = smoothstep(
        uniforms.splat.x - 0.12,
        uniforms.splat.x + 0.12,
        slope + splat_noise * 0.08,
    );
    albedo = mix(albedo, uniforms.rock_color.rgb * (0.8 + 0.4 * splat_noise), rock_amount);

    let snow_amount = smoothstep(
        uniforms.splat.y - 14.0,
        uniforms.splat.y + 14.0,
        height + splat_noise * 18.0,
    ) * (1.0 - smoothstep(0.55, 0.8, slope));
    albedo = mix(albedo, uniforms.snow_color.rgb, snow_amount);

    let grass_mask = (1.0 - rock_amount) * (1.0 - snow_amount);
    let shimmer = grass_wind_sample(
        world_xz,
        uniforms.wind_time.w,
        vec4<f32>(uniforms.wind_time.xyz, 0.5),
    );
    albedo *= 1.0 + (0.10 * shimmer - 0.05) * grass_mask;

    let wind_direction = uniforms.wind_time.xy;
    let wind_perpendicular = vec2<f32>(-wind_direction.y, wind_direction.x);
    let streak_coords = vec2<f32>(
        dot(world_xz, wind_direction) * 0.35,
        dot(world_xz, wind_perpendicular) * 2.2,
    );
    let streaks = terrain_perlin(streak_coords, 311u);
    let streak_fade = 1.0 - smoothstep(900.0, 2200.0, view_distance);
    albedo *= 1.0 + 0.06 * streaks * grass_mask * streak_fade;

    let ambient_occlusion = terrain_cache_sample(
        ambient_cache,
        min(level, u32(uniforms.params.y) - 1u),
        texel_size,
        world_xz,
    );
    let shadow = terrain_shadow(input.world_position, normal, view_distance);

    let light_direction = uniforms.sun_direction.xyz;
    let diffuse_term = max(dot(normal, light_direction), 0.0);
    let color = albedo
        * (uniforms.ambient.rgb * ambient_occlusion
            + uniforms.sun_color.rgb * diffuse_term * shadow * (0.55 + 0.45 * ambient_occlusion)
                / 3.14159265)
        * 0.8;
    var output: TerrainFragmentOutput;
    output.color = vec4<f32>(color, 1.0);
    output.entity_id = bitcast<f32>(TERRAIN_PICK_ID);
    return output;
}
