const RECT: u32 = 0u;
const TEXTURE: u32 = 1u;

@group(0) @binding(0)
var<uniform> window: WindowUniform;
@group(1) @binding(RECT)
var<storage> rects: array<Rect>;
@group(1) @binding(TEXTURE)
var<storage> textures: array<TextureInfo>;

struct Rect {
    color: u32,
    radius: f32,
    thickness: f32,
    inner_radius: f32,
}

struct TextureInfo {
    view_idx: u32,
    sampler_idx: u32,
}

struct Mask {
    x: UiSpan,
    y: UiSpan,
}

struct UiSpan {
    start: UiScalar,
    end: UiScalar,
}

struct UiScalar {
    rel: f32,
    abs: f32,
}

struct UiVec2 {
    rel: vec2<f32>,
    abs: vec2<f32>,
}

@group(2) @binding(0)
var views: binding_array<texture_2d<f32>>;
@group(2) @binding(1)
var samplers: binding_array<sampler>;
@group(2) @binding(2)
var<storage> masks: array<Mask>;

struct WindowUniform {
    dim: vec2<f32>,
};

struct InstanceInput {
    @location(0) x_start: vec2<f32>,
    @location(1) x_end: vec2<f32>,
    @location(2) y_start: vec2<f32>,
    @location(3) y_end: vec2<f32>,
    @location(4) binding: u32,
    @location(5) idx: u32,
    @location(6) mask_idx: u32,
}

struct VertexOutput {
    @location(0) top_left: vec2<f32>,
    @location(1) bot_right: vec2<f32>,
    @location(2) uv: vec2<f32>,
    @location(3) binding: u32,
    @location(4) idx: u32,
    @location(5) mask_idx: u32,
    @builtin(position) clip_position: vec4<f32>,
};

struct Region {
    pos: vec2<f32>,
    uv: vec2<f32>,
    top_left: vec2<f32>,
    bot_right: vec2<f32>,
}

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

    let top_left_rel = vec2(in.x_start.x, in.y_start.x);
    let top_left_abs = vec2(in.x_start.y, in.y_start.y);
    let bot_right_rel = vec2(in.x_end.x, in.y_end.x);
    let bot_right_abs = vec2(in.x_end.y, in.y_end.y);

    let top_left = floor(top_left_rel * window.dim) + floor(top_left_abs);
    let bot_right = floor(bot_right_rel * window.dim) + floor(bot_right_abs);
    let size = bot_right - top_left;

    let uv = vec2<f32>(
        f32(vi % 2u),
        f32(vi / 2u)
    );
    let pos = (top_left + uv * size) / window.dim * 2.0 - 1.0;
    out.clip_position = vec4<f32>(pos.x, -pos.y, 0.0, 1.0);
    out.uv = uv;
    out.binding = in.binding;
    out.idx = in.idx;
    out.top_left = top_left;
    out.bot_right = bot_right;
    out.mask_idx = in.mask_idx;

    return out;
}

@fragment
fn fs_main(
    in: VertexOutput
) -> @location(0) vec4<f32> {
    let pos = in.clip_position.xy;
    let region = Region(pos, in.uv, in.top_left, in.bot_right);
    let i = in.idx;
    var color: vec4<f32>;
    switch in.binding {
        case RECT: {
            color = draw_rounded_rect(region, rects[i]);
        }
        case TEXTURE: {
            color = draw_texture(region, textures[i]);
        }
        default: {
            color = vec4(1.0, 0.0, 1.0, 1.0);
        }
    }
    if in.mask_idx != 4294967295u {
        let mask = masks[in.mask_idx];
        let tl = UiVec2(vec2(mask.x.start.rel, mask.y.start.rel), vec2(mask.x.start.abs, mask.y.start.abs));
        let br = UiVec2(vec2(mask.x.end.rel, mask.y.end.rel), vec2(mask.x.end.abs, mask.y.end.abs));

        let top_left = floor(tl.rel * window.dim) + floor(tl.abs);
        let bot_right = floor(br.rel * window.dim) + floor(br.abs);
        if pos.x < top_left.x || pos.x > bot_right.x || pos.y < top_left.y || pos.y > bot_right.y {
            color *= 0.0;
        }
    }
    return color;
}

// TODO: this seems really inefficient (per frag indexing)?
fn draw_texture(region: Region, info: TextureInfo) -> vec4<f32> {
    return textureSample(views[info.view_idx], samplers[info.sampler_idx], region.uv);
}

fn draw_rounded_rect(region: Region, rect: Rect) -> vec4<f32> {
    var color = unpack4x8unorm(rect.color);

    let edge = 0.5;

    let size = region.bot_right - region.top_left;
    let corner = size / 2.0;
    let center = region.top_left + corner;

    let dist = distance_from_rect(region.pos, center, corner, rect.radius);
    color.a *= 1.0 - smoothstep(-min(edge, rect.radius), edge, dist);

    if rect.thickness > 0.0 {
        let dist2 = distance_from_rect(region.pos, center, corner - rect.thickness, rect.inner_radius);
        color.a *= smoothstep(-min(edge, rect.inner_radius), edge, dist2);
    }

    return color;
}

fn distance_from_rect(pixel_pos: vec2<f32>, rect_center: vec2<f32>, rect_corner: vec2<f32>, radius: f32) -> f32 {
    // vec from center to pixel
    let p = pixel_pos - rect_center;
    // vec from inner rect corner to pixel
    let q = abs(p) - (rect_corner - radius);
    return length(max(q, vec2(0.0))) - radius;
}