const std = @import("std");
const glfw = @import("glfw");
const vk = @import("vulkan");
const resources = @import("resources");
const Allocator = std.mem.Allocator;
const GraphicsContext = @import("graphics_context.zig").GraphicsContext;
const Swapchain = @import("swapchain.zig").Swapchain;
const Vertex = @import("vertex.zig").Vertex;

// TODO: model loading!
const vertices = [_]Vertex{
    .{ .pos = .{ 0, -0.5 }, .color = .{ 1, 0, 0 } },
    .{ .pos = .{ 0.5, 0.5 }, .color = .{ 0, 1, 0 } },
    .{ .pos = .{ -0.5, 0.5 }, .color = .{ 0, 0, 1 } },
};

pub const VulkanRenderer = struct {
    gc: *const GraphicsContext,
    allocator: Allocator,
    window: *glfw.Window,
    swapchain: Swapchain,
    render_pass: vk.RenderPass,
    pipeline_layout: vk.PipelineLayout,
    pipeline: vk.Pipeline,
    cmdbufs: []vk.CommandBuffer,
    framebuffers: []vk.Framebuffer,
    pool: vk.CommandPool,
    buffer: vk.Buffer,
    memory: vk.DeviceMemory,
    viewport: vk.Viewport,
    scissor: vk.Rect2D,
    window_resized: bool,
    new_extent: vk.Extent2D,

    pub fn init(gc: *const GraphicsContext, allocator: Allocator, extent: vk.Extent2D, window: *glfw.Window) !VulkanRenderer {
        var self: VulkanRenderer = undefined;

        self.window = window;
        self.allocator = allocator;
        window.setUserPointer(&self);
        window.setFramebufferSizeCallback(resizeCallback);

        self.gc = gc;
        self.swapchain = try Swapchain.init(gc, allocator, extent);

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

        self.render_pass = try self.createRenderPass();

        self.pipeline = try self.createPipeline();

        self.framebuffers = try self.createFramebuffers();

        self.pool = try gc.vkd.createCommandPool(gc.dev, &.{
            .flags = .{},
            .queue_family_index = gc.graphics_queue.family,
        }, null);

        self.buffer = try gc.vkd.createBuffer(gc.dev, &.{
            .flags = .{},
            .size = @sizeOf(@TypeOf(vertices)),
            .usage = .{ .transfer_dst_bit = true, .vertex_buffer_bit = true },
            .sharing_mode = .exclusive,
            .queue_family_index_count = 0,
            .p_queue_family_indices = undefined,
        }, null);
        const mem_reqs = gc.vkd.getBufferMemoryRequirements(gc.dev, self.buffer);
        self.memory = try gc.allocate(mem_reqs, .{ .device_local_bit = true });
        try gc.vkd.bindBufferMemory(gc.dev, self.buffer, self.memory, 0);

        try self.uploadVertices();

        self.cmdbufs = try self.createCommandBuffers(extent);

        return self;
    }

    pub fn reinit(self: *VulkanRenderer, extent: vk.Extent2D) !void {
        self.swapchain.recreate(extent);
    }

    pub fn deinit(self: *VulkanRenderer) void {
        // todo: log?
        self.swapchain.waitForAllFences() catch {};

        self.destroyCommandBuffers();
        self.gc.vkd.freeMemory(self.gc.dev, self.memory, null);
        self.gc.vkd.destroyBuffer(self.gc.dev, self.buffer, null);
        self.gc.vkd.destroyCommandPool(self.gc.dev, self.pool, null);
        self.destroyFramebuffers();
        self.gc.vkd.destroyPipeline(self.gc.dev, self.pipeline, null);
        self.gc.vkd.destroyRenderPass(self.gc.dev, self.render_pass, null);
        self.gc.vkd.destroyPipelineLayout(self.gc.dev, self.pipeline_layout, null);

        self.swapchain.deinit();
    }

    pub fn drawFrame(self: *VulkanRenderer) !void {
        const cmdbuf = self.cmdbufs[self.swapchain.image_index];

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

        if (state == .suboptimal or self.window_resized) {
            // todo: fn
            self.destroyFramebuffers();

            self.gc.vkd.destroyPipeline(self.gc.dev, self.pipeline, null);
            self.gc.vkd.destroyPipelineLayout(self.gc.dev, self.pipeline_layout, null);
            self.gc.vkd.destroyRenderPass(self.gc.dev, self.render_pass, null);

            try self.swapchain.recreate(self.new_extent);

            // todo: dedupe creation logic
            self.pipeline_layout = try self.gc.vkd.createPipelineLayout(self.gc.dev, &.{
                .flags = .{},
                .set_layout_count = 0,
                .p_set_layouts = undefined,
                .push_constant_range_count = 0,
                .p_push_constant_ranges = undefined,
            }, null);

            self.render_pass = try self.createRenderPass();
            self.pipeline = try self.createPipeline();
            self.framebuffers = try self.createFramebuffers();
            self.window_resized = false;
        }
    }

    fn createRenderPass(self: *VulkanRenderer) !vk.RenderPass {
        const color_attachment = vk.AttachmentDescription{
            .flags = .{},
            .format = self.swapchain.surface_format.format,
            .samples = .{ .@"1_bit" = true },
            .load_op = .clear,
            .store_op = .store,
            .stencil_load_op = .dont_care,
            .stencil_store_op = .dont_care,
            .initial_layout = .@"undefined",
            .final_layout = .present_src_khr,
        };

        const color_attachment_ref = vk.AttachmentReference{
            .attachment = 0,
            .layout = .color_attachment_optimal,
        };

        const subpass = vk.SubpassDescription{
            .flags = .{},
            .pipeline_bind_point = .graphics,
            .input_attachment_count = 0,
            .p_input_attachments = undefined,
            .color_attachment_count = 1,
            .p_color_attachments = @ptrCast([*]const vk.AttachmentReference, &color_attachment_ref),
            .p_resolve_attachments = null,
            .p_depth_stencil_attachment = null,
            .preserve_attachment_count = 0,
            .p_preserve_attachments = undefined,
        };

        return try self.gc.vkd.createRenderPass(self.gc.dev, &.{
            .flags = .{},
            .attachment_count = 1,
            .p_attachments = @ptrCast([*]const vk.AttachmentDescription, &color_attachment),
            .subpass_count = 1,
            .p_subpasses = @ptrCast([*]const vk.SubpassDescription, &subpass),
            .dependency_count = 0,
            .p_dependencies = undefined,
        }, null);
    }

    fn createPipeline(self: *VulkanRenderer) !vk.Pipeline {
        const vert = try self.gc.vkd.createShaderModule(self.gc.dev, &.{
            .flags = .{},
            .code_size = resources.triangle_vert.len,
            .p_code = @ptrCast([*]const u32, resources.triangle_vert),
        }, null);
        defer self.gc.vkd.destroyShaderModule(self.gc.dev, vert, null);

        const frag = try self.gc.vkd.createShaderModule(self.gc.dev, &.{
            .flags = .{},
            .code_size = resources.triangle_frag.len,
            .p_code = @ptrCast([*]const u32, resources.triangle_frag),
        }, null);
        defer self.gc.vkd.destroyShaderModule(self.gc.dev, frag, null);

        const pssci = [_]vk.PipelineShaderStageCreateInfo{
            .{
                .flags = .{},
                .stage = .{ .vertex_bit = true },
                .module = vert,
                .p_name = "main",
                .p_specialization_info = null,
            },
            .{
                .flags = .{},
                .stage = .{ .fragment_bit = true },
                .module = frag,
                .p_name = "main",
                .p_specialization_info = null,
            },
        };

        const pvisci = vk.PipelineVertexInputStateCreateInfo{
            .flags = .{},
            .vertex_binding_description_count = 1,
            .p_vertex_binding_descriptions = @ptrCast([*]const vk.VertexInputBindingDescription, &Vertex.binding_description),
            .vertex_attribute_description_count = Vertex.attribute_description.len,
            .p_vertex_attribute_descriptions = &Vertex.attribute_description,
        };

        const piasci = vk.PipelineInputAssemblyStateCreateInfo{
            .flags = .{},
            .topology = .triangle_list,
            .primitive_restart_enable = vk.FALSE,
        };

        const pvsci = vk.PipelineViewportStateCreateInfo{
            .flags = .{},
            .viewport_count = 1,
            .p_viewports = undefined, // set in createCommandBuffers iwth cmdSetViewport
            .scissor_count = 1,
            .p_scissors = undefined, //set in createCommandBuffers
        };

        const prsci = vk.PipelineRasterizationStateCreateInfo{
            .flags = .{},
            .depth_clamp_enable = vk.FALSE,
            .rasterizer_discard_enable = vk.FALSE,
            .polygon_mode = .fill,
            .cull_mode = .{ .back_bit = true },
            .front_face = .clockwise,
            .depth_bias_enable = vk.FALSE,
            .depth_bias_constant_factor = 0,
            .depth_bias_clamp = 0,
            .depth_bias_slope_factor = 0,
            .line_width = 1,
        };

        const pmsci = vk.PipelineMultisampleStateCreateInfo{
            .flags = .{},
            .rasterization_samples = .{ .@"1_bit" = true },
            .sample_shading_enable = vk.FALSE,
            .min_sample_shading = 1,
            .p_sample_mask = null,
            .alpha_to_coverage_enable = vk.FALSE,
            .alpha_to_one_enable = vk.FALSE,
        };

        const pcbas = 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 pcbsci = vk.PipelineColorBlendStateCreateInfo{
            .flags = .{},
            .logic_op_enable = vk.FALSE,
            .logic_op = .copy,
            .attachment_count = 1,
            .p_attachments = @ptrCast([*]const vk.PipelineColorBlendAttachmentState, &pcbas),
            .blend_constants = [_]f32{ 0, 0, 0, 0 },
        };

        const dynstate = [_]vk.DynamicState{ .viewport, .scissor };
        const pdsci = vk.PipelineDynamicStateCreateInfo{
            .flags = .{},
            .dynamic_state_count = dynstate.len,
            .p_dynamic_states = &dynstate,
        };

        const gpci = vk.GraphicsPipelineCreateInfo{
            .flags = .{},
            .stage_count = 2,
            .p_stages = &pssci,
            .p_vertex_input_state = &pvisci,
            .p_input_assembly_state = &piasci,
            .p_tessellation_state = null,
            .p_viewport_state = &pvsci,
            .p_rasterization_state = &prsci,
            .p_multisample_state = &pmsci,
            .p_depth_stencil_state = null,
            .p_color_blend_state = &pcbsci,
            .p_dynamic_state = &pdsci,
            .layout = self.pipeline_layout,
            .render_pass = self.render_pass,
            .subpass = 0,
            .base_pipeline_handle = .null_handle,
            .base_pipeline_index = -1,
        };

        var pipeline: vk.Pipeline = undefined;
        _ = try self.gc.vkd.createGraphicsPipelines(
            self.gc.dev,
            .null_handle,
            1,
            @ptrCast([*]const vk.GraphicsPipelineCreateInfo, &gpci),
            null,
            @ptrCast([*]vk.Pipeline, &pipeline),
        );
        return pipeline;
    }

    fn createFramebuffers(self: *VulkanRenderer) ![]vk.Framebuffer {
        const framebuffers = try self.allocator.alloc(vk.Framebuffer, self.swapchain.swap_images.len);
        errdefer self.allocator.free(framebuffers);

        var i: usize = 0;
        errdefer for (self.framebuffers[0..i]) |fb| self.gc.vkd.destroyFramebuffer(self.gc.dev, fb, null);

        for (framebuffers) |*fb| {
            fb.* = try self.gc.vkd.createFramebuffer(self.gc.dev, &.{
                .flags = .{},
                .render_pass = self.render_pass,
                .attachment_count = 1,
                .p_attachments = @ptrCast([*]const vk.ImageView, &self.swapchain.swap_images[i].view),
                .width = self.swapchain.extent.width,
                .height = self.swapchain.extent.height,
                .layers = 1,
            }, null);
            i += 1;
        }

        return framebuffers;
    }

    fn createCommandBuffers(self: *VulkanRenderer, extent: vk.Extent2D) ![]vk.CommandBuffer {
        const cmdbufs = try self.allocator.alloc(vk.CommandBuffer, self.framebuffers.len);
        errdefer self.allocator.free(cmdbufs);

        try self.gc.vkd.allocateCommandBuffers(self.gc.dev, &.{
            .command_pool = self.pool,
            .level = .primary,
            .command_buffer_count = @truncate(u32, cmdbufs.len),
        }, cmdbufs.ptr);
        errdefer self.gc.vkd.freeCommandBuffers(self.gc.dev, self.pool, @truncate(u32, cmdbufs.len), cmdbufs.ptr);

        const clear = vk.ClearValue{
            .color = .{ .float_32 = .{ 0, 0, 0, 1 } },
        };

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

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

        for (cmdbufs) |cmdbuf, i| {
            _ = i;
            try self.gc.vkd.beginCommandBuffer(cmdbuf, &.{
                .flags = .{},
                .p_inheritance_info = null,
            });

            self.gc.vkd.cmdSetViewport(cmdbuf, 0, 1, @ptrCast([*]const vk.Viewport, &self.viewport));
            self.gc.vkd.cmdSetScissor(cmdbuf, 0, 1, @ptrCast([*]const vk.Rect2D, &self.scissor));

            // This needs to be a separate definition - see https://github.com/ziglang/zig/issues/7627.
            const render_area = vk.Rect2D{
                .offset = .{ .x = 0, .y = 0 },
                .extent = extent,
            };

            self.gc.vkd.cmdBeginRenderPass(cmdbuf, &.{
                .render_pass = self.render_pass,
                .framebuffer = self.framebuffers[i],
                .render_area = render_area,
                .clear_value_count = 1,
                .p_clear_values = @ptrCast([*]const vk.ClearValue, &clear),
            }, .@"inline");

            self.gc.vkd.cmdBindPipeline(cmdbuf, .graphics, self.pipeline);
            const offset = [_]vk.DeviceSize{0};
            self.gc.vkd.cmdBindVertexBuffers(cmdbuf, 0, 1, @ptrCast([*]const vk.Buffer, &self.buffer), &offset);
            self.gc.vkd.cmdDraw(cmdbuf, vertices.len, 1, 0, 0);

            self.gc.vkd.cmdEndRenderPass(cmdbuf);
            try self.gc.vkd.endCommandBuffer(cmdbuf);
        }

        return cmdbufs;
    }

    fn destroyCommandBuffers(self: *VulkanRenderer) void {
        self.gc.vkd.freeCommandBuffers(self.gc.dev, self.pool, @truncate(u32, self.cmdbufs.len), self.cmdbufs.ptr);
        self.allocator.free(self.cmdbufs);
    }

    fn destroyFramebuffers(self: *VulkanRenderer) void {
        for (self.framebuffers) |fb| self.gc.vkd.destroyFramebuffer(self.gc.dev, fb, null);
        self.allocator.free(self.framebuffers);
    }

    fn copyBuffer(self: *VulkanRenderer, dst: vk.Buffer, src: vk.Buffer, size: vk.DeviceSize) !void {
        var cmdbuf: vk.CommandBuffer = undefined;
        try self.gc.vkd.allocateCommandBuffers(self.gc.dev, &.{
            .command_pool = self.pool,
            .level = .primary,
            .command_buffer_count = 1,
        }, @ptrCast([*]vk.CommandBuffer, &cmdbuf));
        defer self.gc.vkd.freeCommandBuffers(self.gc.dev, self.pool, 1, @ptrCast([*]const vk.CommandBuffer, &cmdbuf));

        try self.gc.vkd.beginCommandBuffer(cmdbuf, &.{
            .flags = .{ .one_time_submit_bit = true },
            .p_inheritance_info = null,
        });

        const region = vk.BufferCopy{
            .src_offset = 0,
            .dst_offset = 0,
            .size = size,
        };
        self.gc.vkd.cmdCopyBuffer(cmdbuf, src, dst, 1, @ptrCast([*]const vk.BufferCopy, &region));

        try self.gc.vkd.endCommandBuffer(cmdbuf);

        const si = vk.SubmitInfo{
            .wait_semaphore_count = 0,
            .p_wait_semaphores = undefined,
            .p_wait_dst_stage_mask = undefined,
            .command_buffer_count = 1,
            .p_command_buffers = @ptrCast([*]const vk.CommandBuffer, &cmdbuf),
            .signal_semaphore_count = 0,
            .p_signal_semaphores = undefined,
        };
        try self.gc.vkd.queueSubmit(self.gc.graphics_queue.handle, 1, @ptrCast([*]const vk.SubmitInfo, &si), .null_handle);
        try self.gc.vkd.queueWaitIdle(self.gc.graphics_queue.handle);
    }

    fn uploadVertices(self: *VulkanRenderer) !void {
        const staging_buffer = try self.gc.vkd.createBuffer(self.gc.dev, &.{
            .flags = .{},
            .size = @sizeOf(@TypeOf(vertices)),
            .usage = .{ .transfer_src_bit = true },
            .sharing_mode = .exclusive,
            .queue_family_index_count = 0,
            .p_queue_family_indices = undefined,
        }, null);

        defer self.gc.vkd.destroyBuffer(self.gc.dev, staging_buffer, null);
        const mem_reqs = self.gc.vkd.getBufferMemoryRequirements(self.gc.dev, staging_buffer);
        const staging_memory = try self.gc.allocate(mem_reqs, .{ .host_visible_bit = true, .host_coherent_bit = true });
        defer self.gc.vkd.freeMemory(self.gc.dev, staging_memory, null);
        try self.gc.vkd.bindBufferMemory(self.gc.dev, staging_buffer, staging_memory, 0);

        {
            const data = try self.gc.vkd.mapMemory(self.gc.dev, staging_memory, 0, vk.WHOLE_SIZE, .{});
            defer self.gc.vkd.unmapMemory(self.gc.dev, staging_memory);

            const gpu_vertices = @ptrCast([*]Vertex, @alignCast(@alignOf(Vertex), data));
            for (vertices) |vertex, i| {
                gpu_vertices[i] = vertex;
            }
        }

        try self.copyBuffer(self.buffer, staging_buffer, @sizeOf(@TypeOf(vertices)));
    }
};

fn resizeCallback(window: glfw.Window, width: u32, height: u32) void {
    if (window.getUserPointer(VulkanRenderer)) |renderer| {
        renderer.window_resized = true;
        renderer.new_extent.width = width;
        renderer.new_extent.height = height;
    }
}