#import "modules/std.sx";
#import "modules/build.sx";
#import "modules/ffi/opengl.sx";
#import "modules/math";
#import "modules/gpu/types.sx";
#import "modules/gpu/api.sx";
#import "modules/ui/types.sx";
#import "modules/ui/render.sx";
#import "modules/ui/glyph_cache.sx";
#import "modules/ui/font.sx";

// Vertex: pos(2) + uv(2) + color(4) + params(4) = 12 floats
UI_VERTEX_FLOATS :s64: 12;
UI_VERTEX_BYTES  :s64: 48;
MAX_UI_VERTICES  :s64: 16384;

UIRenderer :: struct {
    // GL-side handles. Used when `gpu == null` (every non-iOS target today).
    vao: u32;
    vbo: u32;
    shader: u32;
    proj_loc: s32;
    tex_loc: s32;

    // CPU-side vertex scratch buffer — same for both backends.
    vertices: [*]f32;
    vertex_count: s64;
    screen_width: f32;
    screen_height: f32;
    dpi_scale: f32;
    white_texture: u32;     // GL name OR TextureHandle (both are u32-shaped)
    current_texture: u32;
    draw_calls: s64;

    // GPU protocol backend. When set, the renderer routes shader / buffer /
    // texture / draw calls through `gpu` instead of raw GL. The chess game
    // sets this on iOS to a boxed `*MetalGPU`.
    gpu: ?GPU = null;
    mtl_shader: ShaderHandle = 0;
    mtl_vbuf:   BufferHandle = 0;
    // Per-frame byte offset into the Metal vertex buffer. Each flush writes
    // to a fresh slice so concurrent in-flight draws don't trample each
    // other's data — Metal's shared-storage buffer is read at GPU execution
    // time, not at draw-call submission, so re-using offset 0 across flushes
    // would let the last writer win and earlier batches would render as
    // whatever was uploaded last. Reset to 0 in `begin()`.
    mtl_buf_offset: s64 = 0;
    mtl_buf_capacity: s64 = 0;

    init :: (self: *UIRenderer) {
        // Allocate vertex scratch (CPU side) — same for both backends.
        buf_size := MAX_UI_VERTICES * UI_VERTEX_BYTES;
        self.vertices = xx context.allocator.alloc(buf_size);
        memset(self.vertices, 0, buf_size);
        self.vertex_count = 0;
        self.dpi_scale = 1.0;

        if self.gpu != null {
            // ── Metal backend (via GPU protocol) ───────────────────────
            // Oversize the GPU buffer enough to hold many sub-batches per
            // frame without wrapping. With per-flush offset advance, each
            // draw reads from its own slice and can outlive earlier in-
            // flight draws without corruption.
            metal_buf_size := buf_size * 4;
            // Backend-specific shader sources: Metal takes a combined MSL
            // library with `vmain` + `fmain` entry points; GLES3 takes a
            // separate vertex / fragment pair. Caller selects via OS gate.
            inline if OS == .android {
                self.mtl_shader = self.gpu.create_shader(UI_VERT_SRC_ES, UI_FRAG_SRC_ES);
            }
            inline if OS == .ios {
                self.mtl_shader = self.gpu.create_shader(UI_MSL_SRC, "");
            }
            self.mtl_vbuf   = self.gpu.create_buffer(metal_buf_size);
            self.mtl_buf_capacity = metal_buf_size;
            white_px : [4]u8 = .[255, 255, 255, 255];
            self.white_texture = self.gpu.create_texture(1, 1, .rgba8, xx @white_px[0]);
        } else {
            // ── GL backend ─────────────────────────────────────────────
            // Create shader (ES for WASM/WebGL2 + iOS GLES3, Core for desktop GL 3.3)
            inline if OS == .wasm or OS == .ios {
                self.shader = create_program(UI_VERT_SRC_ES, UI_FRAG_SRC_ES);
            } else {
                self.shader = create_program(UI_VERT_SRC_CORE, UI_FRAG_SRC_CORE);
            }
            self.proj_loc = glGetUniformLocation(self.shader, "uProj");
            self.tex_loc = glGetUniformLocation(self.shader, "uTex");

            // Create VAO/VBO
            glGenVertexArrays(1, @self.vao);
            glGenBuffers(1, @self.vbo);
            glBindVertexArray(self.vao);
            glBindBuffer(GL_ARRAY_BUFFER, self.vbo);
            glBufferData(GL_ARRAY_BUFFER, xx buf_size, null, GL_DYNAMIC_DRAW);

            // pos (2 floats)
            glVertexAttribPointer(0, 2, GL_FLOAT, 0, xx UI_VERTEX_BYTES, xx 0);
            glEnableVertexAttribArray(0);
            // uv (2 floats)
            glVertexAttribPointer(1, 2, GL_FLOAT, 0, xx UI_VERTEX_BYTES, xx 8);
            glEnableVertexAttribArray(1);
            // color (4 floats)
            glVertexAttribPointer(2, 4, GL_FLOAT, 0, xx UI_VERTEX_BYTES, xx 16);
            glEnableVertexAttribArray(2);
            // params: corner_radius, border_width, rect_w, rect_h
            glVertexAttribPointer(3, 4, GL_FLOAT, 0, xx UI_VERTEX_BYTES, xx 32);
            glEnableVertexAttribArray(3);

            glBindVertexArray(0);

            // 1x1 white texture for solid rects
            self.white_texture = create_white_texture();
        }
    }

    begin :: (self: *UIRenderer, width: f32, height: f32, font_texture: u32) {
        self.screen_width = width;
        self.screen_height = height;
        self.vertex_count = 0;
        self.current_texture = font_texture;
        self.draw_calls = 0;

        proj := Mat4.ortho(0.0, width, height, 0.0, -1.0, 1.0);

        if self.gpu != null {
            // Reset the per-frame ring offset; this frame's flushes start at 0.
            self.mtl_buf_offset = 0;
            // Pipeline state + vertex buffer + projection + initial texture.
            // Metal blend mode + scissor-cleared defaults are baked into
            // the pipeline state, so no per-frame glEnable/glDisable.
            self.gpu.set_shader(self.mtl_shader);
            self.gpu.set_vertex_buffer(self.mtl_vbuf);
            self.gpu.set_vertex_constants(1, xx proj.data, 64);
            self.gpu.set_texture(0, font_texture);
        } else {
            // GL: bind everything for the frame.
            glUseProgram(self.shader);
            glUniformMatrix4fv(self.proj_loc, 1, 0, proj.data);
            glUniform1i(self.tex_loc, 0);
            glActiveTexture(GL_TEXTURE0);
            glBindTexture(GL_TEXTURE_2D, font_texture);
            glBindVertexArray(self.vao);
            glBindBuffer(GL_ARRAY_BUFFER, self.vbo);
        }
    }

    bind_texture :: (self: *UIRenderer, tex: u32) {
        if tex != self.current_texture {
            self.flush();
            self.current_texture = tex;
        }
    }

    // Emit a quad (2 triangles = 6 vertices)
    push_quad :: (self: *UIRenderer, frame: Frame, color: Color, radius: f32, border_w: f32) {
        if self.vertex_count + 6 > MAX_UI_VERTICES {
            self.flush();
        }

        x0 := frame.origin.x;
        y0 := frame.origin.y;
        x1 := x0 + frame.size.width;
        y1 := y0 + frame.size.height;

        r := color.rf();
        g := color.gf();
        b := color.bf();
        a := color.af();

        w := frame.size.width;
        h := frame.size.height;

        // 6 vertices for quad: TL, TR, BL, TR, BR, BL
        self.write_vertex(x0, y0, 0.0, 0.0, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y0, 1.0, 0.0, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x0, y1, 0.0, 1.0, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y0, 1.0, 0.0, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y1, 1.0, 1.0, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x0, y1, 0.0, 1.0, r, g, b, a, radius, border_w, w, h);
    }

    // Emit a quad with custom UV coordinates (for sprite sheet sub-textures)
    push_quad_uv :: (self: *UIRenderer, frame: Frame, color: Color, radius: f32, border_w: f32, uv_min: Point, uv_max: Point) {
        if self.vertex_count + 6 > MAX_UI_VERTICES {
            self.flush();
        }

        x0 := frame.origin.x;
        y0 := frame.origin.y;
        x1 := x0 + frame.size.width;
        y1 := y0 + frame.size.height;

        r := color.rf();
        g := color.gf();
        b := color.bf();
        a := color.af();

        w := frame.size.width;
        h := frame.size.height;

        u_min := uv_min.x;
        v_min := uv_min.y;
        u_max := uv_max.x;
        v_max := uv_max.y;

        self.write_vertex(x0, y0, u_min, v_min, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y0, u_max, v_min, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x0, y1, u_min, v_max, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y0, u_max, v_min, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x1, y1, u_max, v_max, r, g, b, a, radius, border_w, w, h);
        self.write_vertex(x0, y1, u_min, v_max, r, g, b, a, radius, border_w, w, h);
    }

    write_vertex :: (self: *UIRenderer, x: f32, y: f32, u: f32, v: f32, r: f32, g: f32, b: f32, a: f32, cr: f32, bw: f32, rw: f32, rh: f32) {
        off := self.vertex_count * UI_VERTEX_FLOATS;
        self.vertices[off + 0] = x;
        self.vertices[off + 1] = y;
        self.vertices[off + 2] = u;
        self.vertices[off + 3] = v;
        self.vertices[off + 4] = r;
        self.vertices[off + 5] = g;
        self.vertices[off + 6] = b;
        self.vertices[off + 7] = a;
        self.vertices[off + 8] = cr;
        self.vertices[off + 9] = bw;
        self.vertices[off + 10] = rw;
        self.vertices[off + 11] = rh;
        self.vertex_count += 1;
    }

    // Walk the render tree and emit quads
    process :: (self: *UIRenderer, tree: *RenderTree) {
        i := 0;
        while i < tree.nodes.len {
            node := tree.nodes.items[i];
            if node.type == {
                case .rect: {
                    self.push_quad(node.frame, node.fill_color, 0.0, 0.0);
                }
                case .rounded_rect: {
                    self.push_quad(node.frame, node.fill_color, node.corner_radius, node.stroke_width);
                }
                case .text: {
                    if g_font != null {
                        self.render_text(node);
                    }
                }
                case .image: {
                    self.bind_texture(node.texture_id);
                    neg2 : f32 = 0.0 - 2.0;
                    self.push_quad_uv(node.frame, COLOR_WHITE, neg2, 0.0, node.uv_min, node.uv_max);
                    // Re-bind font atlas after image
                    font := g_font;
                    if font != null {
                        self.bind_texture(font.texture_id);
                    }
                }
                case .clip_push: {
                    self.flush();
                    dpi := self.dpi_scale;
                    if self.gpu != null {
                        // Metal: pixel coords, top-left origin (no Y flip).
                        self.gpu.set_scissor(
                            xx (node.frame.origin.x * dpi),
                            xx (node.frame.origin.y * dpi),
                            xx (node.frame.size.width * dpi),
                            xx (node.frame.size.height * dpi),
                        );
                    } else {
                        // GL: pixel coords, bottom-left origin — flip Y.
                        glEnable(GL_SCISSOR_TEST);
                        glScissor(
                            xx (node.frame.origin.x * dpi),
                            xx ((self.screen_height - node.frame.origin.y - node.frame.size.height) * dpi),
                            xx (node.frame.size.width * dpi),
                            xx (node.frame.size.height * dpi)
                        );
                    }
                }
                case .clip_pop: {
                    self.flush();
                    if self.gpu != null {
                        self.gpu.disable_scissor();
                    } else {
                        glDisable(GL_SCISSOR_TEST);
                    }
                }
                case .opacity_push: {}
                case .opacity_pop: {}
            }
            i += 1;
        }
    }

    flush :: (self: *UIRenderer) {
        if self.vertex_count == 0 { return; }

        upload_size : s64 = self.vertex_count * UI_VERTEX_BYTES;

        if self.gpu != null {
            // Mirror the GL path: bind current texture before drawing.
            // current_texture may have changed since the last flush.
            self.gpu.set_texture(0, self.current_texture);

            // Write this batch to a fresh slice of the GPU buffer and draw
            // it from there. Re-using offset 0 would race against earlier
            // still-in-flight draws (see `mtl_buf_offset` comment in the
            // struct).
            if self.mtl_buf_offset + upload_size > self.mtl_buf_capacity {
                // Frame overflowed the GPU buffer; wrap to 0. Previous in-
                // flight batches from this frame will likely render wrong,
                // but the alternative (skipping the draw) would render
                // even less. Practical UIs should never hit this.
                self.mtl_buf_offset = 0;
            }
            byte_off := self.mtl_buf_offset;
            self.gpu.update_buffer_at(self.mtl_vbuf, xx self.vertices, upload_size, byte_off);
            vertex_off : s32 = xx (byte_off / UI_VERTEX_BYTES);
            self.gpu.draw_triangles(vertex_off, xx self.vertex_count);
            // Each batch starts on a vertex boundary so `vertex_off =
            // byte_off / UI_VERTEX_BYTES` lands on a whole vertex (otherwise
            // the shader reads partway through the previous vertex and the
            // text quads end up reading garbled UVs/colors). upload_size is
            // already vertex_count × UI_VERTEX_BYTES so the running offset
            // stays vertex-aligned without an extra `% UI_VERTEX_BYTES` pad.
            self.mtl_buf_offset = byte_off + upload_size;
        } else {
            // Only re-bind the current texture (program, projection, VAO
            // already bound in begin()). glBufferData orphans the old buffer
            // to avoid GPU sync stalls.
            glBindTexture(GL_TEXTURE_2D, self.current_texture);
            glBufferData(GL_ARRAY_BUFFER, xx upload_size, self.vertices, GL_DYNAMIC_DRAW);
            glDrawArrays(GL_TRIANGLES, 0, xx self.vertex_count);
        }

        self.vertex_count = 0;
        self.draw_calls += 1;
    }

    render_text :: (self: *UIRenderer, node: RenderNode) {
        font := g_font;
        if font == null { return; }

        // Shape text into positioned glyphs. This may rasterize new glyphs
        // AND grow the atlas, which creates a new GPU texture under
        // `font.texture_id`. Re-bind the font atlas after shaping so the
        // upcoming text quads sample the texture that actually holds the
        // glyphs we just rasterized.
        font.shape_text(node.text, node.font_size);

        // Flush any new glyphs to the atlas texture
        font.flush();

        // If the atlas grew (or otherwise changed handle), switch the
        // current texture so the next flush samples the right one.
        self.bind_texture(font.texture_id);

        r := node.text_color.rf();
        g := node.text_color.gf();
        b := node.text_color.bf();
        a := node.text_color.af();

        ascent := font.get_ascent(node.font_size);
        raster_size := node.font_size * font.dpi_scale;
        inv_dpi := font.inv_dpi;

        i : s64 = 0;
        while i < font.shaped_buf.len {
            shaped := font.shaped_buf.items[i];
            cached := font.get_or_rasterize(shaped.glyph_index, raster_size);

            if cached != null {
                if cached.width > 0.0 {
                    // Scale physical pixel dimensions back to logical units
                    gx0 := node.frame.origin.x + shaped.x + cached.offset_x * inv_dpi;
                    gy0 := node.frame.origin.y + ascent + shaped.y + cached.offset_y * inv_dpi;
                    gx1 := gx0 + cached.width * inv_dpi;
                    gy1 := gy0 + cached.height * inv_dpi;

                    u_min := cached.uv_x;
                    v_min := cached.uv_y;
                    u_max := cached.uv_x + cached.uv_w;
                    v_max := cached.uv_y + cached.uv_h;


                    if self.vertex_count + 6 > MAX_UI_VERTICES {
                        self.flush();
                    }

                    // corner_radius = -1.0 signals "text mode" to the fragment shader
                    neg1 : f32 = 0.0 - 1.0;
                    self.write_vertex(gx0, gy0, u_min, v_min, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                    self.write_vertex(gx1, gy0, u_max, v_min, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                    self.write_vertex(gx0, gy1, u_min, v_max, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                    self.write_vertex(gx1, gy0, u_max, v_min, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                    self.write_vertex(gx1, gy1, u_max, v_max, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                    self.write_vertex(gx0, gy1, u_min, v_max, r, g, b, a, neg1, 0.0, 0.0, 0.0);
                }
            }
            i += 1;
        }

        // Flush any glyphs rasterized during this text draw
        font.flush();
    }
}

create_white_texture :: () -> u32 {
    tex : u32 = 0;
    glGenTextures(1, @tex);
    glBindTexture(GL_TEXTURE_2D, tex);
    pixel : [4]u8 = .[255, 255, 255, 255];
    glTexImage2D(GL_TEXTURE_2D, 0, xx GL_RGBA, 1, 1, 0, GL_RGBA, GL_UNSIGNED_BYTE, @pixel);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, xx GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, xx GL_NEAREST);
    tex
}

// --- UI Shaders ---

// --- Desktop (Core Profile 3.3) shaders ---

UI_VERT_SRC_CORE :: #string GLSL
#version 330 core
layout(location = 0) in vec2 aPos;
layout(location = 1) in vec2 aUV;
layout(location = 2) in vec4 aColor;
layout(location = 3) in vec4 aParams;

uniform mat4 uProj;

out vec2 vUV;
out vec4 vColor;
out vec4 vParams;

void main() {
    gl_Position = uProj * vec4(aPos, 0.0, 1.0);
    vUV = aUV;
    vColor = aColor;
    vParams = aParams;
}
GLSL;

UI_FRAG_SRC_CORE :: #string GLSL
#version 330 core
in vec2 vUV;
in vec4 vColor;
in vec4 vParams;

uniform sampler2D uTex;

out vec4 FragColor;

float roundedBoxSDF(vec2 center, vec2 half_size, float radius) {
    vec2 q = abs(center) - half_size + vec2(radius);
    return length(max(q, vec2(0.0))) + min(max(q.x, q.y), 0.0) - radius;
}

void main() {
    float mode = vParams.x;
    float border = vParams.y;
    vec2 rectSize = vParams.zw;

    if (mode < -1.5) {
        // Image mode (mode == -2.0): sample texture
        FragColor = texture(uTex, vUV) * vColor;
    } else if (mode < 0.0) {
        // Text mode (mode == -1.0): sample glyph atlas .r as alpha
        float alpha = texture(uTex, vUV).r;
        float ew = fwidth(alpha) * 0.7;
        alpha = smoothstep(0.5 - ew, 0.5 + ew, alpha);
        FragColor = vec4(vColor.rgb, vColor.a * pow(alpha, 0.9));
    } else if (mode > 0.0 || border > 0.0) {
        // Rounded rect: SDF alpha, vertex color only (no texture sample)
        vec2 half_size = rectSize * 0.5;
        vec2 center = (vUV - vec2(0.5)) * rectSize;
        float dist = roundedBoxSDF(center, half_size, mode);
        float aa = fwidth(dist);
        float alpha = 1.0 - smoothstep(-aa, aa, dist);

        if (border > 0.0) {
            float inner = roundedBoxSDF(center, half_size - vec2(border), max(mode - border, 0.0));
            float border_alpha = smoothstep(-aa, aa, inner);
            alpha = alpha * max(border_alpha, 0.0);
        }

        FragColor = vec4(vColor.rgb, vColor.a * alpha);
    } else {
        // Plain rect: vertex color only (no texture sample)
        FragColor = vColor;
    }
}
GLSL;

// --- WASM (ES 3.0 / WebGL2) shaders ---

UI_VERT_SRC_ES :: #string GLSL
#version 300 es
precision mediump float;
layout(location = 0) in vec2 aPos;
layout(location = 1) in vec2 aUV;
layout(location = 2) in vec4 aColor;
layout(location = 3) in vec4 aParams;

uniform mat4 uProj;

out vec2 vUV;
out vec4 vColor;
out vec4 vParams;

void main() {
    gl_Position = uProj * vec4(aPos, 0.0, 1.0);
    vUV = aUV;
    vColor = aColor;
    vParams = aParams;
}
GLSL;

UI_FRAG_SRC_ES :: #string GLSL
#version 300 es
precision mediump float;
in vec2 vUV;
in vec4 vColor;
in vec4 vParams;

uniform sampler2D uTex;

out vec4 FragColor;

float roundedBoxSDF(vec2 center, vec2 half_size, float radius) {
    vec2 q = abs(center) - half_size + vec2(radius);
    return length(max(q, vec2(0.0))) + min(max(q.x, q.y), 0.0) - radius;
}

void main() {
    float mode = vParams.x;
    float border = vParams.y;
    vec2 rectSize = vParams.zw;

    if (mode < -1.5) {
        // Image mode (mode == -2.0): sample texture
        FragColor = texture(uTex, vUV) * vColor;
    } else if (mode < 0.0) {
        // Text mode (mode == -1.0): sample glyph atlas .r as alpha
        float alpha = texture(uTex, vUV).r;
        float ew = fwidth(alpha) * 0.7;
        alpha = smoothstep(0.5 - ew, 0.5 + ew, alpha);
        FragColor = vec4(vColor.rgb, vColor.a * pow(alpha, 0.9));
    } else if (mode > 0.0 || border > 0.0) {
        // Rounded rect: SDF alpha, vertex color only
        vec2 half_size = rectSize * 0.5;
        vec2 center = (vUV - vec2(0.5)) * rectSize;
        float dist = roundedBoxSDF(center, half_size, mode);
        float aa = fwidth(dist);
        float alpha = 1.0 - smoothstep(-aa, aa, dist);

        if (border > 0.0) {
            float inner = roundedBoxSDF(center, half_size - vec2(border), max(mode - border, 0.0));
            float border_alpha = smoothstep(-aa, aa, inner);
            alpha = alpha * max(border_alpha, 0.0);
        }

        FragColor = vec4(vColor.rgb, vColor.a * alpha);
    } else {
        // Plain rect: vertex color only
        FragColor = vColor;
    }
}
GLSL;

// --- Metal (MSL) — single library with vmain/fmain entry points ---
//
// `packed_float2 / packed_float4` keep the 12-float interleaved vertex
// layout (pos2 / uv2 / color4 / params4 = 48 bytes) without padding —
// MSL's default `float4` has 16-byte alignment and would force a 64-byte
// struct (see examples/63-metal-clear.sx for the gotcha).
//
// Uniform passing: GL uses `glUniformMatrix4fv("uProj", proj)`; Metal
// receives the projection via `setVertexBytes:length:atIndex:1` (slot 0
// is the vertex buffer). Texture binding goes through
// `setFragmentTexture:atIndex:0`.

UI_MSL_SRC :: #string MSL
#include <metal_stdlib>
using namespace metal;

struct UIVertex {
    packed_float2 pos;
    packed_float2 uv;
    packed_float4 color;
    packed_float4 params;
};

struct VOut {
    float4 position [[position]];
    float2 uv;
    float4 color;
    float4 params;
};

vertex VOut vmain(uint vid [[vertex_id]],
                  constant UIVertex* verts [[buffer(0)]],
                  constant float4x4& proj [[buffer(1)]]) {
    UIVertex v = verts[vid];
    VOut o;
    o.position = proj * float4(v.pos, 0.0, 1.0);
    o.uv = float2(v.uv);
    o.color = float4(v.color);
    o.params = float4(v.params);
    return o;
}

static float roundedBoxSDF(float2 center, float2 half_size, float radius) {
    float2 q = abs(center) - half_size + float2(radius);
    return length(max(q, float2(0.0))) + min(max(q.x, q.y), 0.0) - radius;
}

fragment float4 fmain(VOut in [[stage_in]],
                      texture2d<float> tex [[texture(0)]]) {
    constexpr sampler s(coord::normalized, address::clamp_to_edge, filter::linear);

    float mode = in.params.x;
    float border = in.params.y;
    float2 rectSize = in.params.zw;

    if (mode < -1.5) {
        // Image mode (mode == -2.0): sample texture
        return tex.sample(s, in.uv) * in.color;
    } else if (mode < 0.0) {
        // Text mode (mode == -1.0): the glyph atlas stores R8 alpha coverage.
        float alpha = tex.sample(s, in.uv).r;
        return float4(in.color.rgb, in.color.a * alpha);
    } else if (mode > 0.0 || border > 0.0) {
        // Rounded rect: SDF alpha, vertex color only
        float2 half_size = rectSize * 0.5;
        float2 center = (in.uv - float2(0.5)) * rectSize;
        float dist = roundedBoxSDF(center, half_size, mode);
        float aa = fwidth(dist);
        float alpha = 1.0 - smoothstep(-aa, aa, dist);
        if (border > 0.0) {
            float inner = roundedBoxSDF(center, half_size - float2(border), max(mode - border, 0.0));
            float border_alpha = smoothstep(-aa, aa, inner);
            alpha = alpha * max(border_alpha, 0.0);
        }
        return float4(in.color.rgb, in.color.a * alpha);
    } else {
        // Plain rect: vertex color only
        return in.color;
    }
}
MSL;