const std = @import("std");
const vk = @import("vk");

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

fn _choose_format(core: *const Core, alloc: std.mem.Allocator) !vk.SurfaceFormatKHR {
    const formats = try core.i.getPhysicalDeviceSurfaceFormatsAllocKHR(
        core.pdev,
        core.surface,
        alloc,
    );
    defer alloc.free(formats);

    for (formats) |format| {
        if (format.color_space == .srgb_nonlinear_khr) return format;
    } else {
        return formats[0];
    }
}

fn _choose_mode(core: *const Core, alloc: std.mem.Allocator) !vk.PresentModeKHR {
    _ = core;
    _ = alloc;

    return .fifo_khr;
}

const Self = @This();

pub const Target = struct {
    image_index: u32,
    flight_index: u32,
    image: vk.Image,
    view: vk.ImageView,
    acquired: vk.Semaphore, // this semaphore will be signaled when the target is acquired
    complete: vk.Semaphore, // this semaphore should be signaled when the render is complete
    available: vk.Fence, // this fence should be signaled when the target flight is available

    pub fn top_of_pipe_barrier(target: Target) vk.ImageMemoryBarrier {
        return 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 = target.image,
            .subresource_range = .{
                .aspect_mask = .{ .color_bit = true },
                .base_mip_level = 0,
                .level_count = 1,
                .base_array_layer = 0,
                .layer_count = 1,
            },
        };
    }

    pub fn bottom_of_pipe_barrier(target: Target) vk.ImageMemoryBarrier {
        return 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 = target.image,
            .subresource_range = .{
                .aspect_mask = .{ .color_bit = true },
                .base_mip_level = 0,
                .level_count = 1,
                .base_array_layer = 0,
                .layer_count = 1,
            },
        };
    }
};

const Sync = struct {
    acquired: vk.Semaphore = .null_handle,
    complete: vk.Semaphore = .null_handle,
    available: vk.Fence = .null_handle,
};

const View = struct {
    image: vk.Image,
    view: vk.ImageView,
};

core: *const Core,
alloc: std.mem.Allocator,

frames_in_flight: u32,
flight_index: usize = 0,
flight_syncs: std.MultiArrayList(Sync) = .{},

cinfo: vk.SwapchainCreateInfoKHR = undefined,
handle: vk.SwapchainKHR = .null_handle,

chain: std.MultiArrayList(View) = .{},

pub fn init(core: *const Core, alloc: std.mem.Allocator, frames_in_flight: usize) !Self {
    var self: Self = .{
        .core = core,
        .alloc = alloc,
        .frames_in_flight = @intCast(frames_in_flight),
    };
    errdefer self.deinit();

    try self.flight_syncs.resize(alloc, frames_in_flight);
    for (self.flight_syncs.items(.acquired)) |*sem|
        sem.* = try core.d.createSemaphore(&.{}, null);
    for (self.flight_syncs.items(.complete)) |*sem|
        sem.* = try core.d.createSemaphore(&.{}, null);
    for (self.flight_syncs.items(.available)) |*fnc|
        fnc.* = try core.d.createFence(&.{ .flags = .{ .signaled_bit = true } }, null);

    const caps = try core.i.getPhysicalDeviceSurfaceCapabilitiesKHR(core.pdev, core.surface);
    const format = try _choose_format(core, alloc);
    const mode = try _choose_mode(core, alloc);

    self.cinfo = .{
        .surface = core.surface,
        .min_image_count = std.math.clamp(
            @min(3, caps.min_image_count + 1),
            caps.min_image_count,
            if (caps.max_image_count > 0) caps.max_image_count else 127,
        ),
        .image_format = format.format,
        .image_color_space = format.color_space,
        .image_extent = undefined, // set in rebuild
        .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 = mode,
        .clipped = vk.TRUE,
        .old_swapchain = .null_handle,
    };

    // try self.rebuild();

    return self;
}

pub fn rebuild(self: *Self) !void {
    const caps = try self.core.i.getPhysicalDeviceSurfaceCapabilitiesKHR(self.core.pdev, self.core.surface);
    self.cinfo.image_extent = caps.current_extent;

    self.handle = try self.core.d.createSwapchainKHR(&self.cinfo, null);
    self.core.d.destroySwapchainKHR(self.cinfo.old_swapchain, null);
    self.cinfo.old_swapchain = self.handle;

    for (self.chain.items(.view)) |view| self.core.d.destroyImageView(view, null);
    @memset(self.chain.items(.view), .null_handle);
    @memset(self.chain.items(.image), .null_handle);

    var count: u32 = undefined;
    std.debug.assert(
        .success == try self.core.d.getSwapchainImagesKHR(
            self.handle,
            &count,
            null,
        ),
    );
    try self.chain.resize(self.alloc, count);
    std.debug.assert(
        .success == try self.core.d.getSwapchainImagesKHR(
            self.handle,
            &count,
            self.chain.items(.image).ptr,
        ),
    );

    for (self.chain.items(.image), self.chain.items(.view)) |image, *view| {
        view.* = try self.core.d.createImageView(&vk.ImageViewCreateInfo{
            .image = image,
            .view_type = .@"2d",
            .format = self.cinfo.image_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);
    }
}

pub fn deinit(self: *Self) void {
    // The easiest way to ensure fences and semaphores are not in use for deletion. If it fails, the device must
    // have failed somehow and sync is not necessary so just continue with cleanup.
    self.core.d.deviceWaitIdle() catch {};

    // images are owned by swapchain and not explicitly destroyed
    for (self.chain.items(.view)) |view| self.core.d.destroyImageView(view, null);
    self.chain.deinit(self.alloc);

    self.core.d.destroySwapchainKHR(self.handle, null);

    for (self.flight_syncs.items(.acquired)) |sem| self.core.d.destroySemaphore(sem, null);
    for (self.flight_syncs.items(.complete)) |sem| self.core.d.destroySemaphore(sem, null);
    for (self.flight_syncs.items(.available)) |fnc| self.core.d.destroyFence(fnc, null);
    self.flight_syncs.deinit(self.alloc);
}

pub fn acquire(self: *Self) !?Target {
    var target: Target = .{
        .flight_index = @intCast(self.flight_index),
        .acquired = self.flight_syncs.items(.acquired)[self.flight_index],
        .complete = self.flight_syncs.items(.complete)[self.flight_index],
        .available = self.flight_syncs.items(.available)[self.flight_index],
        .image = undefined,
        .view = undefined,
        .image_index = undefined,
    };

    const timeout = std.math.maxInt(u64);

    std.debug.assert(.success == try self.core.d.waitForFences(
        1,
        &.{target.available},
        vk.TRUE,
        std.math.maxInt(u64),
    ));

    // two attempts

    target.image_index = for (0..2) |_| {
        if (self.handle == .null_handle) try self.rebuild();

        if (self.core.d.acquireNextImageKHR(
            self.handle,
            timeout,
            target.acquired,
            .null_handle,
        )) |res| switch (res.result) {
            .success, .suboptimal_khr => break res.image_index,
            else => unreachable,
        } else |err| switch (err) {
            error.OutOfDateKHR => {
                self.handle = .null_handle;
            },
            else => return err,
        }
    } else {
        return null;
    };

    target.image = self.chain.items(.image)[target.image_index];
    target.view = self.chain.items(.view)[target.image_index];

    try self.core.d.resetFences(1, &.{target.available});
    self.flight_index = @mod(self.flight_index + 1, self.flight_syncs.len);

    return target;
}

pub fn present(self: *Self, target: Target) !void {
    if (self.core.q.presentKHR(&vk.PresentInfoKHR{
        .wait_semaphore_count = 1,
        .p_wait_semaphores = &.{target.complete},
        .swapchain_count = 1,
        .p_swapchains = &.{self.handle},
        .p_image_indices = &.{target.image_index},
        .p_results = null,
    })) |res| switch (res) {
        .success => {},
        .suboptimal_khr => self.handle = .null_handle,
        else => unreachable,
    } else |err| switch (err) {
        error.OutOfDateKHR => self.handle = .null_handle,
        else => return err,
    }
}