sokoban-v2/src/platform/graphics.zig

const std = @import("std");
const log = std.log.scoped(.graphics);
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;

const sokol = @import("sokol");
const sg = sokol.gfx;
const sapp = sokol.app; 
const sglue = sokol.glue;

const Math = @import("math");
const Vec2 = Math.Vec2;
const Vec4 = Math.Vec4;
const Mat4 = Math.Mat4;
const Rect = Math.Rect;

const STBRectPack = @import("stb_rect_pack");
const STBTrueType = @import("stb_truetype");
const ImGUI = @import("imgui.zig");
pub const Color = @import("./color.zig");
const shd = @import("shader");
const SlotMapType = @import("slot_map.zig").SlotMapType;

const Frame = struct {
    index: u32,
    quads: std.ArrayList(VertexQuad),
    transforms: std.ArrayList(TransformFrame),

    window_size: Vec2,
    clear_color: Color = .black,
    pass_started: bool = false
};

const State = struct {
    gpa: std.mem.Allocator,

    frame: Frame,

    shader: sg.Shader,
    pipeline: sg.Pipeline,
    bindings: sg.Bindings,

    default_sprite: Sprite.Id,
    error_sprite: Sprite.Id,
    default_spritesheet: Spritesheet.Id,

    linear_sampler: sg.Sampler,
    nearest_sampler: sg.Sampler,
    default_sampler: Sampler,

    quads_buffer: [8192]VertexQuad,
    transforms_buffer: [32]TransformFrame,

    textures_buffer: [Texture.max_textures]Texture.SlotMap.Slot,
    textures: Texture.SlotMap,

    spritesheets_buffer: [Spritesheet.max_spritesheets]Spritesheet.SlotMap.Slot,
    spritesheets: Spritesheet.SlotMap,

    sprites_buffer: [Sprite.max_sprites]Sprite.SlotMap.Slot,
    sprites: Sprite.SlotMap,
};

var g_state: State = undefined;

pub fn init(gpa: std.mem.Allocator, logger: sg.Logger) !void {
    const self = &g_state;

    sg.setup(.{
        .environment = sglue.environment(),
        .logger = logger
    });

    const shader = sg.makeShader(shd.mainShaderDesc(sg.queryBackend()));
    assert(shader.id != sg.invalid_id);

    const linear_sampler = sg.makeSampler(.{
        .min_filter = .LINEAR,
        .mag_filter = .LINEAR,
        .label = "linear-sampler"
    });
    assert(linear_sampler.id != sg.invalid_id);

    const nearest_sampler = sg.makeSampler(.{
        .min_filter = .NEAREST,
        .mag_filter = .NEAREST,
        .label = "nearest-sampler"
    });
    assert(linear_sampler.id != sg.invalid_id);

    const pipeline = sg.makePipeline(.{
        .primitive_type = .TRIANGLES,
        .index_type = .UINT16,
        .colors = init: {
            var s: [8]sg.ColorTargetState = [_]sg.ColorTargetState{.{}} ** 8;
            s[0] = .{
                .blend = .{
                    .enabled = true,
                    .src_factor_rgb = .SRC_ALPHA,
                    .dst_factor_rgb = .ONE_MINUS_SRC_ALPHA,
                    .op_rgb = .ADD,
                    .src_factor_alpha = .ONE,
                    .dst_factor_alpha = .ONE_MINUS_SRC_ALPHA,
                    .op_alpha = .ADD,
                }
            };
            break :init s;
        },
        .shader = shader,
        .layout = init: {
            var l = sg.VertexLayoutState{};
            l.attrs[shd.ATTR_main_position].format = .FLOAT2;
            l.attrs[shd.ATTR_main_texcoord0].format = .FLOAT2;
            l.attrs[shd.ATTR_main_color0].format = .FLOAT4;
            break :init l;
        },
        .label = "main-pipeline"
    });
    assert(pipeline.id != sg.invalid_id);

    var bindings: sg.Bindings = .{};
    {
        const max_quads = self.quads_buffer.len;

        bindings.vertex_buffers[0] = sg.makeBuffer(.{
            .size = @sizeOf(VertexQuad) * max_quads,
            .usage = .{ .vertex_buffer = true, .stream_update = true },
            .label = "quad-vertices"
        });

        const index_count = max_quads * 6;
        var indices: [index_count]u16 = undefined;
        for (0..max_quads) |i| {
            indices[6*i + 0] = @intCast(4*i + 0);
            indices[6*i + 1] = @intCast(4*i + 1);
            indices[6*i + 2] = @intCast(4*i + 2);

            indices[6*i + 3] = @intCast(4*i + 1);
            indices[6*i + 4] = @intCast(4*i + 3);
            indices[6*i + 5] = @intCast(4*i + 2);
        }

        bindings.index_buffer = sg.makeBuffer(.{
            .data = sg.asRange(&indices),
            .usage = .{ .index_buffer = true, },
            .label = "indices"
        });
    }

    self.* = State{
        .gpa = gpa,
        .shader = shader,
        .pipeline = pipeline,
        .bindings = bindings,
        .quads_buffer = undefined,
        .transforms_buffer = undefined,
        .default_sprite = undefined,
        .error_sprite = undefined,
        .default_spritesheet = undefined,

        .linear_sampler = linear_sampler,
        .nearest_sampler = nearest_sampler,
        .default_sampler = .nearest,

        .textures_buffer = undefined,
        .textures = .initBuffer(&self.textures_buffer),

        .sprites_buffer = undefined,
        .sprites = .initBuffer(&self.sprites_buffer),

        .spritesheets_buffer = undefined,
        .spritesheets = .initBuffer(&self.spritesheets_buffer),

        .frame = .{
            .index = 0,
            .quads = .initBuffer(&self.quads_buffer),
            .transforms = .initBuffer(&self.transforms_buffer),
            .clear_color = .black,
            .window_size = .init(0, 0)
        }
    };

    self.default_spritesheet = initSpritesheet(.{});

    {
        const width = 8;
        const height = 8;

        var pixels: [width * height * 4]u8 = undefined;
        @memset(&pixels, 0xFF);

        self.default_sprite = initSprite(.{ .padding = 0 });
        setSprite(self.default_sprite, .{
            .image = ImageData{
                .width = width,
                .height = height,
                .pixels = .{ .rgba8 = &pixels }
            }
        });
    }

    {
        const width = 8;
        const height = 8;

        var pixels: [width * height * 4]u8 = undefined;
        for (0..height) |y| {
            for (0..width) |x| {
                const pixel_index = y * width + x;

                var color: Color = undefined;
                if ((pixel_index+y) % 2 == 0) {
                    color = .purple;
                } else {
                    color = .black;
                }

                const pixel = pixels[(4*pixel_index)..];
                pixel[0] = @intFromFloat(color.r*255);
                pixel[1] = @intFromFloat(color.g*255);
                pixel[2] = @intFromFloat(color.b*255);
                pixel[3] = @intFromFloat(color.a*255);
            }
        }

        self.error_sprite = initSprite(.{ });
        setSprite(self.error_sprite, .{
            .image = .{
                .width = width,
                .height = height,
                .pixels = .{ .rgba8 = &pixels }
            }
        });
    }
}

pub fn deinit() void {
    const self = &g_state;

    var sprite_iter = self.sprites.iterator();
    while (sprite_iter.next()) |sprite_id| {
        const sprite = self.sprites.getAssumeExists(sprite_id);
        if (sprite.data) |data| {
            data.deinit(self.gpa);
        }
    }

    sg.shutdown();
}

fn toSokolColor(color: Color) sokol.gfx.Color {
    // TODO: Assert and do cast
    return sokol.gfx.Color{
        .r = color.r,
        .g = color.g,
        .b = color.b,
        .a = color.a,
    };
}

fn getTransformPtr() ?*TransformFrame {
    const self = &g_state;
    if (self.frame.transforms.items.len == 0) {
        return null;
    }

    return &self.frame.transforms.items[self.frame.transforms.items.len-1];
}

pub fn transformTranslate(offset: Vec2) void {
    const top_frame = getTransformPtr() orelse return;

    top_frame.offset = top_frame.offset.add(offset.multiply(top_frame.scale));
}

pub fn transformScale(scale: Vec2) void {
    const top_frame = getTransformPtr() orelse return;

    top_frame.scale = top_frame.scale.multiply(scale);
}

pub fn transformPush() void {
    const self = &g_state;

    const top_frame = getTransformPtr() orelse return;

    self.frame.transforms.appendBounded(top_frame.*) catch {
        log.warn("Transform stack limit reached!", .{});
    };
}

pub fn transformPop() void {
    const self = &g_state;
    _ = self.frame.transforms.pop();
}

pub fn getTransform() TransformFrame {
    if (getTransformPtr()) |top_frame| {
        return top_frame.*;
    } else {
        return TransformFrame{
            .offset = .init(0, 0),
            .scale = .init(1, 1)
        };
    }
}

pub fn beginFrame() void {
    const self = &g_state;

    self.frame.quads.clearRetainingCapacity();
    self.frame.transforms.clearRetainingCapacity();
    self.frame = Frame{
        .index = self.frame.index,
        .quads = self.frame.quads,
        .transforms = self.frame.transforms,
        .window_size = Vec2.init(sapp.widthf(), sapp.heightf())
    };

    self.frame.transforms.appendAssumeCapacity(TransformFrame{
        .offset = .init(0, 0),
        .scale = .init(1, 1),
    });

    self.bindings.views[shd.VIEW_tex].id = sg.invalid_id;
    self.bindings.samplers[shd.SMP_smp] = self.linear_sampler;
}

pub fn showDebug() void {
    const self = &g_state;

    _ = ImGUI.beginTabBar("graphics tab bar");
    defer ImGUI.endTabBar();

    if (ImGUI.beginTabItem("Spritesheets")) {
        defer ImGUI.endTabItem();

        if(ImGUI.beginTable("spritesheets", 2, 0)) {
            defer ImGUI.endTable();

            ImGUI.tableSetupColumn("ID", 0);
            ImGUI.tableSetupColumn("Size", 0);
            ImGUI.tableHeadersRow();

            var spritesheet_iter = self.spritesheets.iterator();
            while (spritesheet_iter.next()) |spritesheet_id| {
                const spritesheet = self.spritesheets.getAssumeExists(spritesheet_id);

                ImGUI.pushID(.{ .int = spritesheet_id.index });
                defer ImGUI.popID();

                ImGUI.tableNextRow();

                ImGUI.tableNextColumn();
                ImGUI.text("{{ {}, {} }}", .{spritesheet_id.index, spritesheet_id.generation});

                ImGUI.tableNextColumn();
                ImGUI.text("{}x{}", .{spritesheet.size.x, spritesheet.size.x});
            }
        }
    }

    if (ImGUI.beginTabItem("Sprites")) {
        defer ImGUI.endTabItem();

        if(ImGUI.beginTable("sprites", 2, 0)) {
            defer ImGUI.endTable();

            ImGUI.tableSetupColumn("ID", 0);
            ImGUI.tableSetupColumn("Position", 0);
            ImGUI.tableHeadersRow();

            var sprite_iter = self.sprites.iterator();
            while (sprite_iter.next()) |sprite_id| {
                const sprite = self.sprites.getAssumeExists(sprite_id);

                ImGUI.pushID(.{ .int = sprite_id.index });
                defer ImGUI.popID();

                ImGUI.tableNextRow();

                ImGUI.tableNextColumn();
                ImGUI.text("{{ {}, {} }}", .{sprite_id.index, sprite_id.generation});

                ImGUI.tableNextColumn();
                if (sprite.position) |position| {
                    ImGUI.text("{:3}, {:3}", .{position.x, position.y});
                } else {
                    ImGUI.text("null", .{});
                }
            }
        }
    }
}

pub fn setClearColor(color: Color) void {
    const self = &g_state;
    self.frame.clear_color = color;
}

pub fn getWindowSize() Vec2 {
    const self = &g_state;
    return self.frame.window_size;
}

pub fn getFrameIndex() u32 {
    const self = &g_state;
    return self.frame.index;
}

pub fn flush() void {
    const self = &g_state;

    if (self.frame.quads.items.len == 0) {
        return;
    }
    defer self.frame.quads.clearRetainingCapacity();

    if (!self.frame.pass_started) {
        var pass_action: sg.PassAction = .{};
        pass_action.colors[0] = .{
            .load_action = .CLEAR,
            .clear_value = toSokolColor(self.frame.clear_color)
        };

        sg.beginPass(.{
            .action = pass_action,
            .swapchain = sglue.swapchain()
        });
        self.frame.pass_started = true;
    }

    var spritesheet_iter = self.spritesheets.iterator();
    while (spritesheet_iter.next()) |spritesheet_id| {
        const spritesheet = self.spritesheets.getAssumeExists(spritesheet_id);
        const texture = self.textures.get(spritesheet.texture) orelse continue;
        if (isViewBound(texture.view)) {
            rebuildSpritesheetTextureIfNeeded(spritesheet_id);
        }
    }

    if (sg.queryViewState(getBoundView().*) != .VALID) {
        log.warn("Attempt to use view which isn't valid", .{});
        return;
    }

    const vertex_buffer = self.bindings.vertex_buffers[0];

    const data = sg.asRange(self.frame.quads.items);
    if (sg.queryBufferWillOverflow(vertex_buffer, data.size)) {
        log.warn("Vertex buffer overflow", .{});
        return;
    }

    self.bindings.vertex_buffer_offsets[0] = sg.appendBuffer(vertex_buffer, data);

    const window_size = Vec2.init(sapp.widthf(), sapp.heightf());

    var mvp = Mat4.initIdentity();
    {
        // Convert normalized [-1; 1] coordinates to [0; 1]
        mvp = mvp.multiply(Mat4.initTranslate(.{ .x = -1, .y = -1, .z = 0 }));
        mvp = mvp.multiply(Mat4.initScale(.{ .x = 2, .y = 2, .z = 0 }));

        // Make the top-left corner (0,0)
        mvp = mvp.multiply(Mat4.initTranslate(.{ .x = 0, .y = 1, .z = 0 }));
        mvp = mvp.multiply(Mat4.initScale(.{ .x = 1, .y = -1, .z = 0 }));

        // Scretch [0; 1] so that they map 1-to-1 on to screen coordinates.
        mvp = mvp.multiply(Mat4.initScale(.{
            .x = 1/window_size.x,
            .y = 1/window_size.y,
            .z = 0
        }));
    }

    var vs_params: shd.VsParams = .{
        .mvp = mvp
    };
    var fs_params: shd.FsParams = .{
        .texture_mode = 0,
    };
    const image = sg.queryViewImage(getBoundView().*);
    if (sg.queryImagePixelformat(image) == .R8) {
        fs_params.texture_mode = 1;
    }

    sg.applyPipeline(self.pipeline);
    sg.applyBindings(self.bindings);
    sg.applyUniforms(shd.UB_vs_params, sg.asRange(&vs_params));
    sg.applyUniforms(shd.UB_fs_params, sg.asRange(&fs_params));
    assert(self.frame.quads.items.len > 0);
    sg.draw(0, @intCast(self.frame.quads.items.len * 6), 1);
}

pub fn endFrame() void {
    const self = &g_state;

    flush();

    if (self.frame.pass_started) {
        sg.endPass();
        sg.commit();
    }

    var spritesheet_iter = self.spritesheets.iterator();
    while (spritesheet_iter.next()) |spritesheet_id| {
        repackSpritesheetIfNeeded(spritesheet_id);
    }

    if (self.frame.transforms.items.len > 1) {
        log.warn("Too many calls to transformPush()", .{});
    } else if (self.frame.transforms.items.len == 0) {
        log.warn("Too many calls to transformPop()", .{});
    }

    self.frame.index += 1;
}

pub fn drawRectangle(pos: Vec2, size: Vec2, color: Color) void {
    drawQuad(Quad.initRect(pos, size), color);
}

pub fn drawQuad(quad: Quad, color: Color) void {
    const self = &g_state;

    drawSpriteQuad(self.default_sprite, quad, color);
}

pub fn drawRectangleOutline(pos: Vec2, size: Vec2, color: Color, width: f32, alignment: f32) void {
    const self = &g_state;

    const outer_rect = (Rect{ .pos = pos, .size = size }).grow(Vec2.init(width, width).multiplyScalar(1-alignment));
    const inner_rect = (Rect{ .pos = pos, .size = size }).shrink(Vec2.init(width, width).multiplyScalar(alignment));

    drawSpriteQuad(
        self.default_sprite,
        .init(.{
            .init(outer_rect.right(), outer_rect.top()),
            .init(outer_rect.left(), outer_rect.top()),
            .init(inner_rect.right(), inner_rect.top()),
            .init(inner_rect.left(), inner_rect.top()),
        }),
        color
    );
    drawSpriteQuad(
        self.default_sprite,
        .init(.{
            .init(outer_rect.right(), outer_rect.top()),
            .init(inner_rect.right(), inner_rect.top()),
            .init(outer_rect.right(), outer_rect.bottom()),
            .init(inner_rect.right(), inner_rect.bottom()),
        }),
        color
    );
    drawSpriteQuad(
        self.default_sprite,
        .init(.{
            .init(outer_rect.left(), outer_rect.top()),
            .init(inner_rect.left(), inner_rect.top()),
            .init(outer_rect.left(), outer_rect.bottom()),
            .init(inner_rect.left(), inner_rect.bottom()),
        }),
        color
    );
    drawSpriteQuad(
        self.default_sprite,
        .init(.{
            .init(outer_rect.right(), outer_rect.bottom()),
            .init(outer_rect.left(), outer_rect.bottom()),
            .init(inner_rect.right(), inner_rect.bottom()),
            .init(inner_rect.left(), inner_rect.bottom()),
        }),
        color
    );
}

pub fn drawLine(from: Vec2, to: Vec2, color: Color, width: f32) void {
    const self = &g_state;

    const step = to.sub(from).normalized().multiplyScalar(width/2);

    const top_left     = from.add(step.rotateLeft90());
    const bottom_left  = from.add(step.rotateRight90());
    const top_right    = to.add(step.rotateLeft90());
    const bottom_right = to.add(step.rotateRight90());

    drawSpriteQuad(
        self.default_sprite,
        .init(.{ top_right, top_left, bottom_right, bottom_left }),
        color
    );
}

const TextureOptions = struct {
    dynamic: bool = false
};

pub fn initTexture(opts: TextureOptions) Texture.Id {
    const self = &g_state;
    const id = self.textures.insertUndefined() catch {
        log.warn("Failed to create texture, limit reached! limit: {}", .{self.textures.slots.capacity});
        return .nil;
    };

    const view = sg.allocView();
    if (sg.queryViewState(view) != .ALLOC) {
        log.warn("Failed to alloc texture view", .{});
        return .nil;
    }

    const texture = self.textures.getAssumeExists(id);
    texture.* = .{
        .dynamic = opts.dynamic,
        .updated_at = null,
        .image = .{ .id = sg.invalid_id },
        .view = view,
    };
    return id;
}

fn getBoundView() *sg.View {
    const self = &g_state;
    return &self.bindings.views[shd.VIEW_tex];
}

fn isViewBound(view: sg.View) bool {
    return getBoundView().id == view.id;
}

pub fn deinitTexture(id: Texture.Id) void {
    const self = &g_state;
    if (id == Texture.Id.nil) {
        return;
    }

    if (self.textures.get(id)) |texture| {
        const bound_view = getBoundView();
        if (bound_view.id ==  texture.view.id) {
            bound_view.id = sg.invalid_id;
        }

        sg.destroyImage(texture.image);
        sg.destroyView(texture.view);
        self.textures.removeAssumeExists(id);
    }
}

pub fn setTexture(id: Texture.Id, texture_data: ImageData) void {
    const self = &g_state;
    const texture = self.textures.get(id) orelse {
        log.warn("Attempt to set texture that doesn't exist: {}", .{id});
        return;
    };

    const pixel_count = texture_data.width*texture_data.height;

    var image_data: sg.ImageData = .{};
    var pixel_format: sg.PixelFormat = .DEFAULT;

    if (texture_data.pixels) |pixels| {
        switch (pixels) {
            .rgba8 => |rgba8| {
                image_data.mip_levels[0] = sg.asRange(rgba8[0..(pixel_count*4)]);
                pixel_format = .RGBA8;
            },
            .r8 => |r8| {
                image_data.mip_levels[0] = sg.asRange(r8[0..pixel_count]);
                pixel_format = .R8;
            }
        }
    }

    var new_image_desc = sg.ImageDesc{
        .type = ._2D,
        .width = @intCast(texture_data.width),
        .height = @intCast(texture_data.height),
        .num_mipmaps = 1,
        .pixel_format = pixel_format,
    };
    if (texture.dynamic) {
        new_image_desc.usage.dynamic_update = true;
    } else {
        new_image_desc.data = image_data;
    }

    var texture_updated: bool = false;
    if (sg.queryImageState(texture.image) == .VALID) {
        const image_desc = sg.queryImageDesc(texture.image);
        if (image_desc.width == new_image_desc.width and
            image_desc.height == new_image_desc.height and
            image_desc.num_mipmaps == new_image_desc.num_mipmaps and
            image_desc.pixel_format == new_image_desc.pixel_format and
            image_desc.usage.dynamic_update == new_image_desc.usage.dynamic_update and
            image_desc.usage.immutable == new_image_desc.usage.immutable and
            image_desc.usage.dynamic_update
        ) {
            log.debug("[{}] Update texture '{f}'", .{self.frame.index, id});
            if (texture.updated_at == null or texture.updated_at.? != self.frame.index) {
                sg.updateImage(texture.image, image_data);
                texture_updated = true;
                texture.updated_at = self.frame.index;
            } else {
                log.warn("Attempt to update the same texture multiple times per frame, texture: {f}", .{id});
            }
        }
    }

    if (!texture_updated) {
        log.debug("[{}] Create image-view, texture={f}", .{self.frame.index, id});
        const new_image = sg.makeImage(new_image_desc);
        if (sg.queryImageState(new_image) != .VALID) {
            log.warn("makeImage() failed", .{});
            return;
        }

        if (sg.queryViewState(texture.view) == .VALID) {
            sg.uninitView(texture.view);
        }
        sg.initView(texture.view, .{
            .texture = .{ .image = new_image }
        });
        if (sg.queryViewState(texture.view) != .VALID) {
            log.warn("makeView() failed", .{});
            return;
        }

        sg.destroyImage(texture.image);
        texture.image = new_image;
        texture.updated_at = null;

        log.debug("is dynamic {f}, {}", .{id, texture.dynamic});
        if (texture.dynamic) {
            log.debug("update texture {f}", .{id});
            sg.updateImage(new_image, image_data);
            texture.updated_at = self.frame.index;
        }
    }
}

pub fn drawTexture(id: Texture.Id, pos: Vec2, size: Vec2, color: Color) void {
    var quad: VertexQuad = undefined;
    quad.setColor(color);
    quad.setRect(.{ .pos = pos, .size = size });
    quad.setUVRect(.init(0, 0, 1, 1));

    draw(.{
        .texture = id,
        .quad = quad
    });
}

pub fn setDefaultSampler(sampler: Sampler) void {
    const self = &g_state;
    self.default_sampler = sampler;
}

pub const Sampler = enum {
    linear,
    nearest,
};

const DrawOptions = struct {
    sampler: ?Sampler = null,
    texture: Texture.Id,
    quad: VertexQuad
};

pub fn draw(opts: DrawOptions) void {
    const self = &g_state;

    var quad = opts.quad;

    var view: sg.View = .{ .id = sg.invalid_id };
    if (self.textures.get(opts.texture)) |texture| {
        view = texture.view;
    }

    if (sg.queryViewState(view) != .VALID and sg.queryViewState(view) != .ALLOC) {
        if (self.sprites.get(self.error_sprite)) |sprite| {
            if (self.spritesheets.get(sprite.spritesheet)) |spritesheet| {
                if (self.textures.get(spritesheet.texture)) |texture| {
                    if (getSpriteUVRect(self.error_sprite)) |uv| {
                        quad.setUVRect(uv);
                        quad.setColor(.white);
                        view = texture.view;
                    }
                }
            }
        }
    }

    if (sg.queryViewState(view) != .VALID and sg.queryViewState(view) != .ALLOC) {
        log.warn("Failed to bind view", .{});
        return;
    }

    const view_state = sg.queryViewState(view);
    if (view_state == .VALID or view_state == .ALLOC) {
        const bound_view = getBoundView();
        if (bound_view.id != view.id) {
            flush();
            bound_view.* = view;
        }
    } else {
        log.warn("Attempt to bind non-valid view: state={}", .{view_state});
        return;
    }

    const sampler = switch (opts.sampler orelse self.default_sampler) {
        .linear => self.linear_sampler,
        .nearest => self.nearest_sampler,
    };
    const bound_sampler = &self.bindings.samplers[shd.SMP_smp];
    if (bound_sampler.id != sampler.id) {
        flush();
        bound_sampler.* = sampler;
    }

    if (self.frame.quads.items.len == self.frame.quads.capacity) {
        flush();
    }

    if (getTransformPtr()) |transform| {
        for (&quad.vertices) |*vertex| {
            vertex.position = transform.apply(vertex.position);
        }
    }

    self.frame.quads.appendAssumeCapacity(quad);
}

const SpriteOptions = struct {
    spritesheet: ?Spritesheet.Id = null,
    padding: u32 = 0,
};

pub fn initSprite(opts: SpriteOptions) Sprite.Id {
    const self = &g_state;
    const id = self.sprites.insertUndefined() catch {
        log.warn("Failed to create sprite, limit reached! limit: {}", .{self.sprites.slots.capacity});
        return .nil;
    };

    const sprite = self.sprites.getAssumeExists(id);
    sprite.* = .{
        .spritesheet = opts.spritesheet orelse self.default_spritesheet,
        .padding = opts.padding,
        .data = null,
        .position = null
    };

    return id;
}

pub fn deinitSprite(id: Sprite.Id) void {
    const self = &g_state;
    if (id == Sprite.Id.nil) {
        return;
    }

    if (self.sprites.exists(id)) {
        const sprite = self.sprites.getAssumeExists(id);
        if (sprite.data) |data| {
            data.deinit(self.gpa);
        }
        self.sprites.removeAssumeExists(id);

        if (self.spritesheets.get(sprite.spritesheet)) |spritesheet| {
            spritesheet.needs_repack = true;
        }
    }
}

const SetSpriteOptions = struct {
    image: ImageData,
    rect: ?Rect = null
};

pub fn setSprite(id: Sprite.Id, opts: SetSpriteOptions) void {
    const self = &g_state;

    const sprite = self.sprites.get(id) orelse {
        log.warn("Attempt to set sprite that doesn't exist: {}", .{id});
        return;
    };

    const spritesheet = self.spritesheets.get(sprite.spritesheet) orelse {
        log.warn("Attempt to set sprite which doesn't have spritesheet", .{});
        return;
    };

    var sprite_data: ImageData = .none;
    errdefer sprite_data.deinit(self.gpa);

    if (opts.image.pixels) |pixels| {
        const format = @as(ImageData.Format, pixels);
        if (format != spritesheet.format) {
            log.warn("Attempt to use image of format '{}' on a spritesheet with '{}' format", .{format, spritesheet.format});
            return;
        }
    }

    if (opts.rect) |rect| {
        var format: ?ImageData.Format = null;
        if (opts.image.pixels) |pixels| {
            format = @as(ImageData.Format, pixels);
        }

        sprite_data = ImageData.init(self.gpa, @trunc(rect.size.x), @trunc(rect.size.y), format) catch |e| {
            log.warn("Failed to create texture data: {}", .{e});
            return;
        };

        sprite_data.copyFrom(
            opts.image,
            0, 0,
            @trunc(rect.pos.x), @trunc(rect.pos.y),
            @trunc(rect.size.x), @trunc(rect.size.y)
        );
    } else {
        sprite_data = opts.image.clone(self.gpa) catch |e| {
            log.warn("Failed to clone image: {}", .{e});
            return;
        };
    }

    var old_size: ?Vec2 = null;
    if (sprite.data) |old_sprite_data| {
        old_size = Vec2.initFromInt(u32, old_sprite_data.width, old_sprite_data.height);
        old_sprite_data.deinit(self.gpa);
    }
    sprite.data = sprite_data;

    const new_size = Vec2.initFromInt(u32, sprite_data.width, sprite_data.height); 

    log.debug("set sprite on {f}", .{sprite.spritesheet});

    if (old_size == null or !old_size.?.eql(new_size) or sprite.position == null) {
        sprite.position = null;
        spritesheet.needs_repack = true;
    }

    spritesheet.needs_texture_rebuild = true;
}

// WARNING: This will change the UV coordinates of all sprites.
// It's only safe to do this operation between frames where it's guarenteed that sprites aren't being used.
fn repackSpritesheet(id: Spritesheet.Id) void {
    const self = &g_state;

    const spritesheet = self.spritesheets.get(id) orelse {
        log.warn("Attempt to repack spritesheet which doesn't exist", .{});
        return;
    };

    log.debug("[{}] Repack spritesheet: {f}", .{self.frame.index, id});

    // TODO: Add a smarter startegy for picking the initial texture size.
    // One idea is to calculate the sum area of all sprites and pick a atlas size based on that.
    const texture_sizes = [_]u32{
        256, 512, 1024, 2048, 4096
    };

    var sprites_buffer: [Sprite.max_sprites]*Sprite = undefined;
    var sprites = std.ArrayList(*Sprite).initBuffer(&sprites_buffer);

    var sprite_iter = self.sprites.iterator();
    while (sprite_iter.next()) |sprite_id| {
        const sprite = self.sprites.getAssumeExists(sprite_id);
        if (sprite.data != null and sprite.spritesheet == id) {
            sprites.appendAssumeCapacity(sprite);
        }
    }

    var rects: [Sprite.max_sprites]STBRectPack.Rect = undefined;
    @memset(&rects, .{});
    for (0.., sprites.items) |i, sprite| {
        const sprite_data = sprite.data.?;
        rects[i] = .{
            .w = @intCast(sprite_data.width + sprite.padding * 2),
            .h = @intCast(sprite_data.height + sprite.padding * 2),
        };
    }

    spritesheet.packer = undefined;
    const packer = &spritesheet.packer.?;

    for (&texture_sizes) |texture_size| {
        packer.init(texture_size, texture_size, &spritesheet.nodes);
        packer.setupAllowOutOfMem(true);

        if (!packer.packRects(&rects)) {
            continue;
        }

        for (0.., sprites.items) |i, sprite| {
            assert(rects[i].was_packed == 1);
            const x: u32 = @as(u32, @intCast(rects[i].x)) + sprite.padding;
            const y: u32 = @as(u32, @intCast(rects[i].y)) + sprite.padding;
            sprite.position = .initFromInt(u32, x, y);
        }

        spritesheet.size = .initFromInt(u32, texture_size, texture_size);
        spritesheet.needs_texture_rebuild = true;
        spritesheet.needs_repack = false;
        return;
    }

    const max_texture_size = texture_sizes[texture_sizes.len-1];
    log.err("Failed to pack sprites, couldn't create a texture that is large enough. Max size tried: {}x{}", .{max_texture_size, max_texture_size});
    return;
}

fn repackSpritesheetIfNeeded(id: Spritesheet.Id) void {
    const self = &g_state;
    const spritesheet = self.spritesheets.get(id) orelse {
        log.warn("Attempt to pack spritesheet that doesn't exist: {f}", .{id});
        return;
    };

    if (!spritesheet.needs_repack) {
        return;
    }

    repackSpritesheet(id);
}

fn rebuildSpritesheetTextureIfNeeded(id: Spritesheet.Id) void {
    const self = &g_state;
    const spritesheet = self.spritesheets.get(id) orelse return;
    if (!spritesheet.needs_texture_rebuild) {
        return;
    }

    log.debug("[{}] Rebuild spritesheet texture: {f}", .{self.frame.index, id});

    const image_data = ImageData.init(
        self.gpa,
        @intFromFloat(spritesheet.size.x),
        @intFromFloat(spritesheet.size.y),
        spritesheet.format
    ) catch |e| {
        log.err("ImageData allocation failed: {}", .{e});
        return;
    };
    defer image_data.deinit(self.gpa);

    var sprite_iter = self.sprites.iterator();
    while (sprite_iter.next()) |sprite_id| {
        const sprite = self.sprites.getAssumeExists(sprite_id);
        if (sprite.spritesheet != id) {
            continue;
        }

        const sprite_data = sprite.data orelse continue;
        const sprite_position = sprite.position orelse continue;
        const x: u32 = @intFromFloat(sprite_position.x);
        const y: u32 = @intFromFloat(sprite_position.y);
        image_data.setRectZero(
            x - sprite.padding,
            y - sprite.padding,
            sprite_data.width + sprite.padding * 2,
            sprite_data.height + sprite.padding * 2
        );
        image_data.copyFrom(
            sprite_data,
            x,
            y,
            0,
            0,
            sprite_data.width,
            sprite_data.height
        );
    }

    setTexture(spritesheet.texture, image_data);

    spritesheet.needs_texture_rebuild = false;
} 

fn ensureSpritePacked(id: Sprite.Id) bool {
    const self = &g_state;
    const sprite = self.sprites.get(id) orelse return false;

    if (sprite.position == null) {
        const spritesheet = self.spritesheets.get(sprite.spritesheet) orelse return false;
        const sprite_data = sprite.data orelse return false;

        if (spritesheet.packer) |*packer| {
            var rects: [1]STBRectPack.Rect = .{
                .{
                    .w = @intCast(sprite_data.width + sprite.padding * 2),
                    .h = @intCast(sprite_data.height + sprite.padding * 2),
                }
            };
            if (packer.packRects(&rects)) {
                // Suboptimal sprite packing, will repack at end of frame
                assert(rects[0].was_packed == 1);
                const x: u32 = @as(u32, @intCast(rects[0].x)) + sprite.padding;
                const y: u32 = @as(u32, @intCast(rects[0].y)) + sprite.padding;
                sprite.position = .initFromInt(u32, x, y);
                spritesheet.needs_repack = true;
            } else {
                log.debug("Failed to pack a single sprite to spritesheet", .{});
            }
        } else {
            repackSpritesheet(sprite.spritesheet);
        }
    }

    return sprite.position != null;
}

pub fn getSpriteUVRect(id: Sprite.Id) ?Rect {
    const self = &g_state;
    const sprite = self.sprites.get(id) orelse return null;
    const sprite_data = sprite.data orelse return null;

    const spritesheet = self.spritesheets.get(sprite.spritesheet) orelse return null;

    if (!ensureSpritePacked(id)) {
        return null;
    }

    assert(sprite.position != null);


    const sprite_size = Vec2.initFromInt(u32, sprite_data.width, sprite_data.height);
    var uv_rect = Rect{
        .pos = sprite.position.?.divide(spritesheet.size),
        .size = sprite_size.divide(spritesheet.size),
    };

    const texel_size = Vec2.init(1, 1).divide(spritesheet.size);
    uv_rect = uv_rect.shrink(texel_size.divideScalar(8));

    return uv_rect;
}

pub fn drawSpriteQuad(id: Sprite.Id, points: Quad, color: Color) void {
    const self = &g_state;

    var quad: VertexQuad = undefined;
    quad.setColor(color);
    quad.vertices[0].position = points.positions[0];
    quad.vertices[1].position = points.positions[1];
    quad.vertices[2].position = points.positions[2];
    quad.vertices[3].position = points.positions[3];

    var texture: ?Texture.Id = null;
    if (getSpriteUVRect(id)) |sprite_uv| {
        quad.setUVRect(sprite_uv);

        const sprite = self.sprites.getAssumeExists(id);
        const spritesheet = self.spritesheets.getAssumeExists(sprite.spritesheet);
        texture = spritesheet.texture;
    }

    draw(.{
        .texture = texture orelse .nil,
        .quad = quad
    });
}

pub fn drawSprite(id: Sprite.Id, pos: Vec2, size: Vec2, color: Color) void {
    const rect = Quad.initRect(pos, size);
    drawSpriteQuad(id, rect, color);
}

pub const SpriteSheetOptions = struct {
    format: ImageData.Format = .rgba8
};

pub fn initSpritesheet(opts: SpriteSheetOptions) Spritesheet.Id {
    const self = &g_state;
    const id = self.spritesheets.insertUndefined() catch {
        log.warn("Failed to create spritesheet, limit reached! limit: {}", .{self.spritesheets.slots.capacity});
        return .nil;
    };

    const spritesheet = self.spritesheets.getAssumeExists(id);
    spritesheet.* = .{
        .texture = initTexture(.{ .dynamic = true }),
        .size = .init(0, 0),
        .format = opts.format,
        .needs_repack = false,
        .needs_texture_rebuild = false,
        .nodes = undefined,
        .packer = null
    };

    return id;
}

pub fn deinitSpritesheet(id: Spritesheet.Id) void {
    const self = &g_state;
    if (self.spritesheets.get(id)) |spritesheet| {
        deinitTexture(spritesheet.texture);
        self.spritesheets.removeAssumeExists(id);
    }
}

pub const ImageData = struct {
    const Format = enum {
        rgba8,
        r8
    };

    const Pixels = union(Format) {
        rgba8: [*]u8,
        r8: [*]u8,
    };

    pixels: ?Pixels,
    width: u32,
    height: u32,

    pub const none = ImageData{
        .pixels = null,
        .width = 0,
        .height = 0,
    };

    pub fn init(gpa: Allocator, width: u32, height: u32, format: ?Format) !ImageData {
        const pixel_count = width * height;
        var pixels: ?Pixels = null;
        if (format != null) {
            pixels = switch (format.?) {
                .rgba8 => Pixels{
                    .rgba8 = (try gpa.alloc(u8, pixel_count * 4)).ptr
                },
                .r8 => Pixels{
                    .r8 = (try gpa.alloc(u8, pixel_count)).ptr
                }
            };
        }

        return ImageData{
            .pixels = pixels,
            .width = width,
            .height = height
        };
    }

    pub fn initRGBA8(gpa: Allocator, width: u32, height: u32) !ImageData {
        return ImageData.init(gpa, width, height, .rgba8);
    }

    pub fn initR8(gpa: Allocator, width: u32, height: u32) !ImageData {
        return ImageData.init(gpa, width, height, .r8);
    }

    pub fn clone(self: ImageData, gpa: Allocator) !ImageData {
        const pixel_count = self.width * self.height;

        var pixels: ?Pixels = null; 
        if (self.pixels) |self_pixels| {
            pixels = switch (self_pixels) {
                .rgba8 => |rgba8| Pixels{
                    .rgba8 = (try gpa.dupe(u8, rgba8[0..(pixel_count * 4)])).ptr
                },
                .r8 => |r8| Pixels{
                    .r8 = (try gpa.dupe(u8, r8[0..pixel_count])).ptr
                }
            };
        }

        return ImageData{
            .pixels = pixels,
            .width = self.width,
            .height = self.height
        };
    }

    pub fn deinit(self: ImageData, gpa: Allocator) void {
        const pixel_count = self.width * self.height;
        if (self.pixels) |pixels| {
            switch (pixels) {
                .rgba8 => |rgba8| gpa.free(rgba8[0..(pixel_count*4)]),
                .r8 => |r8| gpa.free(r8[0..pixel_count])
            }
        }
    }

    pub fn getPixel(self: ImageData, x: u32, y: u32) []u8 {
        const pixel_index = y * self.width + x;
        if (self.pixels) |pixels| {
            return switch (pixels) {
                .rgba8 => |rgba8| rgba8[(4*pixel_index)..][0..4],
                .r8 => |r8| r8[pixel_index..][0..1]
            };
        } else {
            return &.{};
        }
    }

    pub fn setRect(self: ImageData, x: u32, y: u32, width: u32, height: u32, pixel: []const u8) void {
        if (self.pixels == null) {
            return;
        }

        for (0..height) |oy| {
            const iy = y + oy;
            if (iy >= self.height) {
                continue;
            }

            for (0..width) |ox| {
                const ix = x + ox;
                if (ix >= self.width) {
                    continue;
                }

                const dest_pixel = self.getPixel(@intCast(ix), @intCast(iy));
                @memcpy(dest_pixel, pixel);
            }
        }
    }

    pub fn setRectZero(self: ImageData, x: u32, y: u32, width: u32, height: u32) void {
        if (self.pixels == null) {
            return;
        }

        const pixel = switch (self.pixels.?) {
            .rgba8 => &[_]u8{ 0, 0, 0, 0 },
            .r8 => &[_]u8{ 0 }
        };
        self.setRect(x, y, width, height, pixel);
    }

    pub fn copyFrom(
        self: ImageData,
        source: ImageData,
        cx: u32,
        cy: u32,
        source_cx: u32,
        source_cy: u32,
        source_width: u32,
        source_height: u32,
    ) void {
        if (source.pixels == null or self.pixels == null) {
            return;
        }

        for (0..source_height) |oy| {
            const source_y = source_cy+oy;
            const dest_y = cy+oy; 
            if (dest_y >= self.height) {
                continue;
            }
            if (source_y >= source.height) {
                continue;
            }

            for (0..source_width) |ox| {
                const source_x = source_cx+ox;
                const dest_x = cx+ox; 
                if (dest_x >= self.width) {
                    continue;
                }
                if (source_x >= source.width) {
                    continue;
                }

                const source_pixel = source.getPixel(@intCast(source_x), @intCast(source_y));
                const dest_pixel = self.getPixel(@intCast(dest_x), @intCast(dest_y));

                @memcpy(dest_pixel, source_pixel);
            }
        }
    }
};

pub const Texture = struct {
    dynamic: bool,
    updated_at: ?u64,
    image: sg.Image,
    view: sg.View,

    const max_textures = 64;
    const SlotMap = SlotMapType(std.math.IntFittingRange(0, max_textures-1), u8, Texture);
    pub const Id = SlotMap.Id;
};

pub const Sprite = struct {
    spritesheet: Spritesheet.Id,
    data: ?ImageData,
    position: ?Vec2,
    padding: u32,

    const max_sprites = 256;
    const SlotMap = SlotMapType(std.math.IntFittingRange(0, max_sprites-1), u8, Sprite);
    pub const Id = SlotMap.Id;
};

pub const Spritesheet = struct {
    texture: Texture.Id,
    format: ImageData.Format,
    size: Vec2,

    needs_repack: bool,
    needs_texture_rebuild: bool,

    packer: ?STBRectPack,
    nodes: [256]STBRectPack.Node,

    const max_spritesheets = 64;
    const SlotMap = SlotMapType(std.math.IntFittingRange(0, max_spritesheets-1), u8, Spritesheet);
    pub const Id = SlotMap.Id;
};

pub const Vertex = extern struct {
    position: Vec2,
    texcoord: Vec2,
    color: Vec4,
};

pub const VertexQuad = extern struct {
    vertices: [4]Vertex,

    pub fn init(vertices: [4]Vertex) VertexQuad {
        return VertexQuad{
            .vertices = vertices
        };
    }

    pub fn setColor(self: *VertexQuad, color: Color) void {
        for (&self.vertices) |*vertex| {
            vertex.color = Vec4{
                .x = color.r,
                .y = color.g,
                .z = color.b,
                .w = color.a,
            };
        }
    }

    pub fn setUVRect(self: *VertexQuad, rect: Rect) void {
        const pos = rect.pos;
        const size = rect.size;
        self.vertices[0].texcoord = pos;
        self.vertices[1].texcoord = pos.add(.init(size.x, 0));
        self.vertices[2].texcoord = pos.add(.init(0, size.y));
        self.vertices[3].texcoord = pos.add(size);
    }

    pub fn setRect(self: *VertexQuad, rect: Rect) void {
        const pos = rect.pos;
        const size = rect.size;
        self.vertices[0].position = pos;
        self.vertices[1].position = pos.add(.init(size.x, 0));
        self.vertices[2].position = pos.add(.init(0, size.y));
        self.vertices[3].position = pos.add(size);
    }
};

pub const Quad = struct {
    positions: [4]Vec2,

    pub fn init(positions: [4]Vec2) Quad {
        return Quad{
            .positions = positions
        };
    }

    pub fn initRect(pos: Vec2, size: Vec2) Quad {
        return Quad{
            .positions = .{
                pos,
                pos.add(.init(size.x, 0)),
                pos.add(.init(0, size.y)),
                pos.add(size),
            }
        };
    }

    pub fn applyRotate(self: *Quad, rad: f32, origin: Vec2) void {
        for (&self.positions) |*pos| {
            pos.* = pos.sub(origin).rotate(rad).add(origin);
        }
    }

    pub fn applyOffset(self: *Quad, vec2: Vec2) void {
        for (&self.positions) |*pos| {
            pos.* = pos.add(vec2);
        }
    }
};

pub const TransformFrame = struct {
    offset: Vec2,
    scale: Vec2,

    pub fn apply(self: TransformFrame, position: Vec2) Vec2 {
        return position.multiply(self.scale).add(self.offset);
    }
};