#import "modules/std.sx";
#import "modules/compiler.sx";
#import "modules/opengl.sx";
#import "modules/math";
#import "ui/types.sx";
#import "ui/render.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 {
    vao: u32;
    vbo: u32;
    shader: u32;
    proj_loc: s32;
    tex_loc: s32;
    vertices: [*]f32;
    vertex_count: s64;
    screen_width: f32;
    screen_height: f32;
    white_texture: u32;
    current_texture: u32;

    init :: (self: *UIRenderer) {
        // Create shader (ES for WASM/WebGL2, Core for desktop)
        inline if OS == .wasm {
            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");

        // Allocate vertex buffer (CPU side)
        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;

        // 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) {
        self.screen_width = width;
        self.screen_height = height;
        self.vertex_count = 0;
        self.current_texture = self.white_texture;
    }

    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);
    }

    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 xx g_font != 0 and g_font.char_data != null {
                        self.render_text(node);
                    }
                }
                case .image: {
                    self.bind_texture(node.texture_id);
                    self.push_quad(node.frame, COLOR_WHITE, 0.0, 0.0);
                    self.bind_texture(self.white_texture);
                }
                case .clip_push: {
                    self.flush();
                    glEnable(GL_SCISSOR_TEST);
                    glScissor(
                        xx node.frame.origin.x,
                        xx (self.screen_height - node.frame.origin.y - node.frame.size.height),
                        xx node.frame.size.width,
                        xx node.frame.size.height
                    );
                }
                case .clip_pop: {
                    self.flush();
                    glDisable(GL_SCISSOR_TEST);
                }
                case .opacity_push: {}
                case .opacity_pop: {}
            }
            i += 1;
        }
    }

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

        glUseProgram(self.shader);

        // Orthographic projection: (0,0) top-left, (w,h) bottom-right
        proj := Mat4.ortho(0.0, self.screen_width, self.screen_height, 0.0, -1.0, 1.0);
        glUniformMatrix4fv(self.proj_loc, 1, 0, proj.data);

        // Bind current texture
        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, self.current_texture);
        glUniform1i(self.tex_loc, 0);

        glBindVertexArray(self.vao);
        glBindBuffer(GL_ARRAY_BUFFER, self.vbo);

        upload_size : s64 = self.vertex_count * UI_VERTEX_BYTES;
        glBufferSubData(GL_ARRAY_BUFFER, 0, xx upload_size, self.vertices);

        glDrawArrays(GL_TRIANGLES, 0, xx self.vertex_count);

        glBindVertexArray(0);
        self.vertex_count = 0;
    }

    render_text :: (self: *UIRenderer, node: RenderNode) {
        font := g_font;
        scale := node.font_size / font.font_size;

        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();

        // stbtt_GetBakedQuad works at baked size; we scale output positions
        xpos : f32 = 0.0;
        ypos : f32 = 0.0;
        q : AlignedQuad = ---;
        i : s64 = 0;
        while i < node.text.len {
            ch : s32 = xx node.text[i];
            if ch >= FIRST_CHAR and ch < FIRST_CHAR + NUM_CHARS {
                stbtt_GetBakedQuad(xx font.char_data, ATLAS_W, ATLAS_H, ch - FIRST_CHAR, @xpos, @ypos, xx @q, 1);

                // Scale and offset to frame position
                // ypos=0 means baseline is at y=0; glyphs go above (negative yoff)
                // Add ascent so top of text aligns with frame top
                gx0 := node.frame.origin.x + q.x0 * scale;
                gy0 := node.frame.origin.y + font.ascent * scale + q.y0 * scale;
                gx1 := node.frame.origin.x + q.x1 * scale;
                gy1 := node.frame.origin.y + font.ascent * scale + q.y1 * scale;

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

                // corner_radius = -1.0 signals "text mode" to the fragment shader
                self.write_vertex(gx0, gy0, q.s0, q.t0, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
                self.write_vertex(gx1, gy0, q.s1, q.t0, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
                self.write_vertex(gx0, gy1, q.s0, q.t1, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
                self.write_vertex(gx1, gy0, q.s1, q.t0, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
                self.write_vertex(gx1, gy1, q.s1, q.t1, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
                self.write_vertex(gx0, gy1, q.s0, q.t1, r, g, b, a, 0.0 - 1.0, 0.0, 0.0, 0.0);
            }
            i += 1;
        }

        self.bind_texture(self.white_texture);
    }
}

// Upload font atlas bitmap as GL texture (called after GL init)
upload_font_texture :: (font: *FontAtlas) {
    if font.bitmap == null { return; }
    glGenTextures(1, @font.texture_id);
    glBindTexture(GL_TEXTURE_2D, font.texture_id);
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
    glTexImage2D(GL_TEXTURE_2D, 0, xx GL_R8, ATLAS_W, ATLAS_H, 0, GL_RED, GL_UNSIGNED_BYTE, font.bitmap);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, xx GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, xx GL_LINEAR);
    context.allocator.dealloc(font.bitmap);
    font.bitmap = null;
}

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 radius = vParams.x;
    float border = vParams.y;
    vec2 rectSize = vParams.zw;

    if (radius < 0.0) {
        float textAlpha = texture(uTex, vUV).r;
        FragColor = vec4(vColor.rgb, vColor.a * textAlpha);
    } else if (radius > 0.0 || border > 0.0) {
        vec4 texColor = texture(uTex, vUV);
        vec2 half_size = rectSize * 0.5;
        vec2 center = (vUV - vec2(0.5)) * rectSize;
        float dist = roundedBoxSDF(center, half_size, radius);
        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(radius - border, 0.0));
            float border_alpha = smoothstep(-aa, aa, inner);
            alpha = alpha * max(border_alpha, 0.0);
        }

        FragColor = vec4(texColor.rgb * vColor.rgb, texColor.a * vColor.a * alpha);
    } else {
        vec4 texColor = texture(uTex, vUV);
        FragColor = texColor * 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 radius = vParams.x;
    float border = vParams.y;
    vec2 rectSize = vParams.zw;

    if (radius < 0.0) {
        float textAlpha = texture(uTex, vUV).r;
        FragColor = vec4(vColor.rgb, vColor.a * textAlpha);
    } else if (radius > 0.0 || border > 0.0) {
        vec4 texColor = texture(uTex, vUV);
        vec2 half_size = rectSize * 0.5;
        vec2 center = (vUV - vec2(0.5)) * rectSize;
        float dist = roundedBoxSDF(center, half_size, radius);
        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(radius - border, 0.0));
            float border_alpha = smoothstep(-aa, aa, inner);
            alpha = alpha * max(border_alpha, 0.0);
        }

        FragColor = vec4(texColor.rgb * vColor.rgb, texColor.a * vColor.a * alpha);
    } else {
        vec4 texColor = texture(uTex, vUV);
        FragColor = texColor * vColor;
    }
}
GLSL;