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 Color = @import("./color.zig");
const shd = @import("shader");
const SlotMapType = @import("slot_map.zig").SlotMapType;

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

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

    default_sprite: Sprite.Id,
    error_sprite: Sprite.Id,
    default_sampler: sg.Sampler,

    quads_buffer: [2048]Quad,
    quads: std.ArrayList(Quad),

    clear_color: Color,

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

    spritesheet_texture: Texture.Id,
    spritesheet_size: Vec2,
    spritesheet_stale: bool,
    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 default_sampler = sg.makeSampler(.{
        .min_filter = .NEAREST,
        .mag_filter = .NEAREST,
        .label = "default-sampler"
    });
    assert(default_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(Quad) * 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,
        .default_sampler = default_sampler,
        .quads_buffer = undefined,
        .quads = .initBuffer(&self.quads_buffer),
        .clear_color = .black,
        .textures_buffer = undefined,
        .textures = .init(.initBuffer(&self.textures_buffer)),
        .default_sprite = undefined,
        .error_sprite = undefined,
        .spritesheet_stale = false,
        .sprites_buffer = undefined,
        .sprites = .init(.initBuffer(&self.sprites_buffer)),
        .spritesheet_texture = undefined,
        .spritesheet_size = .init(0, 0)
    };

    {
        const width = 8;
        const height = 8;

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

        self.default_sprite = try initSprite();
        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 = try initSprite();
        setSprite(self.error_sprite, .{
            .image = .{
                .width = width,
                .height = height,
                .pixels = .{ .rgba8 = &pixels }
            }
        });
    }

    self.spritesheet_texture = try initTexture();
}

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.destroyPipeline(self.pipeline);
    sg.destroyShader(self.shader);
    sg.destroySampler(self.default_sampler);
    sg.shutdown();
}

fn createProjectionMatrix(screen_size: Vec2) Mat4 {
    return Mat4.initIdentity()
        // Convert normalized [-1; 1] coordinates to [0; 1]
        .multiply(Mat4.initTranslate(.{ .x = -1, .y = -1, .z = 0 }))
        .multiply(Mat4.initScale(.{ .x = 2, .y = 2, .z = 0 }))

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

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

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

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

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

    sg.beginPass(.{
        .action = pass_action,
        .swapchain = sglue.swapchain()
    });

    const spritesheet_texture = self.textures.getAssumeExists(self.spritesheet_texture);
    self.bindings.views[shd.VIEW_tex] = spritesheet_texture.view;
    self.bindings.samplers[shd.SMP_smp] = self.default_sampler;
}

// TODO: This will always have a 1-frame delay.
// fix this.
pub fn setClearColor(color: Color) void {
    const self = &g_state;

    self.clear_color = color;
}

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

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

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

    const data = sg.asRange(self.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);

    // TODO: Rectify `window_size` vs `gfx.window_size`
    const window_size = Vec2.init(sapp.widthf(), sapp.heightf());

    var vs_params: shd.VsParams = .{
        // TODO: Maybe 'view' matrix is not needed.
        // We probably should only have a combined 'mvp' matrix
        .view = .initIdentity(),
        .projection = createProjectionMatrix(window_size)
    };

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

pub fn endFrame() void {
    flush();

    sg.endPass();
    sg.commit();
}

pub fn drawRectangle(pos: Vec2, size: Vec2, color: Color) void {
    const self = &g_state;

    drawSprite(self.default_sprite, pos, size, color);
}

pub fn drawLine(from: Vec2, to: Vec2, color: Vec4, 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());

    var quad: Quad = undefined;
    quad.setColor(color);
    quad.setUVRect(.init(0, 0, 1, 1));
    quad.vertices[0] = top_right;
    quad.vertices[1] = top_left;
    quad.vertices[2] = bottom_right;
    quad.vertices[3] = bottom_left;

    draw(.{
        .texture = self.default_sprite,
        .quad = quad
    });
}

pub fn initTexture() !Texture.Id {
    const self = &g_state;
    const id = try self.textures.insertBounded();

    const image = sg.allocImage();
    assert(image.id != sg.invalid_id);

    const view = sg.allocView();
    assert(view.id != sg.invalid_id);

    const texture = self.textures.getAssumeExists(id);
    texture.* = .{
        .image = image,
        .view = view
    };
    return id;
}

pub fn deinitTexture(id: Texture.Id) void {
    const self = &g_state;
    if (self.textures.remove(id)) {
        const texture = self.textures_buffer[id.index];
        const bound_view = &self.bindings.views[shd.VIEW_tex];
        if (texture.view.id == bound_view.id) {
            assert(self.textures.exists(self.default_sprite));
            bound_view.* = self.textures_buffer[self.default_sprite.index].view;
        }

        sg.destroyView(texture.view);
        sg.destroyImage(texture.image);
    }
}

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

    if (sg.queryImageState(texture.image) == .VALID) {
        sg.uninitImage(texture.image);
    }

    const pixel_count = texture_data.width*texture_data.height;

    var image_data: sg.ImageData = .{};
    image_data.mip_levels[0] = sg.asRange(texture_data.pixels.rgba8[0..(pixel_count*4)]);
    sg.initImage(texture.image, .{
        .type = ._2D,
        .width = @intCast(texture_data.width),
        .height = @intCast(texture_data.height),
        .num_mipmaps = 1,
        .pixel_format = .RGBA8,
        .data = image_data,
    });
    assert(sg.queryImageState(texture.image) == .VALID);

    if (sg.queryViewState(texture.view) == .ALLOC) {
        sg.initView(texture.view, .{
            .texture = .{ .image = texture.image }
        });
        assert(sg.queryViewState(texture.view) == .VALID);
    }
}

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

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

const DrawOptions = struct {
    texture: Texture.Id,
    quad: Quad
};

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

    var quad = opts.quad;
    var texture = opts.texture;
    if (!self.textures.exists(texture)) {
        if (getSpriteUVRect(self.error_sprite)) |error_uv| {
            texture = self.spritesheet_texture;
            quad.setUVRect(error_uv);
            quad.setColor(.white);
        } else {
            log.err("Failed to draw error texture", .{});
            return;
        }
    }

    const view = self.textures.getAssumeExists(texture).view;

    const bound_view = &self.bindings.views[shd.VIEW_tex];
    if (bound_view.id != view.id) {
        flush();
        bound_view.* = view;
    }

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

    self.quads.appendAssumeCapacity(quad);
}

pub fn initSprite() !Sprite.Id {
    const self = &g_state;
    const id = try self.sprites.insertBounded();

    const sprite = self.sprites.getAssumeExists(id);
    sprite.* = .{
        .data = null,
        .position = null
    };

    return id;
}

pub fn deinitSprite(id: Sprite.Id) void {
    const self = &g_state;
    if (self.sprites.exists(id)) {
        const sprite = self.sprites.getAssumeExists(id);
        if (sprite.data) |data| {
            data.deinit(self.gpa);
        }
        self.sprites.removeAssumeExists(id);
        self.spritesheet_stale = true;
    }
}

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

pub fn setSprite(id: Sprite.Id, opts: SetSpriteOptions) void {
    const self = &g_state;
    if (!self.sprites.exists(id)) {
        log.warn("Attempt to set data for sprite that doesn't exist: {}", .{id});
        return;
    }

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

    if (opts.rect) |rect| {
        sprite_data = ImageData.initRGBA8(self.gpa, @trunc(rect.size.x), @trunc(rect.size.y)) 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;
        };
    }

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

    sprite.data = sprite_data;
    sprite.position = null;
    self.spritesheet_stale = true;
}

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

    // 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{
        64, 128, 256, 512, 1024, 2048, 4096
    };

    var nodes: [32]STBRectPack.Node = undefined;
    var rects: [Sprite.max_sprites]STBRectPack.Rect = undefined;

    for (&texture_sizes) |texture_size| {
        var ctx: STBRectPack = undefined;
        ctx.init(texture_size, texture_size, &nodes);
        ctx.setupAllowOutOfMem(true);

        @memset(&rects, .{});


        {
            var i: usize = 0;
            var sprite_iter = self.sprites.iterator();
            while (sprite_iter.next()) |sprite_id| : (i += 1) {
                const sprite = self.sprites.getAssumeExists(sprite_id);
                const sprite_data = sprite.data orelse continue;
                rects[i] = .{
                    .w = @intCast(sprite_data.width),
                    .h = @intCast(sprite_data.height),
                };
            }
        }

        const success = ctx.packRects(&rects);
        std.debug.print("success {} {}\n", .{texture_size, success});
        if (!success) {
            continue;
        }

        const spritesheet = try ImageData.initRGBA8(self.gpa, texture_size, texture_size);
        defer spritesheet.deinit(self.gpa);

        {
            var i: usize = 0;
            var sprite_iter = self.sprites.iterator();
            while (sprite_iter.next()) |sprite_id| : (i += 1) {
                const sprite = self.sprites.getAssumeExists(sprite_id);
                const sprite_data = sprite.data orelse continue;

                assert(rects[i].was_packed == 1);
                const x: u32 = @intCast(rects[i].x);
                const y: u32 = @intCast(rects[i].y);
                sprite.position = .initFromInt(u32, x, y);
                spritesheet.copyFrom(
                    sprite_data,
                    x,
                    y,
                    0,
                    0,
                    sprite_data.width,
                    sprite_data.height
                );
            }
        }

        setTexture(self.spritesheet_texture, spritesheet);

        self.spritesheet_size = .initFromInt(u32, texture_size, texture_size);
        self.spritesheet_stale = false;
        return;
    }


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

pub fn getSpriteUVRect(id: Sprite.Id) ?Rect {
    const self = &g_state;

    if (!self.sprites.exists(id)) {
        return null;
    }
    const sprite = self.sprites.getAssumeExists(id);
    if (sprite.data == null) {
        return null;
    }

    if (self.spritesheet_stale) {
        packSprites() catch |e| {
            log.warn("Failed to pack sprites: {}", .{e});
        };
    }

    assert(sprite.position != null);

    const sprite_size = Vec2.initFromInt(u32, sprite.data.?.width, sprite.data.?.height);
    return Rect{
        .pos = sprite.position.?.divide(self.spritesheet_size),
        .size = sprite_size.divide(self.spritesheet_size),
    };
}

pub fn drawSprite(id: Sprite.Id, pos: Vec2, size: Vec2, color: Color) void {
    const self = &g_state;

    var quad: Quad = undefined;
    quad.setColor(color);
    quad.setRect(.{ .pos = pos, .size = size });

    if (getSpriteUVRect(id)) |sprite_uv| {
        quad.setUVRect(sprite_uv);
    } else if (getSpriteUVRect(self.error_sprite)) |error_uv| {
        quad.setUVRect(error_uv);
        quad.setColor(.white);
    }

    draw(.{
        .texture = self.spritesheet_texture,
        .quad = quad
    });
}

pub const ImageData = struct {
    const Pixels = union(enum) {
        none,
        rgba8: [*]u8
    };

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

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

    pub fn initRGBA8(gpa: Allocator, width: u32, height: u32) !ImageData {
        const pixel_count = width * height;
        const pixels = try gpa.alloc(u8, pixel_count * 4);

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

    pub fn clone(self: ImageData, gpa: Allocator) !ImageData {
        const pixel_count = self.width * self.height;
        const pixels = switch (self.pixels) {
            .none => Pixels{ .none = {} },
            .rgba8 => |pixels| Pixels{
                .rgba8 = (try gpa.dupe(u8, pixels[0..(pixel_count * 4)])).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;
        switch (self.pixels) {
            .none => {},
            .rgba8 => |pixels| gpa.free(pixels[0..(pixel_count*4)])
        }
    }

    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 == .none or self.pixels == .none) {
            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_index = source_y * source.width + source_x;
                const source_pixel = source.pixels.rgba8[(4*source_pixel_index)..][0..4];

                const dest_pixel_index = dest_y * self.width + dest_x;
                const dest_pixel = self.pixels.rgba8[(4*dest_pixel_index)..][0..4];

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

pub const Texture = struct {
    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 {
    data: ?ImageData,
    position: ?Vec2,

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

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

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

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

    pub fn setColor(self: *Quad, 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: *Quad, 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: *Quad, 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);
    }
};