const std = @import("std");
const nu = @import("nu.zig");
const Uber = @import("Uber.zig");

pub const nu_modules = .{@This()};
pub const main = nu.main;

pub const nu_config: nu.Config = .{
    .window = nu.Window.Config{
        .title = "au",
    },
    .render = .{
        .app_name = "au",
        .frames_in_flight = 2,
    },
};

const im = nu.ImGui;
const ctx = nu.Render.ctx;
const vk = @import("vk");

pub const depends = .{nu.ImGui};

const vertices: []const Uber.Vertex = &.{
    .{ .pos = .{ 0.0, 0.0, 0.0, 1.0 }, .color = .{ 1.0, 1.0, 0.0 } },
    .{ .pos = .{ 1.0, 0.0, 0.0, 1.0 }, .color = .{ 1.0, 0.0, 1.0 } },
    .{ .pos = .{ 0.0, 1.0, 0.0, 1.0 }, .color = .{ 0.0, 1.0, 1.0 } },
    .{ .pos = .{ 1.0, 1.0, 0.0, 1.0 }, .color = .{ 1.0, 1.0, 1.0 } },
};

const Index = u16;
const indices: []const Index = &.{
    1, 0, 2,
    2, 3, 1,
};

const uniform: Uber.Uniform = .{
    .mat = .{
        1, 0, 0, 0,
        0, 1, 0, 0,
        0, 0, 1, 0,
        0, 0, 0, 1,
    },
};

var mem: vk.DeviceMemory = .null_handle;
var vbo: vk.Buffer = .null_handle;
var ibo: vk.Buffer = .null_handle;
var ubo: vk.Buffer = .null_handle;

var pool: vk.DescriptorPool = .null_handle;
var descriptor_set: vk.DescriptorSet = .null_handle;

var cache: vk.PipelineCache = .null_handle;
var uber: Uber = undefined;

pub fn setup(_: std.mem.Allocator) !void {
    errdefer teardown();
    const props = ctx.I.getPhysicalDeviceMemoryProperties(ctx.pdevice.*);
    const memory_type_index: u32 = for (0..props.memory_type_count) |idx| {
        const t = props.memory_types[idx];
        if (t.property_flags.host_coherent_bit and t.property_flags.host_visible_bit) {
            break @intCast(idx);
        }
    } else {
        unreachable;
    };

    const SIZE = 0x10000;
    mem = try ctx.D.allocateMemory(&vk.MemoryAllocateInfo{
        .allocation_size = SIZE,
        .memory_type_index = memory_type_index,
    }, null);

    const raw: [*]u8 = @ptrCast(try ctx.D.mapMemory(mem, 0, vk.WHOLE_SIZE, .{}) orelse unreachable);
    errdefer ctx.D.unmapMemory(mem);

    // todo VMA. This doesn't work for some reason to do with alignment.
    // var fba = std.heap.FixedBufferAllocator.init(raw[0..SIZE]);
    // const aa = fba.allocator();

    var bump: usize = 0;

    const vbytes = std.mem.sliceAsBytes(vertices);
    vbo = try ctx.D.createBuffer(&vk.BufferCreateInfo{
        .queue_family_index_count = 1,
        .p_queue_family_indices = &.{ctx.family.*},
        .sharing_mode = .exclusive,
        .size = vbytes.len,
        .usage = .{ .vertex_buffer_bit = true },
    }, null);
    const vreq = ctx.D.getBufferMemoryRequirements(vbo);
    bump = std.mem.alignForward(usize, bump, vreq.alignment);
    @memcpy(raw + bump, vbytes);
    try ctx.D.bindBufferMemory(vbo, mem, bump);
    bump += vreq.size;

    const ibytes = std.mem.sliceAsBytes(indices);
    ibo = try ctx.D.createBuffer(&vk.BufferCreateInfo{
        .queue_family_index_count = 1,
        .p_queue_family_indices = &.{ctx.family.*},
        .sharing_mode = .exclusive,
        .size = ibytes.len,
        .usage = .{ .index_buffer_bit = true },
    }, null);
    const ireq = ctx.D.getBufferMemoryRequirements(ibo);
    bump = std.mem.alignForward(usize, bump, ireq.alignment);
    @memcpy(raw + bump, ibytes);
    try ctx.D.bindBufferMemory(ibo, mem, bump);
    bump += ireq.size;

    const ubytes = std.mem.asBytes(&uniform);
    ubo = try ctx.D.createBuffer(&vk.BufferCreateInfo{
        .queue_family_index_count = 1,
        .p_queue_family_indices = &.{ctx.family.*},
        .sharing_mode = .exclusive,
        .size = ubytes.len,
        .usage = .{ .uniform_buffer_bit = true },
    }, null);
    const ureq = ctx.D.getBufferMemoryRequirements(ubo);
    bump = std.mem.alignForward(usize, bump, ureq.alignment);
    @memcpy(raw + bump, ubytes);
    try ctx.D.bindBufferMemory(ubo, mem, bump);
    bump += ureq.size;

    cache = try ctx.D.createPipelineCache(&vk.PipelineCacheCreateInfo{}, null);

    uber = try Uber.init(cache);

    const pool_sizes: []const vk.DescriptorPoolSize = &.{vk.DescriptorPoolSize{
        .descriptor_count = 8,
        .type = .uniform_buffer,
    }};

    pool = try ctx.D.createDescriptorPool(&vk.DescriptorPoolCreateInfo{
        .flags = .{ .free_descriptor_set_bit = true },
        .pool_size_count = @intCast(pool_sizes.len),
        .p_pool_sizes = pool_sizes.ptr,
        .max_sets = 32,
    }, null);

    var sets: [1]vk.DescriptorSet = .{.null_handle};
    try ctx.D.allocateDescriptorSets(&vk.DescriptorSetAllocateInfo{
        .descriptor_pool = pool,
        .descriptor_set_count = 1,
        .p_set_layouts = &.{uber.set_layout},
    }, &sets);
    descriptor_set = sets[0];

    ctx.D.updateDescriptorSets(
        1,
        &.{
            vk.WriteDescriptorSet{
                .dst_set = descriptor_set,
                .dst_binding = 0,
                .dst_array_element = 0,
                .descriptor_count = 1,
                .descriptor_type = .uniform_buffer,
                .p_buffer_info = &.{
                    vk.DescriptorBufferInfo{
                        .buffer = ubo,
                        .offset = 0,
                        .range = vk.WHOLE_SIZE,
                    },
                },
                .p_image_info = undefined,
                .p_texel_buffer_view = undefined,
            },
        },
        0,
        null,
    );
}

pub fn teardown() void {
    ctx.Q.waitIdle() catch {};
    uber.deinit();
    ctx.D.destroyPipelineCache(cache, null);
    ctx.D.destroyBuffer(ubo, null);
    ctx.D.destroyBuffer(ibo, null);
    ctx.D.destroyBuffer(vbo, null);
    ctx.D.freeMemory(mem, null);
    ctx.D.freeDescriptorSets(pool, 1, &.{descriptor_set}) catch unreachable;
    ctx.D.destroyDescriptorPool(pool, null);
}

pub fn frame() !void {
    im.igShowMetricsWindow(null);
}

pub fn present(cmd: ctx.CommandBufferProxy) void {
    const w, const h = nu.Window.size();

    cmd.bindPipeline(.graphics, uber.pipeline);
    cmd.setScissor(0, 1, &.{vk.Rect2D{
        .offset = .{ .x = 0, .y = 0 },
        .extent = .{ .width = w, .height = h },
    }});
    cmd.setViewport(0, 1, &.{vk.Viewport{
        .x = 0,
        .y = 0,
        .width = @floatFromInt(w),
        .height = @floatFromInt(h),
        .min_depth = 0,
        .max_depth = 1,
    }});
    cmd.bindIndexBuffer(ibo, 0, .uint16);
    cmd.bindVertexBuffers(0, 1, &.{vbo}, &.{0});
    cmd.bindDescriptorSets(
        .graphics,
        uber.layout,
        0,
        1,
        &.{descriptor_set},
        0,
        null,
    );
    cmd.drawIndexed(@intCast(indices.len), 1, 0, 0, 0);
}
