const std = @import("std");
const sdl = @import("sdl2");
const vk = @import("vulkan");
const Utilities = @import("utilities.zig");
const QueueFamilyIndices = Utilities.QueueFamilyIndices;
const SwapchainDetails = Utilities.SwapchainDetails;

const enable_validation_layers = true;
const validation_layers = [_][*:0]const u8{"VK_LAYER_KHRONOS_validation"};

const apis: []const vk.ApiInfo = &.{
    vk.features.version_1_0,
    vk.features.version_1_1,
    vk.features.version_1_2,
    vk.features.version_1_3,
    vk.extensions.khr_surface,
    vk.extensions.khr_swapchain,
    vk.extensions.ext_debug_utils,
};

const BaseDispatch = vk.BaseWrapper(apis);
const InstanceDispatch = vk.InstanceWrapper(apis);
const DeviceDispatch = vk.DeviceWrapper(apis);

const Instance = vk.InstanceProxy(apis);
const Device = vk.DeviceProxy(apis);
const Queue = vk.QueueProxy(apis);

pub const VulkanRenderer = struct {
    const Self = @This();

    allocator: std.mem.Allocator,

    vkb: BaseDispatch,

    window: sdl.Window,

    instance: Instance,
    physical_device: vk.PhysicalDevice,
    device: Device,

    graphics_queue: Queue,
    presentation_queue: Queue,

    surface: vk.SurfaceKHR,

    swapchain: vk.SwapchainKHR,

    debug_utils: ?vk.DebugUtilsMessengerEXT,

    pub fn init(window: sdl.Window, allocator: std.mem.Allocator) !Self {
        var self: Self = undefined;
        self.window = window;
        self.allocator = allocator;

        self.vkb = try BaseDispatch.load(try sdl.vulkan.getVkGetInstanceProcAddr());

        const instance_handle = try self.createInstance();

        const vki = try allocator.create(InstanceDispatch);
        errdefer allocator.destroy(vki);
        vki.* = try InstanceDispatch.load(instance_handle, self.vkb.dispatch.vkGetInstanceProcAddr);

        self.instance = Instance.init(instance_handle, vki);
        self.surface = try sdl.vulkan.createSurface(self.window, self.instance.handle);

        if (enable_validation_layers) {
            self.debug_utils = try createDebugMessenger(self.instance);
        }

        self.physical_device = try self.getPhysicalDevice();
        const device_handle = try self.createLogicalDevice();

        const vkd = try allocator.create(DeviceDispatch);
        errdefer allocator.destroy(vkd);
        vkd.* = try DeviceDispatch.load(device_handle, self.instance.wrapper.dispatch.vkGetDeviceProcAddr);

        self.device = Device.init(device_handle, vkd);
        const queues = try self.getDeviceQueues();

        self.graphics_queue = Queue.init(queues[0], vkd);
        self.presentation_queue = Queue.init(queues[1], vkd);

        self.swapchain = try self.createSwapChain();

        return self;
    }

    fn createInstance(self: Self) !vk.Instance {
        if (enable_validation_layers and !self.checkValidationLayersSupport()) {
            return error.LayerNotPresent;
        }

        const extensions = try self.getRequiredExtensions();
        defer self.allocator.free(extensions);

        std.debug.print("[Required instance extensions]\n", .{});
        for (extensions) |ext| {
            std.debug.print("\t- {s}\n", .{ext});
        }

        if (!try self.checkInstanceExtensions(&extensions)) {
            return error.ExtensionNotPresent;
        }

        const app_info = vk.ApplicationInfo{
            .p_application_name = "Vulkan SDL Test",
            .application_version = vk.makeApiVersion(0, 0, 1, 0),
            .p_engine_name = "Vulkan SDL Test",
            .engine_version = vk.makeApiVersion(0, 0, 1, 0),
            .api_version = vk.API_VERSION_1_3,
        };

        var instance_create_info: vk.InstanceCreateInfo = .{
            .p_application_info = &app_info,
            .enabled_extension_count = @intCast(extensions.len),
            .pp_enabled_extension_names = @ptrCast(extensions),
        };

        if (enable_validation_layers) {
            const debug_create_info = getDebugUtilsCreateInfo();

            instance_create_info.enabled_layer_count = @intCast(validation_layers.len);
            instance_create_info.pp_enabled_layer_names = &validation_layers;
            instance_create_info.p_next = &debug_create_info;
        }

        return try self.vkb.createInstance(&instance_create_info, null);
    }

    fn createLogicalDevice(self: Self) !vk.Device {
        const indices = try self.getQueueFamilies(self.physical_device);
        // 1 is the highest priority
        const priority = [_]f32{1};

        const qci = [_]vk.DeviceQueueCreateInfo{
            .{
                .queue_family_index = indices.graphics_family.?,
                .queue_count = 1,
                .p_queue_priorities = &priority,
            },
            .{
                .queue_family_index = indices.presentation_family.?,
                .queue_count = 1,
                .p_queue_priorities = &priority,
            },
        };

        const queue_count: u32 = if (indices.graphics_family.? == indices.presentation_family.?)
            1
        else
            2;

        const device_create_info: vk.DeviceCreateInfo = .{
            .queue_create_info_count = queue_count,
            .p_queue_create_infos = &qci,
            .pp_enabled_extension_names = &Utilities.device_extensions,
            .enabled_extension_count = @intCast(Utilities.device_extensions.len),
        };

        return try self.instance.createDevice(self.physical_device, &device_create_info, null);
    }

    fn createSwapChain(self: *Self) !vk.SwapchainKHR {
        const swapchain_details = try self.getSwapchainDetails(self.physical_device);
        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.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 = self.surface,
            .image_sharing_mode = .exclusive,
        };

        // Get queue family indices
        const family_indices = try self.getQueueFamilies(self.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 = 2; // Number of queues to share images between
            swapchain_create_info.p_queue_family_indices = &qfi;
        }

        return try self.device.createSwapchainKHR(&swapchain_create_info, null);
    }

    fn getPhysicalDevice(self: Self) !vk.PhysicalDevice {
        var pdev_count: u32 = 0;
        _ = try self.instance.enumeratePhysicalDevices(&pdev_count, null);

        const pdevs = try self.allocator.alloc(vk.PhysicalDevice, pdev_count);
        defer self.allocator.free(pdevs);

        _ = try self.instance.enumeratePhysicalDevices(&pdev_count, pdevs.ptr);

        for (pdevs) |pdev| {
            if (self.checkDeviceSuitable(pdev)) {
                return pdev;
            }
        }

        // TODO Obviously needs to be something else
        unreachable;
    }

    fn getRequiredExtensions(self: Self) ![][*:0]const u8 {
        var ext_count = sdl.vulkan.getInstanceExtensionsCount(self.window);

        if (enable_validation_layers) {
            ext_count += 1;
        }

        var extensions = try self.allocator.alloc([*:0]const u8, ext_count);
        _ = try sdl.vulkan.getInstanceExtensions(self.window, extensions);

        if (enable_validation_layers) {
            extensions[extensions.len - 1] = vk.extensions.ext_debug_utils.name;
        }

        return extensions;
    }

    fn getQueueFamilies(self: Self, pdev: vk.PhysicalDevice) !QueueFamilyIndices {
        var indices: QueueFamilyIndices = .{ .graphics_family = null };

        var queue_family_count: u32 = 0;
        self.instance.getPhysicalDeviceQueueFamilyProperties(pdev, &queue_family_count, null);

        const queue_family_list = try self.allocator.alloc(vk.QueueFamilyProperties, queue_family_count);
        defer self.allocator.free(queue_family_list);

        self.instance.getPhysicalDeviceQueueFamilyProperties(pdev, &queue_family_count, queue_family_list.ptr);

        for (queue_family_list, 0..) |queue_family, i| {
            if (queue_family.queue_count > 0 and queue_family.queue_flags.graphics_bit) {
                indices.graphics_family = @intCast(i);
            }

            const presentation_support = try self.instance.getPhysicalDeviceSurfaceSupportKHR(pdev, @intCast(i), self.surface);
            if (queue_family.queue_count > 0 and presentation_support == vk.TRUE) {
                indices.presentation_family = @intCast(i);
            }

            if (indices.isValid()) {
                return indices;
            }
        }

        unreachable;
    }

    fn getDeviceQueues(self: Self) ![2]vk.Queue {
        const indices = try self.getQueueFamilies(self.physical_device);

        const graphics_queue = self.device.getDeviceQueue(indices.graphics_family.?, 0);
        const presentation_queue = self.device.getDeviceQueue(indices.presentation_family.?, 0);
        return .{ graphics_queue, presentation_queue };
    }

    fn checkInstanceExtensions(self: Self, required_extensions: *const [][*:0]const u8) !bool {
        var prop_count: u32 = 0;
        _ = try self.vkb.enumerateInstanceExtensionProperties(null, &prop_count, null);

        const props = try self.allocator.alloc(vk.ExtensionProperties, prop_count);
        defer self.allocator.free(props);

        _ = try self.vkb.enumerateInstanceExtensionProperties(null, &prop_count, props.ptr);

        for (required_extensions.*) |required_extension| {
            for (props) |prop| {
                if (std.mem.eql(u8, std.mem.sliceTo(&prop.extension_name, 0), std.mem.span(required_extension))) {
                    break;
                }
            } else {
                return false;
            }
        }

        return true;
    }

    fn checkDeviceExtensions(self: Self, pdev: vk.PhysicalDevice) !bool {
        var prop_count: u32 = 0;
        _ = try self.instance.enumerateDeviceExtensionProperties(pdev, null, &prop_count, null);

        if (prop_count == 0) {
            return false;
        }

        const props = try self.allocator.alloc(vk.ExtensionProperties, prop_count);
        defer self.allocator.free(props);

        _ = try self.instance.enumerateDeviceExtensionProperties(pdev, null, &prop_count, props.ptr);

        for (Utilities.device_extensions) |device_extension| {
            for (props) |prop| {
                if (std.mem.eql(u8, std.mem.sliceTo(&prop.extension_name, 0), std.mem.span(device_extension))) {
                    break;
                }
            } else {
                return false;
            }
        }

        return true;
    }

    fn checkDeviceSuitable(self: Self, pdev: vk.PhysicalDevice) bool {
        const pdev_properties = self.instance.getPhysicalDeviceProperties(pdev);

        if (pdev_properties.device_type == .cpu) {
            return false;
        }

        const queue_family_indices = self.getQueueFamilies(pdev) catch return false;

        const extension_support = self.checkDeviceExtensions(pdev) catch return false;

        const swapchain_details = self.getSwapchainDetails(pdev) catch return false;
        defer self.allocator.free(swapchain_details.formats);
        defer self.allocator.free(swapchain_details.presentation_modes);

        const swapchain_valid = swapchain_details.formats.len != 0 and swapchain_details.formats.len != 0;

        return queue_family_indices.isValid() and extension_support and swapchain_valid;
    }

    fn checkValidationLayersSupport(self: Self) bool {
        var layer_count: u32 = undefined;
        _ = self.vkb.enumerateInstanceLayerProperties(&layer_count, null) catch {
            return false;
        };

        const available_layers = self.allocator.alloc(vk.LayerProperties, layer_count) catch unreachable;
        defer self.allocator.free(available_layers);

        _ = self.vkb.enumerateInstanceLayerProperties(&layer_count, available_layers.ptr) catch {
            return false;
        };

        for (validation_layers) |validation_layer| {
            for (available_layers) |available_layer| {
                if (std.mem.eql(u8, std.mem.span(validation_layer), std.mem.sliceTo(&available_layer.layer_name, 0))) {
                    return true;
                }
            }
        }

        return false;
    }

    fn getSwapchainDetails(self: Self, pdev: vk.PhysicalDevice) !SwapchainDetails {
        // Capabilities
        const surface_capabilities = try self.instance.getPhysicalDeviceSurfaceCapabilitiesKHR(pdev, self.surface);

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

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

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

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

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

    pub fn deinit(self: *Self) void {
        if (enable_validation_layers) {
            self.instance.destroyDebugUtilsMessengerEXT(self.debug_utils.?, null);
        }
        self.device.destroySwapchainKHR(self.swapchain, null);
        self.device.destroyDevice(null);
        self.instance.destroySurfaceKHR(self.surface, null);
        self.instance.destroyInstance(null);

        self.allocator.destroy(self.device.wrapper);
        self.allocator.destroy(self.instance.wrapper);
    }
};

// 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_unorm,
            .color_space = vk.ColorSpaceKHR.srgb_nonlinear_khr,
        };
    }

    for (formats) |format| {
        if ((format.format == vk.Format.r8g8b8a8_unorm or format.format == vk.Format.b8g8r8a8_unorm) 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;
}

// Validation layers stuff
fn createDebugMessenger(instance: Instance) !vk.DebugUtilsMessengerEXT {
    const debug_create_info = getDebugUtilsCreateInfo();

    return try instance.createDebugUtilsMessengerEXT(&debug_create_info, null);
}

fn getDebugUtilsCreateInfo() vk.DebugUtilsMessengerCreateInfoEXT {
    return vk.DebugUtilsMessengerCreateInfoEXT{
        .message_severity = .{ .verbose_bit_ext = true, .warning_bit_ext = true, .error_bit_ext = true },
        .message_type = .{ .general_bit_ext = true, .validation_bit_ext = true, .performance_bit_ext = true },
        .pfn_user_callback = debugCallback,
    };
}

// message_severity: vk.DebugUtilsMessageSeverityFlagsEXT,
// message_types: vk.DebugUtilsMessageTypeFlagsEXT,
// p_callback_data: ?*const vk.DebugUtilsMessengerCallbackDataEXT,
// p_user_data: ?*anyopaque,
fn debugCallback(
    message_severity: vk.DebugUtilsMessageSeverityFlagsEXT,
    message_types: vk.DebugUtilsMessageTypeFlagsEXT,
    p_callback_data: ?*const vk.DebugUtilsMessengerCallbackDataEXT,
    _: ?*anyopaque,
) callconv(vk.vulkan_call_conv) vk.Bool32 {
    const severity = getMessageSeverityLabel(message_severity);
    const message_type = getMessageTypeLabel(message_types);

    std.debug.print("[{s}] ({s}): {s}\n", .{ severity, message_type, p_callback_data.?.p_message.? });
    return vk.TRUE;
}

inline fn getMessageSeverityLabel(message_severity: vk.DebugUtilsMessageSeverityFlagsEXT) []const u8 {
    if (message_severity.verbose_bit_ext) {
        return "VERBOSE";
    } else if (message_severity.info_bit_ext) {
        return "INFO";
    } else if (message_severity.warning_bit_ext) {
        return "WARNING";
    } else if (message_severity.error_bit_ext) {
        return "ERROR";
    } else {
        unreachable;
    }
}

inline fn getMessageTypeLabel(message_types: vk.DebugUtilsMessageTypeFlagsEXT) []const u8 {
    if (message_types.general_bit_ext) {
        return "general";
    } else if (message_types.validation_bit_ext) {
        return "validation";
    } else if (message_types.performance_bit_ext) {
        return "performance";
    } else if (message_types.device_address_binding_bit_ext) {
        return "device_address_binding";
    } else {
        return "unknown";
    }
}