const std = @import("std");
const vk = @import("vk");
const c = @import("c.zig");
const shaders = @import("shaders");
const Allocator = std.mem.Allocator;

const au = @import("au.zig");
const im = @import("cimgui");

const Vertex = extern struct {
    const binding_description = vk.VertexInputBindingDescription{
        .binding = 0,
        .stride = @sizeOf(Vertex),
        .input_rate = .vertex,
    };

    const attribute_description = [_]vk.VertexInputAttributeDescription{
        .{
            .binding = 0,
            .location = 0,
            .format = .r32g32b32a32_sfloat,
            .offset = @offsetOf(Vertex, "pos"),
        },
        .{
            .binding = 0,
            .location = 1,
            .format = .r32g32b32_sfloat,
            .offset = @offsetOf(Vertex, "color"),
        },
    };

    pos: [4]f32,
    color: [3]f32,
};

const Index = u16;

const vertices = [_]Vertex{
    // Vulkan depth range is 0, 1 instead of OpenGL -1, 1
    .{ .pos = .{ -0.5, -0.5, -0.5, 1.0 }, .color = .{ 1, 0, 0 } },
    .{ .pos = .{ -0.5, 0.5, -0.5, 1.0 }, .color = .{ 0, 1, 0 } },
    .{ .pos = .{ 0.5, -0.5, -0.5, 1.0 }, .color = .{ 0, 0, 1 } },
    .{ .pos = .{ 0.5, 0.5, -0.5, 1.0 }, .color = .{ 1, 1, 0 } },
    .{ .pos = .{ -0.5, -0.5, 0.5, 1.0 }, .color = .{ 1, 0, 0 } },
    .{ .pos = .{ -0.5, 0.5, 0.5, 1.0 }, .color = .{ 0, 1, 0 } },
    .{ .pos = .{ 0.5, -0.5, 0.5, 1.0 }, .color = .{ 0, 0, 1 } },
    .{ .pos = .{ 0.5, 0.5, 0.5, 1.0 }, .color = .{ 1, 1, 0 } },
};

const indices = [_]Index{ 4, 5, 6, 6, 5, 7 };

const Frame = struct {
    pub fn init() !Frame {
        return .{};
    }

    pub fn deinit(self: Frame) void {
        _ = self;
    }

    pub fn record_render(
        self: Frame,
        cmd: au.CommandBufferProxy,
        image: vk.Image,
        view: vk.ImageView,
        scissor: vk.Rect2D,
    ) !void {
        _ = self;

        cmd.pipelineBarrier(
            .{ .top_of_pipe_bit = true },
            .{ .color_attachment_output_bit = true },
            .{},
            0,
            null,
            0,
            null,
            1,
            @ptrCast(&vk.ImageMemoryBarrier{
                .src_access_mask = .{},
                .dst_access_mask = .{ .color_attachment_write_bit = true },
                .old_layout = .undefined,
                .new_layout = .color_attachment_optimal,
                .src_queue_family_index = 0,
                .dst_queue_family_index = 0,
                .image = image,
                .subresource_range = .{
                    .aspect_mask = .{ .color_bit = true },
                    .base_mip_level = 0,
                    .level_count = 1,
                    .base_array_layer = 0,
                    .layer_count = 1,
                },
            }),
        );

        const info = vk.RenderingInfoKHR{
            .render_area = scissor,
            .layer_count = 1,
            .view_mask = 0,
            .color_attachment_count = 1,
            .p_color_attachments = &.{vk.RenderingAttachmentInfo{
                .image_view = view,
                .image_layout = .color_attachment_optimal,
                .resolve_mode = .{},
                .resolve_image_view = .null_handle,
                .resolve_image_layout = .undefined,
                .load_op = .clear,
                .store_op = .store,
                .clear_value = .{ .color = .{ .float_32 = .{ 1, 0, 0, 1 } } },
            }},
        };

        cmd.setViewport(0, 1, &.{.{
            .x = @floatFromInt(scissor.offset.x),
            .y = @floatFromInt(scissor.offset.y),
            .width = @floatFromInt(scissor.extent.width),
            .height = @floatFromInt(scissor.extent.height),
            .min_depth = 0,
            .max_depth = 1,
        }});
        cmd.setScissor(0, 1, &.{scissor});

        cmd.beginRendering(&info);

        // todo
        // vkd.cmdBindPipeline(cmdbuf, .graphics, pipeline);
        // const offset = [_]vk.DeviceSize{0};
        // vkd.cmdBindVertexBuffers(cmdbuf, 0, 1, @ptrCast(&vertex_buffer), &offset);
        // vkd.cmdBindIndexBuffer(cmdbuf, index_buffer, 0, .uint16);
        // vkd.cmdDrawIndexed(cmdbuf, indices.len, 1, 0, 0, 0);

        im.c.ImGui_ImplVulkan_RenderDrawData(im.c.igGetDrawData(), @ptrFromInt(@intFromEnum(cmd.handle)), null);

        cmd.endRendering();

        cmd.pipelineBarrier(
            .{ .color_attachment_output_bit = true },
            .{ .bottom_of_pipe_bit = true },
            .{},
            0,
            null,
            0,
            null,
            1,
            @ptrCast(&vk.ImageMemoryBarrier{
                .src_access_mask = .{ .color_attachment_write_bit = true },
                .dst_access_mask = .{},
                .old_layout = .color_attachment_optimal,
                .new_layout = .present_src_khr,
                .src_queue_family_index = 0,
                .dst_queue_family_index = 0,
                .image = image,
                .subresource_range = .{
                    .aspect_mask = .{ .color_bit = true },
                    .base_mip_level = 0,
                    .level_count = 1,
                    .base_array_layer = 0,
                    .layer_count = 1,
                },
            }),
        );
    }
};

pub fn loader_wrapper(procname: [*c]const u8, _: ?*anyopaque) callconv(.C) vk.PfnVoidFunction {
    return c.glfwGetInstanceProcAddress(au.I.handle, procname);
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.detectLeaks();
    const alloc = gpa.allocator();

    try au.init(alloc);
    defer au.deinit();

    var sc = try au.SwapChain.init(alloc);
    defer sc.deinit();

    var flights = try au.Flights(Frame).init(alloc, 3); // FRAMES IN FLIGHT
    defer flights.deinit();

    const ctx = im.c.igCreateContext(null) orelse return error.igCreateContextFailed;
    defer im.c.igDestroyContext(ctx);

    // _ = im.c.ImGui_ImplGlfw_InitForOther(@ptrCast(au.W.handle), true);
    // defer im.c.ImGui_ImplGlfw_Shutdown();

    const descriptorPool = try au.D.createDescriptorPool(&vk.DescriptorPoolCreateInfo{
        .flags = .{ .free_descriptor_set_bit = true },
        .pool_size_count = 1,
        .p_pool_sizes = &.{
            vk.DescriptorPoolSize{ .descriptor_count = 32, .type = .combined_image_sampler },
        },
        .max_sets = 32,
    }, null);
    defer au.D.destroyDescriptorPool(descriptorPool, null);

    _ = im.c.ImGui_ImplVulkan_LoadFunctions(loader_wrapper, null);
    _ = im.c.ImGui_ImplGlfw_InitForVulkan(@ptrCast(au.W.handle), true);
    defer im.c.ImGui_ImplGlfw_Shutdown();

    _ = try sc.rebuild();

    const prci: vk.PipelineRenderingCreateInfo = .{
        .view_mask = 0,
        .depth_attachment_format = .undefined,
        .stencil_attachment_format = .undefined,
        .color_attachment_count = 1,
        .p_color_attachment_formats = &.{au.device_config.format.format},
    };

    var info: im.c.ImGui_ImplVulkan_InitInfo = .{
        .Instance = @ptrFromInt(@intFromEnum(au.I.handle)),
        .PhysicalDevice = @ptrFromInt(@intFromEnum(au.device_config.pdev)),
        .Device = @ptrFromInt(@intFromEnum(au.D.handle)),
        .QueueFamily = au.device_config.family,
        .Queue = @ptrFromInt(@intFromEnum(au.Q.handle)),
        .DescriptorPool = @ptrFromInt(@intFromEnum(descriptorPool)),
        .RenderPass = null,
        .MinImageCount = 2,
        .ImageCount = @intCast(flights.flights.len),
        .PipelineRenderingCreateInfo = @bitCast(prci),
        .MSAASamples = 0,
        .PipelineCache = null,
        .Subpass = 0,
        .UseDynamicRendering = true,
        .Allocator = null,
    };
    _ = im.c.ImGui_ImplVulkan_Init(&info);
    _ = im.c.ImGui_ImplVulkan_CreateFontsTexture();
    defer im.c.ImGui_ImplVulkan_Shutdown();

    const vert = try au.D.createShaderModule(&vk.ShaderModuleCreateInfo{
        .code_size = shaders.triangle_vert.len,
        .p_code = @ptrCast(&shaders.triangle_vert),
    }, null);
    defer au.D.destroyShaderModule(vert, null);

    const frag = try au.D.createShaderModule(&vk.ShaderModuleCreateInfo{
        .code_size = shaders.triangle_frag.len,
        .p_code = @ptrCast(&shaders.triangle_frag),
    }, null);
    defer au.D.destroyShaderModule(frag, null);

    const cache = try au.D.createPipelineCache(&vk.PipelineCacheCreateInfo{}, null);
    defer au.D.destroyPipelineCache(cache, null);

    const layout = try au.D.createPipelineLayout(&vk.PipelineLayoutCreateInfo{
        // todo
    }, null);
    defer au.D.destroyPipelineLayout(layout, null);

    const gpci: vk.GraphicsPipelineCreateInfo = .{
        .stage_count = 2,
        .p_stages = &.{
            vk.PipelineShaderStageCreateInfo{ .stage = .{ .vertex_bit = true }, .module = vert, .p_name = "main" },
            vk.PipelineShaderStageCreateInfo{ .stage = .{ .fragment_bit = true }, .module = frag, .p_name = "main" },
        },
        .p_vertex_input_state = &vk.PipelineVertexInputStateCreateInfo{
            .vertex_binding_description_count = 1,
            .p_vertex_binding_descriptions = @ptrCast(&Vertex.binding_description),
            .vertex_attribute_description_count = Vertex.attribute_description.len,
            .p_vertex_attribute_descriptions = &Vertex.attribute_description,
        },
        .p_input_assembly_state = &vk.PipelineInputAssemblyStateCreateInfo{
            .topology = .triangle_list,
            .primitive_restart_enable = vk.FALSE,
        },
        .p_tessellation_state = null,
        .p_viewport_state = &vk.PipelineViewportStateCreateInfo{
            .viewport_count = 1,
            .scissor_count = 1,
        },
        .p_rasterization_state = &vk.PipelineRasterizationStateCreateInfo{
            .depth_clamp_enable = vk.FALSE,
            .rasterizer_discard_enable = vk.FALSE,
            .polygon_mode = .fill,
            .cull_mode = .{ .back_bit = true },
            .front_face = .counter_clockwise,
            .depth_bias_enable = vk.FALSE,
            .depth_bias_constant_factor = 0.0,
            .depth_bias_clamp = 0.0,
            .depth_bias_slope_factor = 0.0,
            .line_width = 1.0,
        },
        .p_multisample_state = &vk.PipelineMultisampleStateCreateInfo{
            .rasterization_samples = .{ .@"1_bit" = true },
            .sample_shading_enable = vk.FALSE,
            .min_sample_shading = 1,
            .alpha_to_coverage_enable = vk.FALSE,
            .alpha_to_one_enable = vk.FALSE,
        },
        .p_depth_stencil_state = null,
        .p_color_blend_state = &vk.PipelineColorBlendStateCreateInfo{
            .logic_op_enable = vk.FALSE,
            .logic_op = .copy,
            .blend_constants = [_]f32{ 0, 0, 0, 0 },
            .attachment_count = 1,
            .p_attachments = &.{
                vk.PipelineColorBlendAttachmentState{
                    .blend_enable = vk.FALSE,
                    .color_blend_op = .add,
                    .src_color_blend_factor = .one,
                    .dst_color_blend_factor = .zero,
                    .alpha_blend_op = .add,
                    .src_alpha_blend_factor = .one,
                    .dst_alpha_blend_factor = .zero,
                    .color_write_mask = .{ .r_bit = true, .g_bit = true, .b_bit = true, .a_bit = true },
                },
            },
        },
        .p_dynamic_state = &vk.PipelineDynamicStateCreateInfo{
            .flags = .{},
            .dynamic_state_count = 2,
            .p_dynamic_states = &.{
                .viewport,
                .scissor,
            },
        },
        .layout = layout,
        .render_pass = .null_handle, // set via dynamic rendering
        .subpass = 0,
        .base_pipeline_handle = .null_handle,
        .base_pipeline_index = -1,
        .p_next = &vk.PipelineRenderingCreateInfo{
            .color_attachment_count = 1,
            .p_color_attachment_formats = &.{au.device_config.format.format},
            .depth_attachment_format = .undefined,
            .stencil_attachment_format = .undefined,
            .view_mask = 0,
        },
    };

    var pipeline: vk.Pipeline = undefined;
    _ = try au.D.createGraphicsPipelines(
        cache,
        1,
        @ptrCast(&gpci),
        null,
        @ptrCast(&pipeline),
    );
    defer au.D.destroyPipeline(pipeline, null);

    while (!au.W.should_close()) {
        im.c.ImGui_ImplGlfw_NewFrame();
        im.c.ImGui_ImplVulkan_NewFrame();
        im.c.igNewFrame();
        im.c.igShowDemoWindow(null);
        im.c.igEndFrame();
        im.c.igRender();

        const flight = flights.next();

        const events = if (au.W.focused())
            au.poll_events()
        else
            au.wait_events_timeout(0.5);

        for (events) |u| {
            switch (u) {
                .framebufferSize => sc.mark(),
                .cursorPos, .windowPos, .windowSize, .windowRefresh => {},
                else => |e| std.debug.print("{any}\n", .{e}),
            }
        }

        _ = try sc.rebuild();

        _ = try au.D.waitForFences(1, &.{flight.fence}, vk.TRUE, std.math.maxInt(u64));
        try au.D.resetFences(1, &.{flight.fence});
        try au.D.resetCommandPool(flight.pool, .{});

        const acq = try au.D.acquireNextImageKHR(
            sc.handle,
            std.math.maxInt(u64),
            flight.acquire,
            .null_handle,
        );
        const image = sc.getImage(acq.image_index);
        const view = sc.getView(acq.image_index);

        var cmd = au.CommandBufferProxy.init(flight.cmd, au.D.wrapper);

        try cmd.beginCommandBuffer(&.{ .flags = .{ .one_time_submit_bit = true } });

        try flight.ctx.record_render(
            cmd,
            image,
            view,
            vk.Rect2D{ .offset = .{ .x = 0, .y = 0 }, .extent = sc.cinfo.image_extent },
        );

        try cmd.endCommandBuffer();

        au.Q.submit(
            1,
            &.{
                vk.SubmitInfo{
                    .wait_semaphore_count = 1,
                    .p_wait_semaphores = @ptrCast(&flight.acquire),
                    .p_wait_dst_stage_mask = @ptrCast(&vk.PipelineStageFlags{ .color_attachment_output_bit = true }),
                    .command_buffer_count = 1,
                    .p_command_buffers = @ptrCast(&cmd.handle),
                    .signal_semaphore_count = 1,
                    .p_signal_semaphores = @ptrCast(&flight.complete),
                },
            },
            flight.fence,
        ) catch {
            std.debug.print("Failed to submit.\nWaiting for idle...", .{});
            au.D.deviceWaitIdle() catch
                std.debug.print("deviceWaitIdle failed\n", .{});
            @panic("Submission failed");
        };

        _ = try au.Q.presentKHR(&vk.PresentInfoKHR{
            .wait_semaphore_count = 1,
            .p_wait_semaphores = &.{flight.complete},
            .swapchain_count = 1,
            .p_swapchains = &.{sc.handle},
            .p_image_indices = &.{acq.image_index},
            .p_results = null,
        }); // todo suboptimal?
    }

    try au.D.deviceWaitIdle();

    // const pipeline_layout = try dev.vkd.createPipelineLayout(dev.dev, &.{
    //     .flags = .{},
    //     .set_layout_count = 0,
    //     .p_set_layouts = undefined,
    //     .push_constant_range_count = 0,
    //     .p_push_constant_ranges = undefined,
    // }, null);
    // defer dev.vkd.destroyPipelineLayout(dev.dev, pipeline_layout, null);
    //
    // const pipeline = try createPipeline(dev.dev, pipeline_layout, dev.format, dev.vkd);
    // defer dev.vkd.destroyPipeline(dev.dev, pipeline, null);

    // const vertex_buffer = try dev.vkd.createBuffer(dev.dev, &.{
    //     .size = @sizeOf(@TypeOf(vertices)),
    //     .usage = .{ .transfer_dst_bit = true, .vertex_buffer_bit = true },
    //     .sharing_mode = .exclusive,
    // }, null);
    // defer dev.vkd.destroyBuffer(dev.dev, vertex_buffer, null);
    // const vertex_mem_reqs = dev.vkd.getBufferMemoryRequirements(dev.dev, vertex_buffer);
    // const vertex_memory = try device_local.alloc(dev.dev, dev.vkd, vertex_mem_reqs, .{ .device_local_bit = true });
    // defer dev.vkd.freeMemory(dev.dev, vertex_memory, null);
    // try dev.vkd.bindBufferMemory(dev.dev, vertex_buffer, vertex_memory, 0);
    // try gfx.uploadData(Vertex, dev.pdev, inst.vki, dev.dev, dev.vkd, dev.queue, dev.pool, vertex_buffer, &vertices);

    // const index_buffer = try dev.vkd.createBuffer(dev.dev, &.{
    //     .size = @sizeOf(@TypeOf(indices)),
    //     .usage = .{ .transfer_dst_bit = true, .index_buffer_bit = true },
    //     .sharing_mode = .exclusive,
    // }, null);
    // defer dev.vkd.destroyBuffer(dev.dev, index_buffer, null);
    // const index_mem_reqs = dev.vkd.getBufferMemoryRequirements(dev.dev, index_buffer);
    // const index_memory = try device_local.alloc(dev.dev, dev.vkd, index_mem_reqs, .{ .device_local_bit = true });
    // defer dev.vkd.freeMemory(dev.dev, index_memory, null);
    // try dev.vkd.bindBufferMemory(dev.dev, index_buffer, index_memory, 0);
    // try gfx.uploadData(Index, dev.pdev, inst.vki, dev.dev, dev.vkd, dev.queue, dev.pool, index_buffer, &indices);
}

// fn createPipeline(
//     dev: vk.Device,
//     layout: vk.PipelineLayout,
//     format: vk.SurfaceFormatKHR,
//     vkd: gfx.Device.Wrapper,
// ) !vk.Pipeline {
//     const vert = try vkd.createShaderModule(dev, &.{
//         .code_size = shaders.triangle_vert.len,
//         .p_code = @as([*]const u32, @ptrCast(&shaders.triangle_vert)),
//     }, null);
//     defer vkd.destroyShaderModule(dev, vert, null);
//
//     const frag = try vkd.createShaderModule(dev, &.{
//         .code_size = shaders.triangle_frag.len,
//         .p_code = @as([*]const u32, @ptrCast(&shaders.triangle_frag)),
//     }, null);
//     defer vkd.destroyShaderModule(dev, frag, null);
//
//     const pssci = [_]vk.PipelineShaderStageCreateInfo{
//         .{
//             .stage = .{ .vertex_bit = true },
//             .module = vert,
//             .p_name = "main",
//         },
//         .{
//             .stage = .{ .fragment_bit = true },
//             .module = frag,
//             .p_name = "main",
//         },
//     };
//
//     const color_blend_attachment_states = [_]vk.PipelineColorBlendAttachmentState{
//         vk.PipelineColorBlendAttachmentState{
//             .blend_enable = vk.FALSE,
//             .src_color_blend_factor = .one,
//             .dst_color_blend_factor = .zero,
//             .color_blend_op = .add,
//             .src_alpha_blend_factor = .one,
//             .dst_alpha_blend_factor = .zero,
//             .alpha_blend_op = .add,
//             .color_write_mask = .{ .r_bit = true, .g_bit = true, .b_bit = true, .a_bit = true },
//         },
//     };
//
//     const dynamic_states = [_]vk.DynamicState{
//         .viewport,
//         .scissor,
//     };
//
//     const create_infos = [_]vk.GraphicsPipelineCreateInfo{
//         .{
//             .flags = .{},
//             .stage_count = @intCast(pssci.len),
//             .p_stages = &pssci,
//             .p_vertex_input_state = &vk.PipelineVertexInputStateCreateInfo{
//                 .vertex_binding_description_count = 1,
//                 .p_vertex_binding_descriptions = @ptrCast(&Vertex.binding_description),
//                 .vertex_attribute_description_count = Vertex.attribute_description.len,
//                 .p_vertex_attribute_descriptions = &Vertex.attribute_description,
//             },
//             .p_input_assembly_state = &vk.PipelineInputAssemblyStateCreateInfo{
//                 .topology = .triangle_list,
//                 .primitive_restart_enable = vk.FALSE,
//             },
//             .p_tessellation_state = null,
//             .p_viewport_state = &vk.PipelineViewportStateCreateInfo{
//                 .viewport_count = 1,
//                 .p_viewports = undefined, // set in createCommandBuffers with cmdSetViewport
//                 .scissor_count = 1,
//                 .p_scissors = undefined, // set in createCommandBuffers with cmdSetScissor
//             },
//             .p_rasterization_state = &vk.PipelineRasterizationStateCreateInfo{
//                 .depth_clamp_enable = vk.FALSE,
//                 .rasterizer_discard_enable = vk.FALSE,
//                 .polygon_mode = .fill,
//                 .cull_mode = .{ .back_bit = true },
//                 .front_face = .counter_clockwise,
//                 .depth_bias_enable = vk.FALSE,
//                 .depth_bias_constant_factor = 0,
//                 .depth_bias_clamp = 0,
//                 .depth_bias_slope_factor = 0,
//                 .line_width = 1,
//             },
//             .p_multisample_state = &vk.PipelineMultisampleStateCreateInfo{
//                 .rasterization_samples = .{ .@"1_bit" = true },
//                 .sample_shading_enable = vk.FALSE,
//                 .min_sample_shading = 1,
//                 .alpha_to_coverage_enable = vk.FALSE,
//                 .alpha_to_one_enable = vk.FALSE,
//             },
//             .p_depth_stencil_state = null,
//             .p_color_blend_state = &vk.PipelineColorBlendStateCreateInfo{
//                 .logic_op_enable = vk.FALSE,
//                 .logic_op = .copy,
//                 .attachment_count = @intCast(color_blend_attachment_states.len),
//                 .p_attachments = &color_blend_attachment_states,
//                 .blend_constants = [_]f32{ 0, 0, 0, 0 },
//             },
//             .p_dynamic_state = &vk.PipelineDynamicStateCreateInfo{
//                 .flags = .{},
//                 .dynamic_state_count = @intCast(dynamic_states.len),
//                 .p_dynamic_states = &dynamic_states,
//             },
//             .layout = layout,
//             .render_pass = .null_handle,
//             .subpass = 0,
//             .base_pipeline_handle = .null_handle,
//             .base_pipeline_index = -1,
//             .p_next = &vk.PipelineRenderingCreateInfoKHR{
//                 .color_attachment_count = 1,
//                 .p_color_attachment_formats = @ptrCast(&format),
//                 .depth_attachment_format = .undefined,
//                 .stencil_attachment_format = .undefined,
//                 .view_mask = 0,
//             },
//         },
//     };
//
//     var pipelines: [create_infos.len]vk.Pipeline = undefined;
//     _ = try vkd.createGraphicsPipelines(dev, .null_handle, @intCast(create_infos.len), &create_infos, null, &pipelines);
//     std.debug.assert(pipelines.len == 1);
//     return pipelines[0];
// }