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

const ctx = @import("../AppContext.zig");
const vk = @import("vulkan");

const CommandBuffer = @import("CommandBuffer.zig");
const Engine = @import("Engine.zig");
const Texture = @import("Texture.zig");
const WaitSemaphore = @import("WaitSemaphore.zig");

params: Params,

extent: vk.Extent2D = .{ .width = 0, .height = 0 },
swapchain: vk.SwapchainKHR = .null_handle,
/// Allocated with `allocator_general`.
swapchain_images: []SwapchainImage = &.{},
depth_texture: ?Texture = null,
image_index: ?u32 = null,
fence: vk.Fence = .null_handle,
semaphore_image_acquired: vk.Semaphore = .null_handle,

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

pub fn init() !*Swapchain {
    const allocator_persistent = ctx.allocator_persistent;

    const swapchain = try allocator_persistent.create(Swapchain);
    errdefer allocator_persistent.destroy(swapchain);

    swapchain.* = .{
        .params = try .init(),
    };

    try swapchain.recreate();
    return swapchain;
}

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

    std.log.scoped(.deinit).debug("Deinitializing {*}", .{self});

    for (self.swapchain_images) |*swapchain_image| {
        swapchain_image.deinit();
    }
    allocator_general.free(self.swapchain_images);

    if (self.depth_texture) |*depth_texture| {
        depth_texture.deinit();
    }

    if (self.swapchain != .null_handle) {
        engine.destroySwapchain(self.swapchain);
    }

    engine.destroySemaphore(self.semaphore_image_acquired);
    engine.destroyFence(self.fence);

    self.* = undefined;
}

pub fn recreate(self: *Swapchain) !void {
    const allocator_general = ctx.allocator_general;
    const allocator_frame = ctx.allocator_frame;
    const engine = ctx.engine;

    // TODO This is very "LMAO just wait for all" way to synchronize. It might
    // be the only viable option, though. Either way, we should research what
    // is "the way".
    try engine.deviceWaitIdle();

    const old_swapchain = self.swapchain;
    const old_swapchain_images = self.swapchain_images;
    const old_depth_texture = &self.depth_texture;
    const old_semaphore_image_acquired = self.semaphore_image_acquired;
    const old_fence = self.fence;

    const extent = try getCurrentExtent();

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

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

    const surface_capabilities = try engine.getPhysicalDeviceSurfaceCapabilities();
    const new_swapchain = try engine.createSwapchain(.{
        .surface = engine.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 },
        .queue_family_indices = &.{
            engine.graphics_queue.allocation.family,
            engine.presentation_queue.allocation.family,
        },
        .pre_transform = surface_capabilities.current_transform,
        .composite_alpha = .{ .opaque_bit_khr = true },
        .present_mode = .fifo_khr,
        .clipped = .true,
        .old_swapchain = old_swapchain,
    });
    errdefer engine.destroySwapchain(new_swapchain);
    engine.setObjectName(new_swapchain, "SC {d}×{d}", .{ extent.width, extent.height });

    // --- 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();
    }
    allocator_general.free(self.swapchain_images);
    self.swapchain_images = &.{};

    if (old_depth_texture.*) |*depth_texture| {
        depth_texture.deinit();
    }

    if (old_swapchain != .null_handle) {
        engine.destroySwapchain(old_swapchain);
        self.swapchain = .null_handle;
    }

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

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

    if (old_fence != .null_handle) {
        engine.destroyFence(old_fence);
        self.fence = .null_handle;
    }

    // --- CREATE DEPTH TEXTURE ------------------------------------------------

    var new_depth_texture = try Texture.init(.{
        .width = extent.width,
        .height = extent.height,
        .target_queue = .graphics,
        .usage = .depth,
        .name = "@Depth",
    });
    errdefer new_depth_texture.deinit();

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

    const new_swapchain_images = blk: {
        const images = try engine.getSwapchainImagesAlloc(new_swapchain, allocator_frame);

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

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

        for (images, 0..) |image, index| {
            swapchain_images.appendAssumeCapacity(try SwapchainImage.init(.{
                .image = image,
                .format = self.params.surface_format.format,
                .extent = extent,
                .depth_stencil_image_view = new_depth_texture.image_view,
                .index = index,
            }));
        }

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

    // --- CREATE SEMAPHORES AND FENCES ----------------------------------------

    const semaphore_image_acquired = try engine.createSemaphore();
    errdefer engine.destroySemaphore(semaphore_image_acquired);
    engine.setObjectName(semaphore_image_acquired, "S Acquired[E]", .{});

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

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

    self.extent = extent;
    self.swapchain = new_swapchain;
    self.swapchain_images = new_swapchain_images;
    self.depth_texture = new_depth_texture;
    self.image_index = null;
    self.semaphore_image_acquired = semaphore_image_acquired;
    self.fence = fence;
}

pub fn acquire(self: *Swapchain) !void {
    const engine = ctx.engine;

    try engine.waitForFence(self.fence);
    try engine.resetFence(self.fence);

    const res = try engine.acquireNextImage(self.swapchain, .{
        .semaphore = self.semaphore_image_acquired,
    });
    std.mem.swap(vk.Semaphore, &self.swapchain_images[res.image_index].semaphore_image_acquired, &self.semaphore_image_acquired);
    self.image_index = res.image_index;
}

pub fn getCurrentSwapchainImage(self: *Swapchain) *const SwapchainImage {
    return &self.swapchain_images[self.image_index.?];
}

pub const PresentInfo = struct {
    command_buffer: CommandBuffer,
    wait_semaphores: []const WaitSemaphore = &.{},
};

pub fn present(self: *Swapchain, present_info: PresentInfo) !void {
    const allocator_frame = ctx.allocator_frame;
    const engine = ctx.engine;

    const current_swapchain_image = self.getCurrentSwapchainImage();

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

    var wait_semaphores: std.ArrayList(WaitSemaphore) = try .initCapacity(allocator_frame, 1 + present_info.wait_semaphores.len);

    wait_semaphores.appendAssumeCapacity(.{
        .semaphore = current_swapchain_image.semaphore_image_acquired,
        .stage_flags = .{ .top_of_pipe_bit = true },
    });
    wait_semaphores.appendSliceAssumeCapacity(present_info.wait_semaphores);

    try present_info.command_buffer.submit(.{
        .wait_semaphores = wait_semaphores.items,
        .signal_semaphores = &.{current_swapchain_image.semaphore_render_finished},
        .fence = self.fence,
    });

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

    _ = try engine.queuePresent(.{
        .wait_semaphores = &.{current_swapchain_image.semaphore_render_finished},
        .swapchain = self.swapchain,
        .image_index = self.image_index.?,
    });
}

fn getCurrentExtent() !vk.Extent2D {
    const window = ctx.window;
    const engine = ctx.engine;

    const surface_capabilities = try engine.getPhysicalDeviceSurfaceCapabilities();

    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 = 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() !Params {
        const allocator_frame = ctx.allocator_frame;
        const engine = ctx.engine;

        const surface_capabilities = try engine.getPhysicalDeviceSurfaceCapabilities();

        const surface_format = blk: {
            const surface_formats = try engine.getPhysicalDeviceSurfaceFormatsAlloc(allocator_frame);

            // 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,

    const InitProps = struct {
        image: vk.Image,
        format: vk.Format,
        extent: vk.Extent2D,
        depth_stencil_image_view: vk.ImageView,

        index: usize,
    };

    fn init(props: InitProps) !SwapchainImage {
        const engine = ctx.engine;

        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);
        engine.setObjectName(image_view, "IV Swapchain[{d}]", .{props.index});

        const semaphore_image_acquired = try engine.createSemaphore();
        errdefer engine.destroySemaphore(semaphore_image_acquired);
        engine.setObjectName(image_view, "S Acquired[{d}]", .{props.index});

        const semaphore_render_finished = try engine.createSemaphore();
        errdefer engine.destroySemaphore(semaphore_render_finished);
        engine.setObjectName(image_view, "S Rendered[{d}]", .{props.index});

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

    fn deinit(self: *SwapchainImage) void {
        const engine = ctx.engine;

        std.log.scoped(.deinit).debug("Deinitializing {*}", .{self});

        engine.destroySemaphore(self.semaphore_render_finished);
        engine.destroySemaphore(self.semaphore_image_acquired);
        engine.destroyImageView(self.image_view);

        self.* = undefined;
    }
};