const Gui = @This();
const std = @import("std");

const ctx = @import("../AppContext.zig");
const shaders = @import("../shaders.zig");
const vk = @import("vulkan");
const vm = @import("vecmath");

const CommandBuffer = @import("CommandBuffer.zig");
const Engine = @import("Engine.zig");
const GenericBuffer = @import("GenericBuffer.zig").GenericBuffer;
const Swapchain = @import("Swapchain.zig");
const Texture = @import("Texture.zig");
const Textures = @import("Textures.zig");

pub const Draw = struct {
    pub const Box = extern struct {
        background_color: vm.Vector4,
        border_color: vm.Vector4,
        position_sspx: vm.Vector2,
        size_px: vm.Vector2,
        border_width_px: f32,
        border_radius_px: f32,
    };

    pub const Text = extern struct {
        color: vm.Vector4,
        pos_min: vm.Vector2,
        pos_max: vm.Vector2,
        uv_min: vm.Vector2,
        uv_max: vm.Vector2,
        atlas_id: Textures.Id,
    };

    pub const Image = extern struct {
        tint: vm.Vector4,
        position_sspx: vm.Vector2,
        size_px: vm.Vector2,
        uv_min: vm.Vector2,
        uv_max: vm.Vector2,
        texture_id: Textures.Id,
    };
};

pub const Batch = struct {
    draw_type: std.meta.DeclEnum(Draw),
    first_instance: u32,
    instance_count: u32,
};

box_gpu_buffer: GenericBuffer(void, Draw.Box),
text_gpu_buffer: GenericBuffer(void, Draw.Text),
image_gpu_buffer: GenericBuffer(void, Draw.Image),

box_cpu_buffer: std.ArrayList(Draw.Box),
text_cpu_buffer: std.ArrayList(Draw.Text),
image_cpu_buffer: std.ArrayList(Draw.Image),

batches: std.ArrayList(Batch),

sampler: vk.Sampler,
index_buffer: shaders.IndexBuffer,

box_descriptor_set_layout: vk.DescriptorSetLayout,
box_pipeline_layout: vk.PipelineLayout,
box_pipeline: vk.Pipeline,

descriptor_pool: vk.DescriptorPool,
box_descriptor_set: vk.DescriptorSet,

pub const max_box_draws = 1024;
pub const max_image_draws = 1024;
pub const max_text_draws = 4096;

pub const max_batches = 1024;

pub fn init() !*Gui {
    const allocator_persistent = ctx.allocator_persistent;
    const engine = ctx.engine;
    const swapchain = ctx.swapchain;

    const gui = try allocator_persistent.create(Gui);

    var box_gpu_buffer: GenericBuffer(void, Draw.Box) = try .init(.{
        .usage = .storage,
        .target_queue = .graphics,
        .array_capacity = max_box_draws,
        .name = "GUI box draws",
    });
    errdefer box_gpu_buffer.deinit();

    var text_gpu_buffer: GenericBuffer(void, Draw.Text) = try .init(.{
        .usage = .storage,
        .target_queue = .graphics,
        .array_capacity = max_text_draws,
        .name = "GUI text draws",
    });
    errdefer text_gpu_buffer.deinit();

    var image_gpu_buffer: GenericBuffer(void, Draw.Image) = try .init(.{
        .usage = .storage,
        .target_queue = .graphics,
        .array_capacity = max_image_draws,
        .name = "GUI image draws",
    });
    errdefer image_gpu_buffer.deinit();

    const box_cpu_buffer: std.ArrayList(Draw.Box) = try .initCapacity(allocator_persistent, max_box_draws);
    const text_cpu_buffer: std.ArrayList(Draw.Text) = try .initCapacity(allocator_persistent, max_text_draws);
    const image_cpu_buffer: std.ArrayList(Draw.Image) = try .initCapacity(allocator_persistent, max_image_draws);
    const batches: std.ArrayList(Batch) = try .initCapacity(allocator_persistent, max_batches);

    const sampler = try engine.createSampler(.{
        .mag_filter = .linear,
        .min_filter = .linear,
        .mipmap_mode = .linear,
        .address_mode_u = .repeat,
        .address_mode_v = .repeat,
        .address_mode_w = .repeat,
        .mip_lod_bias = 0,
        .anisotropy_enable = .false,
        .max_anisotropy = 0,
        .compare_enable = .false,
        .compare_op = .always,
        .min_lod = 0,
        .max_lod = vk.LOD_CLAMP_NONE,
        .border_color = .float_transparent_black,
        .unnormalized_coordinates = .false,
    });
    errdefer engine.destroySampler(sampler);

    const box_descriptor_set_layout = try engine.createDescriptorSetLayout(.{
        .bindings = &.{
            .{
                .binding = 0,
                .descriptor_type = .uniform_buffer,
                .descriptor_count = 1,
                .stage_flags = .{ .vertex_bit = true, .fragment_bit = true },
            },
            .{
                .binding = 1,
                .descriptor_type = .storage_buffer,
                .descriptor_count = 1,
                .stage_flags = .{ .vertex_bit = true, .fragment_bit = true },
            },
        },
    });
    errdefer engine.destroyDescriptorSetLayout(box_descriptor_set_layout);
    engine.setObjectName(box_descriptor_set_layout, "DSL GUI Box", .{});

    const box_pipeline_layout = try engine.createPipelineLayout(.{
        .set_layouts = &.{
            box_descriptor_set_layout,
        },
    });
    errdefer engine.destroyPipelineLayout(box_pipeline_layout);
    engine.setObjectName(box_pipeline_layout, "PL GUI Box", .{});

    const box_vertex_shader = try engine.createShaderModule(.{ .code = &shaders.gui_box_vert_spv });
    defer engine.destroyShaderModule(box_vertex_shader);
    engine.setObjectName(box_vertex_shader, "SM gui_box_vert", .{});

    const box_fragment_shader = try engine.createShaderModule(.{ .code = &shaders.gui_box_frag_spv });
    defer engine.destroyShaderModule(box_fragment_shader);
    engine.setObjectName(box_fragment_shader, "SM gui_box_frag", .{});

    var index_buffer = try shaders.IndexBuffer.init(.{
        .usage = .index,
        .target_queue = .graphics,
        .array_capacity = 6,
        .name = "QuadIB",
    });
    errdefer index_buffer.deinit();
    try index_buffer.write(.{
        .elements = &.{ 0, 1, 2, 2, 1, 3 },
    });

    const box_pipeline = try engine.createGraphicsPipeline(.{
        .stages = &.{
            .{
                .stage = .{ .vertex_bit = true },
                .module = box_vertex_shader,
                .name = "main",
            },
            .{
                .stage = .{ .fragment_bit = true },
                .module = box_fragment_shader,
                .name = "main",
            },
        },
        .vertex_input_state = .{},
        .input_assembly_state = .{
            .topology = .triangle_list,
            .primitive_restart_enable = .false,
        },
        .viewport_state = .{
            .viewports = &.{undefined}, // dynamic
            .scissors = &.{undefined}, // dynamic
        },
        .rasterization_state = .{
            .depth_clamp_enable = .false,
            .rasterizer_discard_enable = .false,
            .polygon_mode = .fill,
            .cull_mode = .{},
            .front_face = .counter_clockwise,
            .depth_bias_enable = .false,
            .depth_bias_constant_factor = 0,
            .depth_bias_clamp = 0,
            .depth_bias_slope_factor = 0,
            .line_width = 1,
        },
        .multisample_state = .{
            .rasterization_samples = .{ .@"1_bit" = true },
            .sample_shading_enable = .false,
            .min_sample_shading = 1,
            .alpha_to_coverage_enable = .false,
            .alpha_to_one_enable = .false,
        },
        .depth_stencil_state = .{
            .depth_test_enable = .false,
            .depth_write_enable = .false,
            .depth_compare_op = .always,
            .depth_bounds_test_enable = .false,
            .stencil_test_enable = .false,
            .front = .{
                .fail_op = .keep,
                .pass_op = .keep,
                .depth_fail_op = .keep,
                .compare_op = .never,
                .compare_mask = 0,
                .write_mask = 0,
                .reference = 0,
            },
            .back = .{
                .fail_op = .keep,
                .pass_op = .keep,
                .depth_fail_op = .keep,
                .compare_op = .never,
                .compare_mask = 0,
                .write_mask = 0,
                .reference = 0,
            },
            .min_depth_bounds = 0,
            .max_depth_bounds = 1,
        },
        .color_blend_state = .{
            .logic_op_enable = .false,
            .logic_op = .copy,
            .attachments = &.{
                .{
                    .blend_enable = .true,
                    .src_color_blend_factor = .one,
                    .dst_color_blend_factor = .one_minus_src_alpha,
                    .color_blend_op = .add,
                    .src_alpha_blend_factor = .one,
                    .dst_alpha_blend_factor = .one_minus_src_alpha,
                    .alpha_blend_op = .add,
                    .color_write_mask = .{
                        .r_bit = true,
                        .g_bit = true,
                        .b_bit = true,
                        .a_bit = true,
                    },
                },
            },
            .blend_constants = .{ 0, 0, 0, 0 },
        },
        .dynamic_state = .{
            .dynamic_states = &.{ .viewport, .scissor },
        },
        .layout = box_pipeline_layout,
        .view_mask = 0,
        .color_attachment_formats = &.{swapchain.params.surface_format.format},
        .depth_attachment_format = Texture.Usage.depth.vkFormat(),
        .stencil_attachment_format = .undefined,
    });
    errdefer engine.destroyPipeline(box_pipeline);
    engine.setObjectName(box_pipeline, "P GUI Box", .{});

    const descriptor_pool = try engine.createDescriptorPool(.{
        .max_sets = 1,
        .pool_sizes = &.{
            .{
                .type = .uniform_buffer,
                .descriptor_count = 1,
            },
            .{
                .type = .storage_buffer,
                .descriptor_count = 1,
            },
        },
    });
    errdefer engine.destroyDescriptorPool(descriptor_pool);
    engine.setObjectName(descriptor_pool, "DP GUI", .{});

    const box_descriptor_set = try engine.allocateDescriptorSet(.{
        .descriptor_pool = descriptor_pool,
        .set_layout = box_descriptor_set_layout,
    });
    engine.setObjectName(box_descriptor_set, "DS GUI Box", .{});

    try engine.updateDescriptorSets(.{
        .writes = &.{
            .{
                .dst_set = box_descriptor_set,
                .dst_binding = 1,
                .dst_array_element = 0,
                .descriptor_type = .storage_buffer,
                .descriptor_infos = .{
                    .buffer = &.{
                        .{
                            .buffer = box_gpu_buffer.buffer,
                            .offset = 0,
                            .range = vk.WHOLE_SIZE,
                        },
                    },
                },
            },
        },
    });

    gui.* = .{
        .box_gpu_buffer = box_gpu_buffer,
        .text_gpu_buffer = text_gpu_buffer,
        .image_gpu_buffer = image_gpu_buffer,

        .box_cpu_buffer = box_cpu_buffer,
        .text_cpu_buffer = text_cpu_buffer,
        .image_cpu_buffer = image_cpu_buffer,

        .batches = batches,

        .sampler = sampler,
        .index_buffer = index_buffer,

        .box_descriptor_set_layout = box_descriptor_set_layout,
        .box_pipeline_layout = box_pipeline_layout,
        .box_pipeline = box_pipeline,

        .descriptor_pool = descriptor_pool,
        .box_descriptor_set = box_descriptor_set,
    };

    return gui;
}

pub fn deinit(self: *Gui) void {
    const engine = ctx.engine;

    std.log.scoped(.deinit).debug("Deinitializing {*}", .{self});

    engine.destroyDescriptorPool(self.descriptor_pool);
    engine.destroyPipeline(self.box_pipeline);
    self.index_buffer.deinit();
    engine.destroyPipelineLayout(self.box_pipeline_layout);
    engine.destroyDescriptorSetLayout(self.box_descriptor_set_layout);
    engine.destroySampler(self.sampler);

    self.image_gpu_buffer.deinit();
    self.text_gpu_buffer.deinit();
    self.box_gpu_buffer.deinit();

    self.* = undefined;
}

pub fn beginFrame(self: *Gui) void {
    self.box_cpu_buffer.clearRetainingCapacity();
    self.text_cpu_buffer.clearRetainingCapacity();
    self.image_cpu_buffer.clearRetainingCapacity();

    self.batches.clearRetainingCapacity();
}

pub fn pushBox(self: *Gui, box: Draw.Box) void {
    const instance: u32 = @intCast(self.box_cpu_buffer.items.len);
    self.box_cpu_buffer.appendBounded(box) catch return;
    self.extendOrAddBatch(.Box, instance);
}

pub fn pushText(self: *Gui, text: Draw.Text) void {
    const instance: u32 = @intCast(self.text_cpu_buffer.items.len);
    self.text_cpu_buffer.appendBounded(text) catch return;
    self.extendOrAddBatch(.Text, instance);
}

pub fn pushImage(self: *Gui, image: Draw.Image) void {
    const instance: u32 = @intCast(self.image_cpu_buffer.items.len);
    self.image_cpu_buffer.appendBounded(image) catch return;
    self.extendOrAddBatch(.Image, instance);
}

fn extendOrAddBatch(self: *Gui, draw_type: std.meta.DeclEnum(Draw), instance: u32) void {
    if (self.batches.items.len == 0 or self.batches.items[self.batches.items.len - 1].draw_type != draw_type) {
        self.batches.appendBounded(.{
            .draw_type = draw_type,
            .first_instance = instance,
            .instance_count = 1,
        }) catch return;
    } else {
        const batch = &self.batches.items[self.batches.items.len - 1];
        std.debug.assert(batch.first_instance + batch.instance_count == instance);
        batch.instance_count += 1;
    }
}

pub fn draw(self: *const Gui, command_buffer: CommandBuffer) !void {
    command_buffer.bindIndexBuffer(self.index_buffer.buffer, 0, .uint16);

    try self.box_gpu_buffer.write(.{ .elements = self.box_cpu_buffer.items });
    try self.text_gpu_buffer.write(.{ .elements = self.text_cpu_buffer.items });
    try self.image_gpu_buffer.write(.{ .elements = self.image_cpu_buffer.items });

    for (self.batches.items) |batch| {
        switch (batch.draw_type) {
            .Box => {
                command_buffer.bindPipeline(.graphics, self.box_pipeline);
                command_buffer.bindDescriptorSet(.graphics, self.box_pipeline_layout, 0, self.box_descriptor_set, null);
            },
            .Text => {
                std.log.warn("GUI Text not implemented; skipping batch", .{});
                continue;
            },
            .Image => {
                std.log.warn("GUI Image not implemented; skipping batch", .{});
                continue;
            },
        }

        command_buffer.drawIndexed(.{
            .index_count = 6,
            .first_instance = batch.first_instance,
            .instance_count = batch.instance_count,
        });
    }
}