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

const gfx = @import("gfx.zig");

const app_name = "vulkan-zig triangle example";

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

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

    if (c.glfwInit() != c.GLFW_TRUE) return error.GlfwInitFailed;
    defer c.glfwTerminate();

    if (c.glfwVulkanSupported() != c.GLFW_TRUE) {
        std.log.err("GLFW could not find libvulkan", .{});
        return error.NoVulkan;
    }

    var extent = vk.Extent2D{ .width = 800, .height = 600 };

    const window = try gfx.create_window(extent, app_name);
    defer c.glfwDestroyWindow(window);

    const vkb = try gfx.BaseDispatch.load(c.glfwGetInstanceProcAddress);

    const instance, const vki, const messenger = try gfx.create_instance(vkb, app_name);
    defer vki.destroyInstance(instance, null);
    defer if (gfx.use_debug_messenger)
        vki.destroyDebugUtilsMessengerEXT(instance, messenger, null);

    const surface = try gfx.create_surface(instance, window);
    defer vki.destroySurfaceKHR(instance, surface, null);

    const pdev: vk.PhysicalDevice, const dev: vk.Device, const vkd: gfx.DeviceDispatch, const family: u32 =
        try gfx.create_device(ally, instance, surface, vki);
    defer vkd.destroyDevice(dev, null);

    const queue = vkd.getDeviceQueue(dev, family, 0);

    const pool = try vkd.createCommandPool(dev, &.{
        .queue_family_index = family,
    }, null);
    defer vkd.destroyCommandPool(dev, pool, null);

    const preferred_format: vk.SurfaceFormatKHR = .{
        .format = .b8g8r8a8_srgb,
        .color_space = .srgb_nonlinear_khr,
    };
    const format = try gfx.find_surface_format(pdev, vki, surface, preferred_format);
    extent = try gfx.find_swap_extent(pdev, vki, surface, window);

    const present_mode = try gfx.find_present_mode(pdev, vki, surface, .mailbox_khr);

    const swap_image_count = try gfx.find_swap_image_count(pdev, vki, surface);

    var swapchain: vk.SwapchainKHR = .null_handle;
    defer vkd.destroySwapchainKHR(dev, swapchain, null);

    const ChainImage = struct {
        image: vk.Image = .null_handle,
        view: vk.ImageView = .null_handle,
        cmdbuf: vk.CommandBuffer = .null_handle,
        fence: vk.Fence = .null_handle,
        image_available: vk.Semaphore = .null_handle,
        render_finished: vk.Semaphore = .null_handle,
    };

    var chain = std.MultiArrayList(ChainImage){};
    defer chain.deinit(ally);
    defer vkd.freeCommandBuffers(dev, pool, @intCast(chain.len), chain.items(.cmdbuf).ptr);
    defer for (chain.items(.view)) |view| vkd.destroyImageView(dev, view, null);
    defer for (chain.items(.fence)) |fence| vkd.destroyFence(dev, fence, null);
    defer for (chain.items(.image_available)) |sem| vkd.destroySemaphore(dev, sem, null);
    defer for (chain.items(.render_finished)) |sem| vkd.destroySemaphore(dev, sem, null);

    swapchain = try vkd.createSwapchainKHR(dev, &.{
        .surface = surface,
        .min_image_count = swap_image_count,
        .image_format = format.format,
        .image_color_space = format.color_space,
        .image_extent = extent,
        .image_array_layers = 1,
        .image_usage = .{ .color_attachment_bit = true },
        .image_sharing_mode = .exclusive,
        .pre_transform = .{ .identity_bit_khr = true },
        .composite_alpha = .{ .opaque_bit_khr = true },
        .present_mode = present_mode,
        .clipped = vk.TRUE,
        .old_swapchain = swapchain,
    }, null);

    var image_count: u32 = undefined;
    _ = try vkd.getSwapchainImagesKHR(dev, swapchain, &image_count, null);
    try chain.resize(ally, image_count);
    _ = try vkd.getSwapchainImagesKHR(dev, swapchain, &image_count, chain.items(.image).ptr);

    for (chain.items(.image), chain.items(.view)) |image, *view| {
        view.* = try vkd.createImageView(dev, &.{
            .image = image,
            .view_type = .@"2d",
            .format = format.format,
            .components = .{ .r = .identity, .g = .identity, .b = .identity, .a = .identity },
            .subresource_range = .{
                .aspect_mask = .{ .color_bit = true },
                .base_mip_level = 0,
                .level_count = 1,
                .base_array_layer = 0,
                .layer_count = 1,
            },
        }, null);
    }

    for (chain.items(.fence)) |*fence| {
        fence.* = try vkd.createFence(dev, &.{ .flags = .{ .signaled_bit = true } }, null);
    }

    for (chain.items(.image_available)) |*sem| {
        sem.* = try vkd.createSemaphore(dev, &.{}, null);
    }

    for (chain.items(.render_finished)) |*sem| {
        sem.* = try vkd.createSemaphore(dev, &.{}, null);
    }

    try vkd.allocateCommandBuffers(dev, &.{
        .command_buffer_count = @intCast(chain.len),
        .command_pool = pool,
        .level = .primary,
    }, chain.items(.cmdbuf).ptr);

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

    const pipeline = try createPipeline(dev, pipeline_layout, format, vkd);
    defer vkd.destroyPipeline(dev, pipeline, null);

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

    try gfx.uploadData(Vertex, pdev, vki, dev, vkd, queue, pool, vertex_buffer, &vertices);

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

    try gfx.uploadData(Index, pdev, vki, dev, vkd, queue, pool, index_buffer, &indices);

    for (chain.items(.image), chain.items(.view), chain.items(.cmdbuf)) |image, view, cmdbuf| {
        try record_cmdbuf(cmdbuf, vkd, image, view, extent, pipeline, vertex_buffer, index_buffer);
    }

    var index: u32 = 0;

    while (c.glfwWindowShouldClose(window) == c.GLFW_FALSE) {
        var w: c_int = undefined;
        var h: c_int = undefined;
        c.glfwGetFramebufferSize(window, &w, &h);

        // Don't present or resize swapchain while the window is minimized
        if (w == 0 or h == 0) {
            c.glfwPollEvents();
            continue;
        }

        const frame: ChainImage = chain.get(index);
        // const next_frame: ChainImage = chain.get((index + 1) % chain.len);

        _ = try vkd.waitForFences(dev, 1, @ptrCast(&frame.fence), vk.TRUE, std.math.maxInt(u64));
        try vkd.resetFences(dev, 1, @ptrCast(&frame.fence));

        // var index: u32 = undefined;
        // try vkd.acquireNextImageKHR(dev, swapchain, std.math.maxInt(u64), frame., fence);
        const result = try vkd.acquireNextImageKHR(dev, swapchain, std.math.maxInt(u64), frame.image_available, .null_handle);

        // std.log.debug("frame {d}", .{result.image_index});
        // const frame = chain.get(result.image_index);

        try vkd.queueSubmit(queue, 1, @ptrCast(&vk.SubmitInfo{
            .wait_semaphore_count = 1,
            .p_wait_semaphores = @ptrCast(&frame.image_available),
            .p_wait_dst_stage_mask = @ptrCast(&vk.PipelineStageFlags{ .color_attachment_output_bit = true }),
            .command_buffer_count = 1,
            .p_command_buffers = @ptrCast(&frame.cmdbuf),
            .signal_semaphore_count = 1,
            .p_signal_semaphores = @ptrCast(&frame.render_finished),
        }), frame.fence);

        _ = try vkd.queuePresentKHR(queue, &.{
            .wait_semaphore_count = 1,
            .p_wait_semaphores = @ptrCast(&frame.render_finished),
            .swapchain_count = 1,
            .p_swapchains = @ptrCast(&swapchain),
            .p_image_indices = @ptrCast(&result.image_index),
            .p_results = null,
        });

        // const cmdbuf = cmdbufs[swapchain.image_index];

        // const state = swapchain.present(cmdbuf) catch |err| switch (err) {
        //     error.OutOfDateKHR => Swapchain.PresentState.suboptimal,
        //     else => |narrow| return narrow,
        // };

        // if (state == .suboptimal or extent.width != @as(u32, @intCast(w)) or extent.height != @as(u32, @intCast(h))) {
        //     extent.width = @intCast(w);
        //     extent.height = @intCast(h);
        //     try swapchain.recreate(extent, format);
        //
        //     destroyCommandBuffers(&gc, pool, ally, cmdbufs);
        //
        //     cmdbufs = try createCommandBuffers(
        //         &gc,
        //         pool,
        //         ally,
        //         vertex_buffer,
        //         index_buffer,
        //         pipeline,
        //         swapchain,
        //     );
        // }

        c.glfwPollEvents();

        index = @intCast((index + 1) % chain.len);
    }

    // try swapchain.waitForAllFences();
    try vkd.deviceWaitIdle(dev);
}

fn record_cmdbuf(
    cmdbuf: vk.CommandBuffer,
    vkd: gfx.DeviceDispatch,
    image: vk.Image,
    view: vk.ImageView,
    extent: vk.Extent2D,
    pipeline: vk.Pipeline,
    vertex_buffer: vk.Buffer,
    index_buffer: vk.Buffer,
) !void {
    const clear = vk.ClearValue{
        .color = .{ .float_32 = .{ 0, 0, 0, 1 } },
    };

    const viewport = vk.Viewport{
        .x = 0,
        .y = 0,
        .width = @floatFromInt(extent.width),
        .height = @floatFromInt(extent.height),
        .min_depth = 0,
        .max_depth = 1,
    };

    const scissor = vk.Rect2D{
        .offset = .{ .x = 0, .y = 0 },
        .extent = extent,
    };

    try vkd.beginCommandBuffer(cmdbuf, &.{});

    vkd.cmdPipelineBarrier(
        cmdbuf,
        .{ .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,
            },
        }),
    );

    vkd.cmdSetViewport(cmdbuf, 0, 1, @ptrCast(&viewport));
    vkd.cmdSetScissor(cmdbuf, 0, 1, @ptrCast(&scissor));

    const color_attachments = [_]vk.RenderingAttachmentInfoKHR{
        .{
            .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 = clear,
        },
    };

    const render_info = vk.RenderingInfoKHR{
        .render_area = scissor, // since we always do full-frame changes
        .layer_count = 1,
        .view_mask = 0,
        .color_attachment_count = color_attachments.len,
        .p_color_attachments = &color_attachments,
    };

    vkd.cmdBeginRenderingKHR(cmdbuf, &render_info);

    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);

    vkd.cmdEndRenderingKHR(cmdbuf);

    vkd.cmdPipelineBarrier(
        cmdbuf,
        .{ .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,
            },
        }),
    );

    try vkd.endCommandBuffer(cmdbuf);
}

fn createPipeline(dev: vk.Device, layout: vk.PipelineLayout, format: vk.SurfaceFormatKHR, vkd: gfx.DeviceDispatch) !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];
}