struct UiRectUniforms {
    projection: mat4x4<f32>,
}

struct RectInstance {
    color: vec4<f32>,
    rect: vec4<f32>,
    border_color: vec4<f32>,
    corner_radius: f32,
    border_width: f32,
    depth: f32,
    _padding: f32,
}

@group(0) @binding(0) var<uniform> uniforms: UiRectUniforms;
@group(1) @binding(0) var<uniform> instance: RectInstance;

struct VertexInput {
    @location(0) position: vec2<f32>,
}

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

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

    let rect_pos = instance.rect.xy;
    let rect_size = instance.rect.zw;

    let local = vertex.position * rect_size;
    let world_pos = rect_pos + local;

    output.position = uniforms.projection * vec4<f32>(world_pos, 0.0, 1.0);
    output.position.z = instance.depth;
    output.local_pos = local;
    output.rect_size = rect_size;

    return output;
}

fn rounded_rect_sdf(pos: vec2<f32>, size: vec2<f32>, radius: f32) -> f32 {
    let half_size = size * 0.5;
    let center_pos = pos - half_size;
    let clamped_radius = min(radius, min(half_size.x, half_size.y));
    let q = abs(center_pos) - half_size + clamped_radius;
    return length(max(q, vec2<f32>(0.0))) + min(max(q.x, q.y), 0.0) - clamped_radius;
}

@fragment
fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {
    let dist = rounded_rect_sdf(input.local_pos, input.rect_size, instance.corner_radius);

    let edge_softness = 1.0;
    var alpha = 1.0 - smoothstep(-edge_softness, 0.0, dist);

    var final_color = instance.color;

    if instance.border_width > 0.0 {
        let inner_dist = dist + instance.border_width;
        let border_alpha = 1.0 - smoothstep(-edge_softness, 0.0, inner_dist);
        let is_border = alpha - border_alpha;
        final_color = mix(final_color, instance.border_color, is_border);
    }

    if alpha < 0.01 {
        discard;
    }

    return vec4<f32>(final_color.rgb, final_color.a * alpha);
}
