const sokol = @import("sokol");
const sg = sokol.gfx;
const sglue = sokol.glue;
const slog = sokol.log;
const sapp = sokol.app;
const simgui = sokol.imgui;
const sgl = sokol.gl;

const Math = @import("./math.zig");
const Line = Math.Line;
const Vec2 = Math.Vec2;
const Vec4 = Math.Vec4;
const rgb = Math.rgb;
const Rect = Math.Rect;

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

const imgui = @import("imgui.zig");
const tracy = @import("tracy");
const fontstash = @import("./fontstash/root.zig");
pub const Font = fontstash.Font;

const GraphicsFrame = @import("./frame.zig").Graphics;

// TODO: Seems that there is a vertical jitter bug when resizing a window in OpenGL. Seems like a driver bug.
// From other peoples research it seems that disabling vsync when a resize event occurs fixes it.
// Maybe a patch for sokol could be made?
// More info:
// * https://github.com/libsdl-org/SDL/issues/11618
// * https://github.com/nimgl/nimgl/issues/59

const Options = struct {
    const ImguiFont = struct { ttf_data: []const u8, size: f32 = 16 };

    allocator: std.mem.Allocator,
    logger: sg.Logger = .{},
    imgui_font: ?ImguiFont = null,
};

pub const Command = union(enum) {
    pub const DrawRectangle = struct {
        rect: Rect,
        color: Vec4,
        sprite: ?Sprite = null,
        rotation: f32 = 0,
        origin: Vec2 = .init(0, 0),
        uv_flip_diagonal: bool = false,
        uv_flip_horizontal: bool = false,
        uv_flip_vertical: bool = false
    };

    pub const DrawCircle = struct {
        center: Vec2,
        radius: f32,
        color: Vec4,
    };

    set_scissor: Rect,
    draw_rectangle: DrawRectangle,
    draw_line: struct { pos1: Vec2, pos2: Vec2, color: Vec4, width: f32 },
    draw_text: struct {
        pos: Vec2,
        text: []const u8,
        font: Font.Id,
        size: f32,
        color: Vec4,
    },
    push_transformation: struct { translation: Vec2, scale: Vec2 },
    pop_transformation: void,
    draw_circle: DrawCircle
};

const Texture = struct {
    image: sg.Image,
    view: sg.View,
    info: Info,

    const Info = struct {
        width: u32,
        height: u32,
    };

    const Id = enum(u32) { _ };
    const Data = struct { width: u32, height: u32, rgba: [*]u8 };
};
pub const TextureId = Texture.Id;
pub const TextureInfo = Texture.Info;

pub const Sprite = struct {
    texture: TextureId,
    uv: Rect,

    pub fn flipHorizontal(self: Sprite) Sprite {
        var uv = self.uv;
        uv.pos.x += uv.size.x;
        uv.size.x *= -1;

        return Sprite{
            .texture = self.texture,
            .uv = uv
        };
    }

    pub fn getSize(self: Sprite) Vec2 {
        const texture_info = getTextureInfo(self.texture);
        const texture_size = Vec2.initFromInt(u32, texture_info.width, texture_info.height);
        return self.uv.size.multiply(texture_size);
    }
};

var gpa: std.mem.Allocator = undefined;

var main_pipeline: sgl.Pipeline = .{};
var linear_sampler: sg.Sampler = .{};
var nearest_sampler: sg.Sampler = .{};
var font_context: fontstash.Context = undefined;
var textures: std.ArrayList(Texture) = .empty;

var scale_stack_buffer: [32]Vec2 = undefined;
var scale_stack: std.ArrayList(Vec2) = .empty;

pub fn init(options: Options) !void {
    gpa = options.allocator;

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

    sgl.setup(.{ .logger = .{ .func = options.logger.func, .user_data = options.logger.user_data } });

    main_pipeline = sgl.makePipeline(.{
        .colors = init: {
            var colors: [8]sg.ColorTargetState = @splat(.{});
            colors[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 colors;
        },
    });

    imgui.setup(options.allocator, .{
        .logger = .{ .func = options.logger.func, .user_data = options.logger.user_data },
        .no_default_font = options.imgui_font != null,

        // TODO: Figure out a way to make imgui play nicely with UI
        // Ideally when mouse is inside a Imgui window, then the imgui cursor should be used.
        // Otherwise our own cursor should be used.
        .disable_set_mouse_cursor = true,
    });

    if (options.imgui_font) |imgui_font| {
        imgui.addFont(imgui_font.ttf_data, imgui_font.size);
    }

    linear_sampler = sg.makeSampler(.{
        .min_filter = .LINEAR,
        .mag_filter = .LINEAR,
        .mipmap_filter = .LINEAR,
        .label = "linear-sampler",
    });

    nearest_sampler = sg.makeSampler(.{
        .min_filter = .NEAREST,
        .mag_filter = .NEAREST,
        .mipmap_filter = .NEAREST,
        .label = "nearest-sampler",
    });

    const dpi_scale = sapp.dpiScale();
    const atlas_size = 512;
    const atlas_dim = std.math.ceilPowerOfTwoAssert(u32, @intFromFloat(atlas_size * dpi_scale));
    font_context = try fontstash.Context.init(.{
        .width = @intCast(atlas_dim),
        .height = @intCast(atlas_dim),
    });
}

pub fn deinit() void {
    textures.deinit(gpa);
    imgui.shutdown();
    font_context.deinit();
    sgl.shutdown();
    sg.shutdown();
}

pub fn drawCommand(command: Command) void {
    switch (command) {
        .push_transformation => |opts| {
            pushTransform(opts.translation, opts.scale);
            // font_resolution_scale = font_resolution_scale.multiply(opts.scale);
        },
        .pop_transformation => {
            popTransform();
        },
        .draw_rectangle => |opts| {
            drawRectangle(opts);
        },
        .set_scissor => |opts| {
            sgl.scissorRectf(opts.pos.x, opts.pos.y, opts.size.x, opts.size.y, true);
        },
        .draw_line => |opts| {
            drawLine(opts.pos1, opts.pos2, opts.color, opts.width);
        },
        .draw_text => |opts| {
            const font_resolution_scale = scale_stack.getLast();

            sgl.pushMatrix();
            defer sgl.popMatrix();

            sgl.scale(1 / font_resolution_scale.x, 1 / font_resolution_scale.y, 1);

            font_context.setFont(opts.font);
            font_context.setSize(opts.size * font_resolution_scale.y);
            font_context.setAlign(.{ .x = .left, .y = .top });
            font_context.setSpacing(0);

            const r: u8 = @intFromFloat(opts.color.x * 255);
            const g: u8 = @intFromFloat(opts.color.y * 255);
            const b: u8 = @intFromFloat(opts.color.z * 255);
            const a: u8 = @intFromFloat(opts.color.w * 255);
            const color: u32 = r | (@as(u32, g) << 8) | (@as(u32, b) << 16) | (@as(u32, a) << 24);
            font_context.setColor(color);
            font_context.drawText(opts.pos.x * font_resolution_scale.x, opts.pos.y * font_resolution_scale.y, opts.text);
        },
        .draw_circle => |opts| {
            sgl.beginTriangleStrip();
            defer sgl.end();

            const angle_step: f32 = std.math.pi * 2.0 / 32.0;

            sgl.c4f(opts.color.x, opts.color.y, opts.color.z, opts.color.w);

            sgl.v2f(opts.center.x, opts.center.y);
            var angle: f32 = 0;
            while (angle < std.math.pi * 2) {
                const point_on_circle = Vec2.initAngle(angle).multiplyScalar(opts.radius).add(opts.center);
                sgl.v2f(point_on_circle.x, point_on_circle.y);
                sgl.v2f(opts.center.x, opts.center.y);

                angle += angle_step;
            }
            sgl.v2f(opts.center.x + opts.radius, opts.center.y);
        }
    }
}

pub fn drawCommands(commands: []const Command) void {
    for (commands) |command| {
        drawCommand(command);
    }
}

pub fn beginFrame() void {
    const zone = tracy.initZone(@src(), .{});
    defer zone.deinit();

    imgui.newFrame(.{ .width = sapp.width(), .height = sapp.height(), .delta_time = sapp.frameDuration(), .dpi_scale = sapp.dpiScale() });

    scale_stack = .initBuffer(&scale_stack_buffer);
    scale_stack.appendAssumeCapacity(.init(1, 1));

    font_context.clearState();
    sgl.defaults();
    sgl.matrixModeProjection();
    sgl.ortho(0, sapp.widthf(), sapp.heightf(), 0, -1, 1);
    sgl.loadPipeline(main_pipeline);
}

pub fn endFrame(clear_color: Vec4) void {
    const zone = tracy.initZone(@src(), .{});
    defer zone.deinit();

    var pass_action: sg.PassAction = .{};

    pass_action.colors[0] = sg.ColorAttachmentAction{ .load_action = .CLEAR, .clear_value = .{ .r = clear_color.x, .g = clear_color.y, .b = clear_color.z, .a = clear_color.w } };

    font_context.flush();

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

        sgl.draw();

        imgui.render();
    }
    sg.commit();
}

fn pushTransform(translation: Vec2, scale: Vec2) void {
    sgl.pushMatrix();
    sgl.translate(translation.x, translation.y, 0);
    sgl.scale(scale.x, scale.y, 1);

    scale_stack.appendAssumeCapacity(scale_stack.getLast().multiply(scale));
}

fn popTransform() void {
    sgl.popMatrix();
    _ = scale_stack.pop().?;
}

const Vertex = struct { pos: Vec2, uv: Vec2 };

fn drawQuad(quad: [4]Vertex, color: Vec4, texture_id: TextureId) void {
    sgl.enableTexture();
    defer sgl.disableTexture();

    const view = textures.items[@intFromEnum(texture_id)].view;
    // TODO: Make sampler configurable
    sgl.texture(view, nearest_sampler);

    sgl.beginQuads();
    defer sgl.end();

    sgl.c4f(color.x, color.y, color.z, color.w);
    for (quad) |vertex| {
        const pos = vertex.pos;
        const uv = vertex.uv;
        sgl.v2fT2f(pos.x, pos.y, uv.x, uv.y);
    }
}

fn drawQuadNoUVs(quad: [4]Vec2, color: Vec4) void {
    sgl.beginQuads();
    defer sgl.end();

    sgl.c4f(color.x, color.y, color.z, color.w);
    for (quad) |pos| {
        sgl.v2f(pos.x, pos.y);
    }
}

fn drawRectangle(opts: Command.DrawRectangle) void {
    const pos = opts.rect.pos;
    const size = opts.rect.size;

    const top_left = Vec2.init(0, 0).rotateAround(opts.rotation, opts.origin);
    const top_right = Vec2.init(size.x, 0).rotateAround(opts.rotation, opts.origin);
    const bottom_right = size.rotateAround(opts.rotation, opts.origin);
    const bottom_left = Vec2.init(0, size.y).rotateAround(opts.rotation, opts.origin);

    if (opts.sprite) |sprite| {
        const uv = sprite.uv;
        var uv_top_left = Vec2.init(uv.left(), uv.top());
        var uv_top_right = Vec2.init(uv.right(), uv.top());
        var uv_bottom_right = Vec2.init(uv.right(), uv.bottom());
        var uv_bottom_left = Vec2.init(uv.left(), uv.bottom());

        if (opts.uv_flip_diagonal) {
            std.mem.swap(Vec2, &uv_bottom_left, &uv_top_right);
        }
        if (opts.uv_flip_horizontal) {
            std.mem.swap(Vec2, &uv_top_left, &uv_top_right);
            std.mem.swap(Vec2, &uv_bottom_left, &uv_bottom_right);
        }
        if (opts.uv_flip_vertical) {
            std.mem.swap(Vec2, &uv_top_left, &uv_bottom_left);
            std.mem.swap(Vec2, &uv_top_right, &uv_bottom_right);
        }

        const quad = [4]Vertex{
            .{
                .pos = pos.add(top_left),
                .uv = uv_top_left
            },
            .{
                .pos = pos.add(top_right),
                .uv = uv_top_right,
            },
            .{
                .pos = pos.add(bottom_right),
                .uv = uv_bottom_right
            },
            .{
                .pos = pos.add(bottom_left),
                .uv = uv_bottom_left
            }
        };
        drawQuad(quad, opts.color, sprite.texture);
    } else {
        const quad = .{ pos.add(top_left), pos.add(top_right), pos.add(bottom_right), pos.add(bottom_left) };
        drawQuadNoUVs(quad, opts.color);
    }
}

fn drawLine(from: Vec2, to: Vec2, color: Vec4, width: f32) void {
    const line = Line{
        .p0 = from,
        .p1 = to,
    };

    const quad = line.getQuad(width);

    drawQuadNoUVs(
        .{
            quad.top_right,
            quad.top_left,
            quad.bottom_left,
            quad.bottom_right
        },
        color
    );
}

pub fn addFont(name: [*c]const u8, data: []const u8) !Font.Id {
    return try font_context.addFont(name, data);
}

fn makeView(image: sg.Image) !sg.View {
    const image_view = sg.makeView(.{ .texture = .{ .image = image } });
    if (image_view.id == sg.invalid_id) {
        return error.InvalidView;
    }
    return image_view;
}

fn makeImageWithMipMaps(mipmaps: []const Texture.Data) !sg.Image {
    if (mipmaps.len == 0) {
        return error.NoMipMaps;
    }

    var data: sg.ImageData = .{};
    var mip_levels: std.ArrayListUnmanaged(sg.Range) = .initBuffer(&data.mip_levels);

    for (mipmaps) |mipmap| {
        try mip_levels.appendBounded(.{ .ptr = mipmap.rgba, .size = mipmap.width * mipmap.height * 4 });
    }

    const image = sg.makeImage(.{ .width = @intCast(mipmaps[0].width), .height = @intCast(mipmaps[0].height), .pixel_format = .RGBA8, .usage = .{ .immutable = true }, .num_mipmaps = @intCast(mip_levels.items.len), .data = data });
    if (image.id == sg.invalid_id) {
        return error.InvalidImage;
    }

    return image;
}

pub fn addTexture(mipmaps: []const Texture.Data) !TextureId {
    const image = try makeImageWithMipMaps(mipmaps);
    errdefer sg.deallocImage(image);

    const view = try makeView(image);
    errdefer sg.deallocView(view);

    assert(mipmaps.len > 0);
    const index = textures.items.len;
    try textures.append(gpa, .{ .image = image, .view = view, .info = .{
        .width = mipmaps[0].width,
        .height = mipmaps[0].height,
    } });

    return @enumFromInt(index);
}

pub fn getTextureInfo(id: TextureId) TextureInfo {
    const texture = textures.items[@intFromEnum(id)];
    return texture.info;
}

pub fn getTextureSize(id: TextureId) Vec2 {
    const texture_info = getTextureInfo(id);
    return Vec2.initFromInt(u32, texture_info.width, texture_info.height);
}