#import "modules/std.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;
    vertices: [*]f32;
    vertex_count: s64;
    screen_width: f32;
    screen_height: f32;
    white_texture: u32;

    init :: (self: *UIRenderer) {
        // Create shader
        self.shader = create_program(UI_VERT_SRC, UI_FRAG_SRC);

        self.proj_loc = glGetUniformLocation(self.shader, "uProj");

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

    // 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: {
                    // TODO: text rendering via glyph atlas
                    // For now, render a placeholder rect
                    self.push_quad(node.frame, node.text_color, 0.0, 0.0);
                }
                case .image: {
                    // TODO: textured quad
                    self.push_quad(node.frame, COLOR_WHITE, 0.0, 0.0);
                }
                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);

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

// Additional GL functions needed for UI renderer
GL_SCISSOR_TEST     :u32: 0x0C11;
GL_DYNAMIC_DRAW     :u32: 0x88E8;
GL_TEXTURE_2D       :u32: 0x0DE1;
GL_TEXTURE_MIN_FILTER :u32: 0x2801;
GL_TEXTURE_MAG_FILTER :u32: 0x2800;
GL_NEAREST          :u32: 0x2600;
GL_RGBA             :u32: 0x1908;
GL_UNSIGNED_BYTE    :u32: 0x1401;
GL_SRC_ALPHA        :u32: 0x0302;
GL_ONE_MINUS_SRC_ALPHA :u32: 0x0303;

glScissor            : (s32, s32, s32, s32) -> void = ---;
glBufferSubData      : (u32, s64, s64, *void) -> void = ---;
glGenTextures        : (s32, *u32) -> void = ---;
glBindTexture        : (u32, u32) -> void = ---;
glTexImage2D         : (u32, s32, s32, s32, s32, s32, u32, u32, *void) -> void = ---;
glTexParameteri      : (u32, u32, s32) -> void = ---;
glBlendFunc          : (u32, u32) -> void = ---;
glReadPixels         : (s32, s32, s32, s32, u32, u32, *void) -> void = ---;

// Load additional GL functions (call after load_gl)
load_gl_ui :: (get_proc: ([:0]u8) -> *void) {
    glScissor       = xx get_proc("glScissor");
    glBufferSubData = xx get_proc("glBufferSubData");
    glGenTextures   = xx get_proc("glGenTextures");
    glBindTexture   = xx get_proc("glBindTexture");
    glTexImage2D    = xx get_proc("glTexImage2D");
    glTexParameteri = xx get_proc("glTexParameteri");
    glBlendFunc     = xx get_proc("glBlendFunc");
    glReadPixels    = xx get_proc("glReadPixels");
}

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 ---

UI_VERT_SRC :: #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 :: #string GLSL
#version 330 core
in vec2 vUV;
in vec4 vColor;
in vec4 vParams; // corner_radius, border_width, rect_w, rect_h

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 || border > 0.0) {
        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(vColor.rgb, vColor.a * alpha);
    } else {
        FragColor = vColor;
    }
}
GLSL;