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

const vk = @import("vulkan");

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

engine: *Engine,
params: Params,
render_pass: vk.RenderPass,
swapchain: vk.SwapchainKHR = .null_handle,
swapchain_images: []SwapchainImage = &.{},

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 = .dont_care,
                .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 = .{
        .engine = engine,
        .params = params,
        .render_pass = render_pass,
    };

    try recreate(&swapchain);
    return swapchain;
}

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

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

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

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

    self.* = undefined;
}

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

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

    const extent = try getCurrentExtent(self.engine);

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

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

    const qsm = QSM.resolve(self.engine.graphics_queue.allocation.family, mode.presentation_queue.allocation.family);
    const new_swapchain = try self.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,
    }, &self.engine.vk_allocator.interface);
    errdefer self.engine.device.destroySwapchainKHR(new_swapchain, &self.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(self.engine);
    }
    allocator.free(self.swapchain_images);
    self.swapchain_images = &.{};

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

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

    const new_swapchain_images = blk: {
        const images = try self.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(self.engine);

        for (images) |image| {
            swapchain_images.appendAssumeCapacity(try SwapchainImage.init(self.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(self.engine);
        }
        allocator.free(new_swapchain_images);
    }

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

    self.swapchain = new_swapchain;
    self.swapchain_images = new_swapchain_images;
}

fn getCurrentExtent(self: *Engine) !vk.Extent2D {
    const mode = &self.mode.surface;
    const surface_capabilities = try self.instance.getPhysicalDeviceSurfaceCapabilitiesKHR(self.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),
        );

        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,

    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.device.createImageView(&.{
            .image = props.image,
            .view_type = .@"2d",
            .format = props.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,
            },
        }, &engine.vk_allocator.interface);
        errdefer engine.device.destroyImageView(image_view, &engine.vk_allocator.interface);

        const semaphore_image_acquired = try engine.device.createSemaphore(&.{}, &engine.vk_allocator.interface);
        errdefer engine.device.destroySemaphore(semaphore_image_acquired, &engine.vk_allocator.interface);

        const semaphore_render_finished = try engine.device.createSemaphore(&.{}, &engine.vk_allocator.interface);
        errdefer engine.device.destroySemaphore(semaphore_render_finished, &engine.vk_allocator.interface);

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

        const attachments = [_]vk.ImageView{image_view};

        const framebuffer = try engine.device.createFramebuffer(&.{
            .render_pass = props.render_pass,
            .attachment_count = attachments.len,
            .p_attachments = &attachments,
            .width = props.extent.width,
            .height = props.extent.height,
            .layers = 1,
        }, &engine.vk_allocator.interface);
        errdefer engine.device.destroyFramebuffer(framebuffer, &engine.vk_allocator.interface);

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

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

    fn waitForFence(self: SwapchainImage, engine: *Engine) !void {
        const fences = [_]vk.Fence{self.fence};
        _ = try engine.device.waitForFences(fences.len, &fences, .true, std.math.maxInt(u64));
    }
};