@group(0) @binding(0) var<uniform> uniforms: GrassUniforms;
@group(0) @binding(1) var<storage, read> types: array<GrassType>;
@group(0) @binding(2) var<storage, read> blades: array<GrassBlade>;
@group(0) @binding(3) var<storage, read> far_tiles: array<GrassTileEntry>;

struct BladeVarying {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) world_position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) along_blade: f32,
    @location(3) blade_data: vec3<f32>,
};

fn blade_vertex(instance_index: u32, vertex_index: u32, lod: u32) -> BladeVarying {
    let blade = blades[min(instance_index, GRASS_MAX_BLADES - 1u)];
    let type_index = u32(blade.facing_type.w);
    let blade_type = types[type_index];
    var height = blade.position_height.w;
    let folded = lod == 0u && height < GRASS_FOLD_HEIGHT;

    var t_raw: f32;
    var side: f32;
    var sub_blade = 0u;
    if lod == 1u {
        var t_table = array<f32, 7>(0.0, 0.0, 0.45, 0.45, 0.8, 0.8, 1.0);
        t_raw = t_table[vertex_index];
        side = select(f32(vertex_index & 1u) * 2.0 - 1.0, 0.0, vertex_index == 6u);
    } else if folded {
        var t_table = array<f32, 15>(
            0.0, 0.0, 0.34, 0.34, 0.68, 0.68, 1.0, 1.0,
            0.0, 0.0, 0.5, 0.5, 0.8, 0.8, 1.0,
        );
        t_raw = t_table[vertex_index];
        sub_blade = select(0u, 1u, vertex_index >= 8u);
        side = select(f32(vertex_index & 1u) * 2.0 - 1.0, 0.0, vertex_index == 14u);
    } else {
        var t_table = array<f32, 15>(
            0.0, 0.0, 0.142857, 0.142857, 0.285714, 0.285714, 0.428571, 0.428571,
            0.428571, 0.428571, 0.714286, 0.714286, 0.857143, 0.857143, 1.0,
        );
        t_raw = t_table[vertex_index];
        side = select(f32(vertex_index & 1u) * 2.0 - 1.0, 0.0, vertex_index == 14u);
    }
    let t = pow(t_raw, blade_type.misc2.y);

    let tile_origin = vec2<f32>(uniforms.camera_tile.xy) * GRASS_TILE_SIZE;
    var base = vec3<f32>(
        tile_origin.x + blade.position_height.x,
        blade.position_height.y - 0.02,
        tile_origin.y + blade.position_height.z,
    );
    var facing = blade.facing_type.xy;
    if folded {
        let splay = select(-0.55, 0.55, sub_blade == 1u);
        let cos_splay = cos(splay);
        let sin_splay = sin(splay);
        facing = vec2<f32>(
            facing.x * cos_splay - facing.y * sin_splay,
            facing.x * sin_splay + facing.y * cos_splay,
        );
        let perpendicular = vec2<f32>(-blade.facing_type.y, blade.facing_type.x);
        let shift = select(-1.0, 1.0, sub_blade == 1u) * blade.shape.z;
        base.x += perpendicular.x * shift;
        base.z += perpendicular.y * shift;
        height *= 0.85;
    }

    let tilt = blade.shape.x;
    let bend = blade.shape.y;
    let time = uniforms.params.w;
    let phase = blade.facing_type.z * GRASS_TAU;
    let sway = blade.shape.w * blade_type.misc.y * (0.6 + 0.4 * sin(time * 1.7 + phase));
    let bob = sin(time * 2.3 + phase + t * 1.4) * 0.07 * blade.shape.w * blade_type.misc.y;

    let lean = vec3<f32>(facing.x, 0.0, facing.y);
    let up = vec3<f32>(0.0, 1.0, 0.0);
    let segment = height / 3.0;
    let angle_1 = tilt * 0.35;
    let angle_2 = tilt * 0.75 + bend * 0.6;
    let angle_3 = min(tilt + bend * 1.4, 1.9);
    let p0 = base;
    let p1 = p0 + (lean * sin(angle_1) + up * cos(angle_1)) * segment;
    var p2 = p1 + (lean * sin(angle_2) + up * cos(angle_2)) * segment;
    var p3 = p2 + (lean * sin(angle_3) + up * cos(angle_3)) * segment;
    let wind_direction = uniforms.wind.xy;
    let wind_push = vec3<f32>(wind_direction.x, 0.0, wind_direction.y) * (sway * height * 0.7 + bob);
    p2 += wind_push * 0.5;
    p3 += wind_push;

    var position = grass_bezier(p0, p1, p2, p3, t);
    if lod == 0u && !folded {
        let base_distance = distance(base, uniforms.camera_position.xyz);
        let shape_blend = smoothstep(uniforms.radii.x * 0.7, uniforms.radii.x, base_distance);
        if shape_blend > 0.001 {
            var knot_low = 0.0;
            var knot_high = 0.45;
            if t_raw > 0.8 {
                knot_low = 0.8;
                knot_high = 1.0;
            } else if t_raw > 0.45 {
                knot_low = 0.45;
                knot_high = 0.8;
            }
            let knot_fraction = (t_raw - knot_low) / (knot_high - knot_low);
            let exponent = blade_type.misc2.y;
            let coarse_a = grass_bezier(p0, p1, p2, p3, pow(knot_low, exponent));
            let coarse_b = grass_bezier(p0, p1, p2, p3, pow(knot_high, exponent));
            position = mix(position, mix(coarse_a, coarse_b, knot_fraction), shape_blend);
        }
    }
    let tangent = normalize(grass_bezier_derivative(p0, p1, p2, p3, max(t, 0.01)));

    let side_direction = normalize(vec3<f32>(-facing.y, 0.0, facing.x));
    let width = blade.shape.z * (1.0 - 0.8 * t) * select(1.0, 0.7, folded);
    position += side_direction * side * width * 0.5;

    var normal = normalize(cross(side_direction, tangent));
    normal = normalize(normal + side_direction * side * 0.42);

    let view_direction = normalize(uniforms.camera_position.xyz - position);
    let orthogonality = 1.0 - abs(dot(view_direction, normal));
    let thickening = orthogonality * orthogonality * (1.0 - 0.4 * orthogonality);
    let view_side = sign(dot(side_direction, uniforms.camera_right.xyz) + 1e-4);
    position += uniforms.camera_right.xyz * side * view_side * width * 0.9 * thickening;

    var output: BladeVarying;
    output.clip_position = uniforms.view_projection * vec4<f32>(position, 1.0);
    output.world_position = position;
    output.normal = normal;
    output.along_blade = t;
    output.blade_data = vec3<f32>(f32(type_index), blade.extra.x, blade.facing_type.z);
    return output;
}

@vertex
fn vs_high(
    @builtin(vertex_index) vertex_index: u32,
    @builtin(instance_index) instance_index: u32,
) -> BladeVarying {
    return blade_vertex(instance_index, vertex_index, 0u);
}

@vertex
fn vs_low(
    @builtin(vertex_index) vertex_index: u32,
    @builtin(instance_index) instance_index: u32,
) -> BladeVarying {
    return blade_vertex(instance_index, vertex_index, 1u);
}

@fragment
fn fs_blade(input: BladeVarying) -> @location(0) vec4<f32> {
    let blade_type = types[u32(input.blade_data.x)];
    let t = input.along_blade;
    var albedo = mix(blade_type.color_base.rgb, blade_type.color_tip.rgb, pow(t, 1.3));
    albedo *= 0.9 + 0.2 * input.blade_data.y;

    let view_direction = normalize(uniforms.camera_position.xyz - input.world_position);
    var normal = normalize(input.normal);
    if dot(normal, view_direction) < 0.0 {
        normal = -normal;
    }
    let view_distance = distance(uniforms.camera_position.xyz, input.world_position);
    let flatten = smoothstep(18.0, 80.0, view_distance);
    let shading_normal = normalize(mix(normal, vec3<f32>(0.0, 1.0, 0.0), flatten * 0.75));

    let light_direction = uniforms.sun_direction.xyz;
    let diffuse_term = max(dot(shading_normal, light_direction), 0.0);
    let half_vector = normalize(light_direction + view_direction);
    let gloss = blade_type.misc.x * (1.0 - flatten * 0.85);
    let specular = pow(max(dot(shading_normal, half_vector), 0.0), mix(12.0, 48.0, blade_type.misc.x)) * gloss * 0.1;
    let occlusion = mix(0.22, 1.0, t * t);

    let color = albedo * (uniforms.ambient.rgb + uniforms.sun_color.rgb * diffuse_term / GRASS_PI) * occlusion
        + uniforms.sun_color.rgb * specular * occlusion;
    return vec4<f32>(color, 1.0);
}
