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

const Context = @import("Context.zig");
const Instance = Context.Instance;
const QueueUtils = @import("queue_utils.zig");
const Utilities = @import("utilities.zig");

pub const SwapchainDetails = struct {
    surface_capabilities: vk.SurfaceCapabilitiesKHR,
    formats: []vk.SurfaceFormatKHR,
    presentation_modes: []vk.PresentModeKHR,
};

pub const SwapchainImage = struct {
    image: vk.Image,
    image_view: vk.ImageView,
};

const Self = @This();

allocator: std.mem.Allocator,

ctx: Context,

handle: vk.SwapchainKHR,

swapchain_images: []SwapchainImage,
swapchain_framebuffers: []vk.Framebuffer,

swapchain_image_format: vk.Format,
extent: vk.Extent2D,

pub fn create(allocator: std.mem.Allocator, context: Context) !Self {
    var self: Self = undefined;

    self.allocator = allocator;
    self.ctx = context;

    const swapchain_details = try getSwapchainDetails(allocator, context.instance, context.physical_device, context.surface);
    defer self.allocator.free(swapchain_details.formats);
    defer self.allocator.free(swapchain_details.presentation_modes);

    // 1. Choose best surface format
    const surface_format = chooseBestSurfaceFormat(swapchain_details.formats);
    // 2. Choose best presentation mode
    const present_mode = chooseBestPresentationMode(swapchain_details.presentation_modes);
    // 3. Choose swapchain image resolution
    const extent = chooseSwapExtent(&self.ctx.window, swapchain_details.surface_capabilities);

    // How many images are in the swapchain? Get 1 more than the minimum to allow triple buffering
    var image_count: u32 = swapchain_details.surface_capabilities.min_image_count + 1;
    const max_image_count = swapchain_details.surface_capabilities.max_image_count;

    // Clamp down if higher
    // If 0, it means it's limitless
    if (max_image_count != 0 and image_count > max_image_count) {
        image_count = max_image_count;
    }

    var swapchain_create_info: vk.SwapchainCreateInfoKHR = .{
        .image_format = surface_format.format,
        .image_color_space = surface_format.color_space,
        .present_mode = present_mode,
        .image_extent = extent,
        .min_image_count = image_count,
        .image_array_layers = 1, // Number of layers for each image
        .image_usage = .{ .color_attachment_bit = true }, // What attachment will images be used as
        .pre_transform = swapchain_details.surface_capabilities.current_transform, // Transform to perform on swapchain images
        .composite_alpha = .{ .opaque_bit_khr = true }, // How to handle blending images with external graphics (e.g.: other windows)
        .clipped = vk.TRUE, // Whether to clip parts of images not in view (e.g.: behind another window, off-screen, etc...)
        .old_swapchain = .null_handle, // Links old one to quickly share responsibilities in case it's been destroyed and replaced
        .surface = context.surface,
        .image_sharing_mode = .exclusive,
    };

    // Get queue family indices
    const family_indices = try QueueUtils.getQueueFamilies(self.ctx, self.ctx.physical_device);

    // If graphic and presentation families are different, then swapchain must let images be shared between families

    if (family_indices.graphics_family.? != family_indices.presentation_family.?) {
        const qfi = [_]u32{
            family_indices.graphics_family.?,
            family_indices.presentation_family.?,
        };

        swapchain_create_info.image_sharing_mode = .concurrent;
        swapchain_create_info.queue_family_index_count = @intCast(qfi.len); // Number of queues to share images between
        swapchain_create_info.p_queue_family_indices = &qfi;
    }

    self.handle = try self.ctx.device.createSwapchainKHR(&swapchain_create_info, null);
    self.swapchain_image_format = surface_format.format;
    self.extent = extent;

    // Swapchain images
    var swapchain_image_count: u32 = 0;
    _ = try self.ctx.device.getSwapchainImagesKHR(self.handle, &swapchain_image_count, null);

    const images = try self.allocator.alloc(vk.Image, swapchain_image_count);
    defer self.allocator.free(images);

    _ = try self.ctx.device.getSwapchainImagesKHR(self.handle, &swapchain_image_count, images.ptr);

    self.swapchain_images = try self.allocator.alloc(SwapchainImage, swapchain_image_count);

    for (images, 0..) |image, i| {
        self.swapchain_images[i] = .{
            .image = image,
            .image_view = try self.createImageView(image, self.swapchain_image_format, .{ .color_bit = true }),
        };
    }

    return self;
}

pub fn createImage(
    self: *Self,
    width: u32,
    height: u32,
    format: vk.Format,
    tiling: vk.ImageTiling,
    use_flags: vk.ImageUsageFlags,
    prop_flags: vk.MemoryPropertyFlags,
    image_memory: *vk.DeviceMemory,
) !vk.Image {
    // -- Create Image --
    const image_create_info: vk.ImageCreateInfo = .{
        .image_type = .@"2d", // Type of image (1D, 2D or 3D)
        .extent = .{
            .width = width, // Width of image extent
            .height = height, // Height of image extent
            .depth = 1, // Depth of image (just 1, no 3D aspecct)
        },
        .mip_levels = 1, // Number of mipmap levels
        .array_layers = 1, // Number of level in image array
        .format = format, // Format type of image
        .tiling = tiling, // How image data should be tiled (arranged for optimal reading)
        .initial_layout = .undefined, // Layout of image data on creation
        .usage = use_flags, // Bit flags defining what image will be used for
        .samples = .{ .@"1_bit" = true }, // Number of samples for multi-sampling
        .sharing_mode = .exclusive, // Whether image can be shared between queues
    };

    const image = try self.ctx.device.createImage(&image_create_info, null);

    // -- Create memory for image --
    // Get memory requirements for a type of image
    const memory_requirements = self.ctx.device.getImageMemoryRequirements(image);

    // Allocate memory using image requirements and user-defined properties
    const memory_alloc_info: vk.MemoryAllocateInfo = .{
        .allocation_size = memory_requirements.size,
        .memory_type_index = Utilities.findMemoryTypeIndex(self.ctx.physical_device, self.ctx.instance, memory_requirements.memory_type_bits, prop_flags),
    };

    image_memory.* = try self.ctx.device.allocateMemory(&memory_alloc_info, null);

    // Connect memory to image
    try self.ctx.device.bindImageMemory(image, image_memory.*, 0);

    return image;
}

pub fn createImageView(self: Self, image: vk.Image, format: vk.Format, aspect_flags: vk.ImageAspectFlags) !vk.ImageView {
    const image_view_create_info: vk.ImageViewCreateInfo = .{
        .image = image,
        .format = format,
        .view_type = .@"2d",
        .components = .{
            // Used for remapping rgba values to other rgba values
            .r = .identity,
            .g = .identity,
            .b = .identity,
            .a = .identity,
        },
        .subresource_range = .{
            .aspect_mask = aspect_flags, // Which aspect of image to view (e.g.: colour, depth, stencil, etc...)
            .base_mip_level = 0, // Start mipmap level to view from
            .level_count = 1, // Number of mipmap levels to view
            .base_array_layer = 0, // Start array level to view from
            .layer_count = 1, // Number of array levels to view
        },
    };

    return try self.ctx.device.createImageView(&image_view_create_info, null);
}

pub fn getSwapchainDetails(allocator: std.mem.Allocator, instance: Instance, pdev: vk.PhysicalDevice, surface: vk.SurfaceKHR) !SwapchainDetails {
    // Capabilities
    const surface_capabilities = try instance.getPhysicalDeviceSurfaceCapabilitiesKHR(pdev, surface);

    // Formats
    var format_count: u32 = 0;
    _ = try instance.getPhysicalDeviceSurfaceFormatsKHR(pdev, surface, &format_count, null);

    const formats = try allocator.alloc(vk.SurfaceFormatKHR, format_count);
    _ = try instance.getPhysicalDeviceSurfaceFormatsKHR(pdev, surface, &format_count, formats.ptr);

    // Presentation modes
    var present_mode_count: u32 = 0;
    _ = try instance.getPhysicalDeviceSurfacePresentModesKHR(pdev, surface, &present_mode_count, null);

    const presentation_modes = try allocator.alloc(vk.PresentModeKHR, format_count);
    _ = try instance.getPhysicalDeviceSurfacePresentModesKHR(pdev, surface, &present_mode_count, presentation_modes.ptr);

    return .{
        .surface_capabilities = surface_capabilities,
        .formats = formats,
        .presentation_modes = presentation_modes,
    };
}

// Format: VK_FORMAT_R8G8B8A8_UNORM (VK_FORMAT_B8G8R8A8_UNORM as backup)
// Color space: VK_COLOR_SPACE_SRGB_NONLINEAR_KHR
fn chooseBestSurfaceFormat(formats: []vk.SurfaceFormatKHR) vk.SurfaceFormatKHR {
    // If only one format available and is undefined, then this means all formats are available
    if (formats.len == 1 and formats[0].format == vk.Format.undefined) {
        return .{
            .format = vk.Format.r8g8b8a8_srgb,
            .color_space = vk.ColorSpaceKHR.srgb_nonlinear_khr,
        };
    }

    for (formats) |format| {
        if ((format.format == vk.Format.r8g8b8a8_srgb or format.format == vk.Format.b8g8r8a8_srgb) and format.color_space == vk.ColorSpaceKHR.srgb_nonlinear_khr) {
            return format;
        }
    }

    return formats[0];
}

fn chooseBestPresentationMode(presentation_modes: []vk.PresentModeKHR) vk.PresentModeKHR {
    for (presentation_modes) |presentation_mode| {
        if (presentation_mode == vk.PresentModeKHR.mailbox_khr) {
            return presentation_mode;
        }
    }

    // Use FIFO as Vulkan spec says it must be present
    return vk.PresentModeKHR.fifo_khr;
}

fn chooseSwapExtent(window: *sdl.Window, surface_capabilities: vk.SurfaceCapabilitiesKHR) vk.Extent2D {
    // If the current extent is at max value, the extent can vary. Otherwise it's the size of the window
    if (surface_capabilities.current_extent.width != std.math.maxInt(u32)) {
        return surface_capabilities.current_extent;
    }

    // If value can very, need to set the extent manually
    const framebuffer_size = sdl.vulkan.getDrawableSize(window);

    var extent: vk.Extent2D = .{
        .width = @intCast(framebuffer_size.width),
        .height = @intCast(framebuffer_size.height),
    };

    // Surface also defines max and min, so make sure it's within boundaries by clamping values
    extent.width = @max(surface_capabilities.min_image_extent.width, @min(surface_capabilities.max_image_extent.width, extent.width));
    extent.height = @max(surface_capabilities.min_image_extent.height, @min(surface_capabilities.max_image_extent.height, extent.height));

    return extent;
}
