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

const vk = @import("vulkan");

const Engine = @import("Engine.zig");
const QSM = @import("QueueSharingMode.zig");

params: Params,
render_pass: vk.RenderPass,

extent: vk.Extent2D = .{ .width = 0, .height = 0 },
swapchain: vk.SwapchainKHR = .null_handle,
swapchain_images: []SwapchainImage = &.{},
image_index: u32 = 0,
semaphore_image_acquired: vk.Semaphore = .null_handle,

pub const PresentResult = enum {
    optimal,
    suboptimal,
};

pub fn init(engine: *Engine) !Swapchain {
    const params: Params = try .init(engine);

    const render_pass = blk: {
        const attachments = [_]vk.AttachmentDescription{
            .{
                .format = params.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_attachments = [_]vk.AttachmentReference{
            .{
                .attachment = 0,
                .layout = .color_attachment_optimal,
            },
        };

        const subpasses = [_]vk.SubpassDescription{
            .{
                .pipeline_bind_point = .graphics,
                .color_attachment_count = color_attachments.len,
                .p_color_attachments = &color_attachments,
            },
        };

        break :blk try engine.device.createRenderPass(&.{
            .attachment_count = attachments.len,
            .p_attachments = &attachments,
            .subpass_count = subpasses.len,
            .p_subpasses = &subpasses,
        }, &engine.vk_allocator.interface);
    };
    errdefer engine.device.destroyRenderPass(render_pass, &engine.vk_allocator.interface);

    var swapchain: Swapchain = .{
        .params = params,
        .render_pass = render_pass,
    };

    try recreate(&swapchain, engine);
    return swapchain;
}

pub fn deinit(self: *Swapchain, engine: *Engine) void {
    const allocator = engine.vk_allocator.allocator;

    for (self.swapchain_images) |swapchain_image| {
        swapchain_image.deinit(engine);
    }
    allocator.free(self.swapchain_images);

    if (self.swapchain != .null_handle) {
        engine.device.destroySwapchainKHR(self.swapchain, &engine.vk_allocator.interface);
    }

    engine.device.destroyRenderPass(self.render_pass, &engine.vk_allocator.interface);

    self.* = undefined;
}

pub fn recreate(self: *Swapchain, engine: *Engine) !void {
    const mode = &engine.mode.surface;
    const allocator = engine.vk_allocator.allocator;

    const old_swapchain = self.swapchain;
    const old_swapchain_images = self.swapchain_images;
    const old_semaphore_image_acquired = self.semaphore_image_acquired;

    const extent = try getCurrentExtent(engine);

    std.log.debug("Recreating swapchain with extent of {d}×{d}...", .{ extent.width, extent.height });

    // --- CREATE NEW SWAPCHAIN ------------------------------------------------

    const surface_capabilities = try engine.instance.getPhysicalDeviceSurfaceCapabilitiesKHR(engine.physical_device, mode.surface);

    const qsm = QSM.resolve(engine.graphics_queue.allocation.family, mode.presentation_queue.allocation.family);
    const new_swapchain = try engine.device.createSwapchainKHR(&.{
        .surface = mode.surface,
        .min_image_count = self.params.image_count,
        .image_format = self.params.surface_format.format,
        .image_color_space = self.params.surface_format.color_space,
        .image_extent = extent,
        .image_array_layers = 1,
        .image_usage = .{ .color_attachment_bit = true },
        .image_sharing_mode = qsm.sharing_mode,
        .queue_family_index_count = qsm.queue_family_index_count,
        .p_queue_family_indices = qsm.p_queue_family_indices,
        .pre_transform = surface_capabilities.current_transform,
        .composite_alpha = .{ .opaque_bit_khr = true },
        .present_mode = .fifo_khr,
        .clipped = .true,
        .old_swapchain = old_swapchain,
    }, &engine.vk_allocator.interface);
    errdefer engine.device.destroySwapchainKHR(new_swapchain, &engine.vk_allocator.interface);

    // --- DESTROY OLD SWAPCHAIN AND ITS IMAGES --------------------------------

    // NOTE At this point, the old swapchain has been recycled. There is no way
    // to revert on error. Instead, we set the current swapchain and images to
    // null and deinit the old swapchain and images.

    for (old_swapchain_images) |swapchain_image| {
        swapchain_image.deinit(engine);
    }
    allocator.free(self.swapchain_images);
    self.swapchain_images = &.{};

    if (old_swapchain != .null_handle) {
        engine.device.destroySwapchainKHR(old_swapchain, &engine.vk_allocator.interface);
        self.swapchain = .null_handle;
    }

    self.extent = .{ .width = 0, .height = 0 };

    if (old_semaphore_image_acquired != .null_handle) {
        engine.destroySemaphore(old_semaphore_image_acquired);
        self.semaphore_image_acquired = .null_handle;
    }

    // --- CREATE NEW SWAPCHAIN IMAGES -----------------------------------------

    const new_swapchain_images = blk: {
        const images = try engine.device.getSwapchainImagesAllocKHR(new_swapchain, allocator);
        defer allocator.free(images);

        var swapchain_images: std.ArrayList(SwapchainImage) = try .initCapacity(allocator, images.len);
        errdefer swapchain_images.deinit(allocator);

        errdefer for (swapchain_images.items) |x| x.deinit(engine);

        for (images) |image| {
            swapchain_images.appendAssumeCapacity(try SwapchainImage.init(engine, .{
                .image = image,
                .format = self.params.surface_format.format,
                .render_pass = self.render_pass,
                .extent = extent,
            }));
        }

        break :blk try swapchain_images.toOwnedSlice(allocator);
    };
    errdefer {
        for (new_swapchain_images) |swapchain_image| {
            swapchain_image.deinit(engine);
        }
        allocator.free(new_swapchain_images);
    }

    // --- ACQUIRE NEXT IMAGE --------------------------------------------------

    var semaphore_image_acquired = try engine.createSemaphore();
    errdefer engine.destroySemaphore(semaphore_image_acquired);

    const res = try engine.device.acquireNextImageKHR(new_swapchain, std.math.maxInt(u64), semaphore_image_acquired, .null_handle);
    if (res.result == .not_ready or res.result == .timeout) {
        return error.ImageAcquireFailed;
    }

    std.mem.swap(vk.Semaphore, &new_swapchain_images[res.image_index].semaphore_image_acquired, &semaphore_image_acquired);

    // --- COMMIT --------------------------------------------------------------

    self.extent = extent;
    self.swapchain = new_swapchain;
    self.swapchain_images = new_swapchain_images;
    self.image_index = res.image_index;
    self.semaphore_image_acquired = semaphore_image_acquired;

    std.log.debug("Swapchain recreated.", .{});
}

pub fn present(self: *Swapchain, engine: *Engine, command_buffer: vk.CommandBuffer) !PresentResult {
    const device = engine.device;
    const mode = &engine.mode.surface;

    std.log.debug("Presenting command buffer {X}.", .{@intFromEnum(command_buffer)});

    // --- WAIT FOR CURRENT FRAME TO FINISH ------------------------------------

    const current_swapchain_image = &self.swapchain_images[self.image_index];
    try engine.waitForFence(current_swapchain_image.fence);
    if (current_swapchain_image.command_buffer != .null_handle) {
        device.freeCommandBuffers(engine.graphics_command_pool, 1, @ptrCast(&current_swapchain_image.command_buffer));
        current_swapchain_image.command_buffer = .null_handle;
    }
    try engine.resetFence(current_swapchain_image.fence);

    // --- SUBMIT COMMAND BUFFER -----------------------------------------------

    const wait_semaphores = [_]vk.Semaphore{
        current_swapchain_image.semaphore_image_acquired,
    };

    const pipeline_stages_flags = [_]vk.PipelineStageFlags{
        .{ .top_of_pipe_bit = true },
    };

    std.debug.assert(wait_semaphores.len == pipeline_stages_flags.len);

    const signal_semaphores = [_]vk.Semaphore{
        current_swapchain_image.semaphore_render_finished,
    };

    current_swapchain_image.command_buffer = command_buffer;
    try device.queueSubmit(engine.graphics_queue.handle, 1, &.{
        .{
            .wait_semaphore_count = wait_semaphores.len,
            .p_wait_semaphores = &wait_semaphores,
            .p_wait_dst_stage_mask = &pipeline_stages_flags,
            .command_buffer_count = 1,
            .p_command_buffers = @ptrCast(&command_buffer),
            .signal_semaphore_count = signal_semaphores.len,
            .p_signal_semaphores = &signal_semaphores,
        },
    }, current_swapchain_image.fence);

    // --- PRESENT CURRENT FRAME -----------------------------------------------

    _ = try device.queuePresentKHR(mode.presentation_queue.handle, &.{
        .wait_semaphore_count = 1,
        .p_wait_semaphores = @ptrCast(&current_swapchain_image.semaphore_render_finished),
        .swapchain_count = 1,
        .p_swapchains = @ptrCast(&self.swapchain),
        .p_image_indices = @ptrCast(&self.image_index),
    });

    // --- ACQUIRE NEXT FRAME --------------------------------------------------

    const res = try device.acquireNextImageKHR(self.swapchain, std.math.maxInt(u64), self.semaphore_image_acquired, .null_handle);
    std.mem.swap(vk.Semaphore, &self.swapchain_images[res.image_index].semaphore_image_acquired, &self.semaphore_image_acquired);
    self.image_index = res.image_index;

    return switch (res.result) {
        .success => .optimal,
        .suboptimal_khr => .suboptimal,
        else => unreachable,
    };
}

fn getCurrentExtent(engine: *Engine) !vk.Extent2D {
    const mode = &engine.mode.surface;
    const surface_capabilities = try engine.instance.getPhysicalDeviceSurfaceCapabilitiesKHR(engine.physical_device, mode.surface);

    if (surface_capabilities.current_extent.width != std.math.maxInt(u32) and surface_capabilities.current_extent.height != std.math.maxInt(u32)) {
        return surface_capabilities.current_extent;
    }

    const framebuffer_width, const framebuffer_height = mode.window.getFramebufferSize();

    return .{
        .width = std.math.clamp(
            @as(u32, @intCast(framebuffer_width)),
            surface_capabilities.min_image_extent.width,
            surface_capabilities.max_image_extent.width,
        ),
        .height = std.math.clamp(
            @as(u32, @intCast(framebuffer_height)),
            surface_capabilities.min_image_extent.height,
            surface_capabilities.max_image_extent.height,
        ),
    };
}

const Params = struct {
    surface_format: vk.SurfaceFormatKHR,
    image_count: u32,

    pub fn init(engine: *Engine) !Params {
        const mode = &engine.mode.surface;
        const allocator = engine.vk_allocator.allocator;

        const surface_capabilities = try engine.instance.getPhysicalDeviceSurfaceCapabilitiesKHR(engine.physical_device, mode.surface);

        const surface_format = blk: {
            const surface_formats = try engine.instance.getPhysicalDeviceSurfaceFormatsAllocKHR(engine.physical_device, mode.surface, allocator);
            defer allocator.free(surface_formats);

            // Look for 8-bit BGRA sRGB surface format.

            for (surface_formats) |surface_format| {
                if (surface_format.format == .b8g8r8a8_srgb and surface_format.color_space == .srgb_nonlinear_khr) {
                    break :blk surface_format;
                }
            }

            if (surface_formats.len == 0) {
                return error.NoSurfaceFormats;
            }

            break :blk surface_formats[0];
        };

        const image_count = @min(
            surface_capabilities.min_image_count + 1,
            if (surface_capabilities.max_image_count > 0) surface_capabilities.max_image_count else std.math.maxInt(u32),
        );

        std.log.debug("Using surface format \"{s}\" and color space \"{s}\" with {d} images.", .{ @tagName(surface_format.format), @tagName(surface_format.color_space), image_count });

        return .{
            .surface_format = surface_format,
            .image_count = image_count,
        };
    }
};

const SwapchainImage = struct {
    image: vk.Image,
    image_view: vk.ImageView,
    semaphore_image_acquired: vk.Semaphore,
    semaphore_render_finished: vk.Semaphore,
    fence: vk.Fence,
    framebuffer: vk.Framebuffer,
    command_buffer: vk.CommandBuffer,

    const InitProps = struct {
        image: vk.Image,
        format: vk.Format,
        render_pass: vk.RenderPass,
        extent: vk.Extent2D,
    };

    fn init(engine: *Engine, props: InitProps) !SwapchainImage {
        const image_view = try engine.createImageView(.{
            .image = props.image,
            .view_type = .@"2d",
            .format = props.format,
            .subresource_range = .{
                .aspect_mask = .{ .color_bit = true },
                .base_mip_level = 0,
                .level_count = 1,
                .base_array_layer = 0,
                .layer_count = 1,
            },
        });
        errdefer engine.destroyImageView(image_view);

        const semaphore_image_acquired = try engine.createSemaphore();
        errdefer engine.destroySemaphore(semaphore_image_acquired);

        const semaphore_render_finished = try engine.createSemaphore();
        errdefer engine.destroySemaphore(semaphore_render_finished);

        const fence = try engine.createFence(.{
            .flags = .{ .signaled_bit = true },
        });
        errdefer engine.destroyFence(fence);

        const framebuffer = try engine.createFramebuffer(.{
            .render_pass = props.render_pass,
            .attachments = &.{image_view},
            .width = props.extent.width,
            .height = props.extent.height,
            .layers = 1,
        });
        errdefer engine.destroyFramebuffer(framebuffer);

        return .{
            .image = props.image,
            .image_view = image_view,
            .semaphore_image_acquired = semaphore_image_acquired,
            .semaphore_render_finished = semaphore_render_finished,
            .fence = fence,
            .framebuffer = framebuffer,
            .command_buffer = .null_handle,
        };
    }

    fn deinit(self: SwapchainImage, engine: *Engine) void {
        engine.waitForFence(self.fence) catch {};
        engine.destroyFramebuffer(self.framebuffer);
        engine.destroyFence(self.fence);
        engine.destroySemaphore(self.semaphore_render_finished);
        engine.destroySemaphore(self.semaphore_image_acquired);
        engine.destroyImageView(self.image_view);
    }
};