//! ----------------------------------------------------
//! `🗲` Vulkan Swapchain `🗲`
//! ----------------------------------------------------

const vk = @import("vulkan");
const hp = @import("../helper.zig");
const std = @import("std");
const gtl = @import("gtl");
const cmn = @import("../../common.zig");
const Self = @This();

const Adapter = @import("adapter.zig");
const Device = @import("device.zig");
const Surface = @import("surface.zig");

//
// FIELDS
//

raw: *vk.VkSwapchainKHR_T,

images: []vk.VkImage,
image_views: []vk.VkImageView,

depth_images: []vk.VkImage = &.{},
depth_memories: []vk.VkDeviceMemory = &.{},
depth_views: []vk.VkImageView = &.{},
depth_format: vk.VkFormat = vk.VK_FORMAT_UNDEFINED,

format: vk.VkFormat,
extent: vk.VkExtent2D,

device: *const Device,
alloc: std.mem.Allocator,

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn init(
    adapter: *const Adapter,
    device: *const Device,
    surface: *const Surface,
    config: cmn.SwapchainConfig,
    alloc: std.mem.Allocator,
) !Self {
    // --- # ---
    var capabilities: vk.VkSurfaceCapabilitiesKHR = .{};
    _ = vk.vkGetPhysicalDeviceSurfaceCapabilitiesKHR(adapter.raw, surface.raw, &capabilities);

    // --- # ---
    var format_count: u32 = 0;
    _ = vk.vkGetPhysicalDeviceSurfaceFormatsKHR(adapter.raw, surface.raw, &format_count, null);

    const formats = try alloc.alloc(vk.VkSurfaceFormatKHR, format_count);
    defer alloc.free(formats);
    _ = vk.vkGetPhysicalDeviceSurfaceFormatsKHR(adapter.raw, surface.raw, &format_count, formats.ptr);

    // --- # ---
    var present_mode_count: u32 = 0;
    _ = vk.vkGetPhysicalDeviceSurfacePresentModesKHR(adapter.raw, surface.raw, &present_mode_count, null);

    const present_modes = try alloc.alloc(vk.VkPresentModeKHR, present_mode_count);
    defer alloc.free(present_modes);
    _ = vk.vkGetPhysicalDeviceSurfacePresentModesKHR(adapter.raw, surface.raw, &present_mode_count, present_modes.ptr);

    // --- # ---
    const extent: vk.VkExtent2D = blk: {
        if (capabilities.currentExtent.width != std.math.maxInt(u32)) break :blk capabilities.currentExtent;
        if (config.size) |size| break :blk .{ .width = size[0], .height = size[1] };
        break :blk .{ .width = 0, .height = 0 };
    };

    // --- # ---
    const format: vk.VkSurfaceFormatKHR = blk: {
        for (formats) |v| {
            // if (v.format == vk.VK_FORMAT_B8G8R8A8_SRGB and v.colorSpace == vk.VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
            //     break :blk v;
            // }
            if (v.format == hp.toRawPixelFormat(config.format)) {
                break :blk v;
            }
        }
        break :blk formats[0];
    };

    // --- # ---
    var image_count: u32 = capabilities.minImageCount + 1;
    if (capabilities.maxImageCount > 0 and image_count > capabilities.maxImageCount) {
        image_count = capabilities.maxImageCount;
    }

    var create_info: vk.VkSwapchainCreateInfoKHR = .{
        .sType = vk.VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .surface = surface.raw,
        .minImageCount = image_count,
        .imageFormat = format.format,
        .imageColorSpace = format.colorSpace,
        .imageExtent = extent,
        .imageArrayLayers = 1,
        .imageUsage = vk.VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,

        .preTransform = capabilities.currentTransform,
        .compositeAlpha = vk.VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = present_modes[0],
        .clipped = 1,

        .imageSharingMode = vk.VK_SHARING_MODE_EXCLUSIVE,
        .queueFamilyIndexCount = 0,
        .pQueueFamilyIndices = null,
    };

    var swapchain: vk.VkSwapchainKHR = null;
    if (!hp.check(vk.vkCreateSwapchainKHR(device.raw, &create_info, null, &swapchain))) {
        return error.FailedToCreateSwapchainKHR;
    }

    // --- GET IMAGES ---
    _ = vk.vkGetSwapchainImagesKHR(device.raw, swapchain, &image_count, null);
    const images = try alloc.alloc(vk.VkImage, image_count);
    _ = vk.vkGetSwapchainImagesKHR(device.raw, swapchain, &image_count, images.ptr);

    // --- GET IMAGE VIEWS ---
    const image_views = try alloc.alloc(vk.VkImageView, images.len);
    for (images, 0..) |image, i| {
        var image_view_create_info: vk.VkImageViewCreateInfo = .{
            .sType = vk.VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = image,
            .viewType = vk.VK_IMAGE_VIEW_TYPE_2D,
            .format = format.format,
            .components = .{
                .r = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .g = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .b = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .a = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
            },
            .subresourceRange = .{
                .aspectMask = vk.VK_IMAGE_ASPECT_COLOR_BIT,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
        };

        if (vk.vkCreateImageView(device.raw, &image_view_create_info, null, &image_views[i]) != vk.VK_SUCCESS) {
            return error.FailedToCreateImageView;
        }
    }

    return .{
        .raw = swapchain.?,
        .images = images,
        .image_views = image_views,

        .extent = extent,
        .format = format.format,

        .device = device,
        .alloc = alloc,
    };
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn deinit(self: *Self) void {
    self.destroyDepth();

    for (self.image_views) |view| {
        vk.vkDestroyImageView(self.device.raw, view, null);
    }

    vk.vkDestroySwapchainKHR(self.device.raw, self.raw, null);

    self.alloc.free(self.images);
    self.alloc.free(self.image_views);
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn ensureDepth(self: *Self, format: vk.VkFormat) !void {
    if (self.depth_format == format and self.depth_views.len > 0) return;

    self.destroyDepth();

    const count = self.images.len;
    const images = try self.alloc.alloc(vk.VkImage, count);
    errdefer self.alloc.free(images);
    const memories = try self.alloc.alloc(vk.VkDeviceMemory, count);
    errdefer self.alloc.free(memories);
    const views = try self.alloc.alloc(vk.VkImageView, count);

    var created: usize = 0;
    errdefer {
        for (0..created) |i| {
            if (views[i]) |view| vk.vkDestroyImageView(self.device.raw, view, null);
            if (memories[i]) |memory| vk.vkFreeMemory(self.device.raw, memory, null);
            if (images[i]) |image| vk.vkDestroyImage(self.device.raw, image, null);
        }
        self.alloc.free(views);
    }

    const aspect: vk.VkImageAspectFlags = if (isDepthStencilFormat(format))
        vk.VK_IMAGE_ASPECT_DEPTH_BIT | vk.VK_IMAGE_ASPECT_STENCIL_BIT
    else
        vk.VK_IMAGE_ASPECT_DEPTH_BIT;

    const image_info: vk.VkImageCreateInfo = .{
        .sType = vk.VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        .imageType = vk.VK_IMAGE_TYPE_2D,
        .extent = .{
            .width = self.extent.width,
            .height = self.extent.height,
            .depth = 1,
        },
        .mipLevels = 1,
        .arrayLayers = 1,
        .format = format,
        .tiling = vk.VK_IMAGE_TILING_OPTIMAL,
        .initialLayout = vk.VK_IMAGE_LAYOUT_UNDEFINED,
        .usage = vk.VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
        .samples = vk.VK_SAMPLE_COUNT_1_BIT,
        .sharingMode = vk.VK_SHARING_MODE_EXCLUSIVE,
    };

    for (0..count) |i| {
        created = i;

        if (vk.vkCreateImage(self.device.raw, &image_info, null, &images[i]) != vk.VK_SUCCESS) {
            return error.FailedToCreateDepthImage;
        }

        var mem_req: vk.VkMemoryRequirements = .{};
        vk.vkGetImageMemoryRequirements(self.device.raw, images[i], &mem_req);

        const mem_type = findMemoryType(
            self.device.physical_device,
            mem_req.memoryTypeBits,
            vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        ) orelse return error.FailedToFindSuitableMemoryType;

        const alloc_info: vk.VkMemoryAllocateInfo = .{
            .sType = vk.VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
            .allocationSize = mem_req.size,
            .memoryTypeIndex = mem_type,
        };

        if (vk.vkAllocateMemory(self.device.raw, &alloc_info, null, &memories[i]) != vk.VK_SUCCESS) {
            return error.FailedToAllocateDepthMemory;
        }
        _ = vk.vkBindImageMemory(self.device.raw, images[i], memories[i], 0);

        const view_info: vk.VkImageViewCreateInfo = .{
            .sType = vk.VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = images[i],
            .viewType = vk.VK_IMAGE_VIEW_TYPE_2D,
            .format = format,
            .components = .{
                .r = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .g = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .b = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
                .a = vk.VK_COMPONENT_SWIZZLE_IDENTITY,
            },
            .subresourceRange = .{
                .aspectMask = aspect,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
        };
        if (vk.vkCreateImageView(self.device.raw, &view_info, null, &views[i]) != vk.VK_SUCCESS) {
            return error.FailedToCreateDepthImageView;
        }

        created = i + 1;
    }

    self.depth_images = images;
    self.depth_memories = memories;
    self.depth_views = views;
    self.depth_format = format;
}

/// ----------------------------------------------------
/// ----------------------------------------------------
fn destroyDepth(self: *Self) void {
    if (self.depth_views.len > 0) {
        for (self.depth_views) |view| {
            if (view) |v| vk.vkDestroyImageView(self.device.raw, v, null);
        }
        self.alloc.free(self.depth_views);
        self.depth_views = &.{};
    }
    if (self.depth_memories.len > 0) {
        for (self.depth_memories) |memory| {
            if (memory) |m| vk.vkFreeMemory(self.device.raw, m, null);
        }
        self.alloc.free(self.depth_memories);
        self.depth_memories = &.{};
    }
    if (self.depth_images.len > 0) {
        for (self.depth_images) |image| {
            if (image) |im| vk.vkDestroyImage(self.device.raw, im, null);
        }
        self.alloc.free(self.depth_images);
        self.depth_images = &.{};
    }
    self.depth_format = vk.VK_FORMAT_UNDEFINED;
}

/// ----------------------------------------------------
/// ----------------------------------------------------
fn isDepthStencilFormat(format: vk.VkFormat) bool {
    return format == vk.VK_FORMAT_D16_UNORM_S8_UINT or
        format == vk.VK_FORMAT_D24_UNORM_S8_UINT or
        format == vk.VK_FORMAT_D32_SFLOAT_S8_UINT;
}

/// ----------------------------------------------------
/// ----------------------------------------------------
fn findMemoryType(
    physical_device: vk.VkPhysicalDevice,
    type_bits: u32,
    properties: vk.VkMemoryPropertyFlags,
) ?u32 {
    var mem_props: vk.VkPhysicalDeviceMemoryProperties = .{};
    vk.vkGetPhysicalDeviceMemoryProperties(physical_device, &mem_props);

    for (0..mem_props.memoryTypeCount) |i| {
        if ((type_bits & (@as(u32, 1) << @intCast(i))) != 0 and
            (mem_props.memoryTypes[i].propertyFlags & properties) == properties)
        {
            return @intCast(i);
        }
    }

    return null;
}