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

const c = @import("../const.zig");
const glfw = @import("zglfw");
const vk = @import("vulkan");

const Queue = @import("Queue.zig");
const VkAllocator = @import("VkAllocator.zig");

pub const Mode = union(enum) {
    headless: void,
    surface: struct {
        window: *glfw.Window,
        surface: vk.SurfaceKHR,
        presentation_queue: Queue,
    },

    pub fn deinit(self: *Mode, instance: vk.InstanceProxy, vk_allocator: *VkAllocator) void {
        switch (self.*) {
            .headless => {},
            .surface => |x| {
                instance.destroySurfaceKHR(x.surface, &vk_allocator.interface);
            },
        }

        self.* = undefined;
    }
};

pub const ModeTag = std.meta.Tag(Mode);

const QueueAllocator = struct {
    queue_families_properties_buffer: [max_queue_families]vk.QueueFamilyProperties,
    queue_families_in_use: [max_queue_families]u32,

    queue_families_count: u32,

    const max_queue_families = 16;

    pub fn init(instance: vk.InstanceProxy, physical_device: vk.PhysicalDevice) QueueAllocator {
        var queue_families_properties_buffer: [max_queue_families]vk.QueueFamilyProperties = undefined;
        var queue_families_count: u32 = max_queue_families;

        instance.getPhysicalDeviceQueueFamilyProperties(physical_device, &queue_families_count, &queue_families_properties_buffer);

        var queue_families_in_use: [max_queue_families]u32 = undefined;
        @memset(queue_families_in_use[0..queue_families_count], 0);

        return .{
            .queue_families_properties_buffer = queue_families_properties_buffer,
            .queue_families_in_use = queue_families_in_use,
            .queue_families_count = queue_families_count,
        };
    }

    pub fn allocateQueueIndex(self: *QueueAllocator, queue_family: u32) u32 {
        std.debug.assert(queue_family < self.queue_families_count);
        if (self.queue_families_in_use[queue_family] < self.queue_families_properties_buffer[queue_family].queue_count) {
            const queue_index = self.queue_families_in_use[queue_family];
            self.queue_families_in_use[queue_family] += 1;
            return queue_index;
        } else {
            return 0;
        }
    }

    pub fn queueFamiliesProperties(self: *const QueueAllocator) []const vk.QueueFamilyProperties {
        return self.queue_families_properties_buffer[0..self.queue_families_count];
    }

    pub fn queueInUse(self: *const QueueAllocator) []const u32 {
        return self.queue_families_in_use[0..self.queue_families_count];
    }
};

const QueueAllocations = struct {
    graphics_queue: Queue.Allocation,
    compute_queue: Queue.Allocation,
    transfer_queue: Queue.Allocation,
    /// Defined only in `surface` mode, `undefined` otherwise.
    presentation_queue: Queue.Allocation,
};

const DebugUtilsMessenger = if (debug) vk.DebugUtilsMessengerEXT else void;

mode: Mode,
vk_allocator: *VkAllocator,

base: vk.BaseWrapper,
instance: vk.InstanceProxy,
device: vk.DeviceProxy,

debug_utils_messenger: DebugUtilsMessenger,
physical_device: vk.PhysicalDevice,
memory_types: std.ArrayList(vk.MemoryType),
memory_heaps: std.ArrayList(vk.MemoryHeap),

graphics_queue: Queue,
compute_queue: Queue,
transfer_queue: Queue,

const debug = @import("builtin").mode == .Debug;

const vk_application_info: vk.ApplicationInfo = .{
    .p_application_name = c.app_name,
    .application_version = @bitCast(vk_version),
    .p_engine_name = c.app_name,
    .engine_version = @bitCast(vk_version),
    .api_version = @bitCast(vk.API_VERSION_1_2),
};

const vk_debug_utils_messenger_create_info: vk.DebugUtilsMessengerCreateInfoEXT = .{
    .message_severity = .{
        .verbose_bit_ext = false,
        .info_bit_ext = false,
        .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 = &debugUtilsMessengerCallback,
    .p_user_data = null,
};

const vk_version: vk.Version = vk.makeApiVersion(
    0,
    c.app_version.major,
    c.app_version.minor,
    c.app_version.patch,
);

pub fn init(allocator: std.mem.Allocator, maybe_window: ?*glfw.Window) !Engine {
    const vk_allocator = try VkAllocator.init(allocator);
    errdefer vk_allocator.deinit();

    // --- LOAD BASE DISPATCH --------------------------------------------------

    const base = vk.BaseWrapper.load(glfw.getInstanceProcAddress);

    // --- CREATE INSTANCE, ITS WRAPPER AND PROXY ------------------------------

    const instance_handle = blk: {
        var enabled_extensions_buffer: [16][*:0]const u8 = undefined;
        var enabled_extensions: std.ArrayList([*:0]const u8) = .initBuffer(&enabled_extensions_buffer);

        if (debug) {
            try enabled_extensions.appendBounded(vk.extensions.ext_debug_utils.name);
        }

        if (maybe_window != null) {
            const glfw_instance_extensions = try glfw.getRequiredInstanceExtensions();
            try enabled_extensions.appendSliceBounded(glfw_instance_extensions);
        }

        const enabled_layers: []const [*:0]const u8 = if (debug) &.{"VK_LAYER_KHRONOS_validation"} else &.{};

        break :blk try base.createInstance(&.{
            .p_application_info = &vk_application_info,
            .enabled_layer_count = enabled_layers.len,
            .pp_enabled_layer_names = enabled_layers.ptr,
            .enabled_extension_count = @intCast(enabled_extensions.items.len),
            .pp_enabled_extension_names = enabled_extensions.items.ptr,
            .p_next = if (debug) &vk_debug_utils_messenger_create_info else null,
        }, &vk_allocator.interface);
    };

    const instance_wrapper_ptr = try allocator.create(vk.InstanceWrapper);
    errdefer allocator.destroy(instance_wrapper_ptr);

    instance_wrapper_ptr.* = .load(instance_handle, base.dispatch.vkGetInstanceProcAddr.?);
    const instance = vk.InstanceProxy.init(instance_handle, instance_wrapper_ptr);
    errdefer instance.destroyInstance(&vk_allocator.interface);

    // --- CREATE DEBUG UTILS MESSENGER ----------------------------------------

    const debug_utils_messenger: DebugUtilsMessenger = if (debug) try instance.createDebugUtilsMessengerEXT(&vk_debug_utils_messenger_create_info, &vk_allocator.interface) else {};
    errdefer {
        if (debug) instance.destroyDebugUtilsMessengerEXT(debug_utils_messenger, &vk_allocator.interface);
    }

    // --- CHOOSE PHYSICAL DEVICE, GET ITS MEMORY PROPERTIES -------------------

    const physical_device = blk: {
        var physical_devices_buffer: [16]vk.PhysicalDevice = undefined;
        var physical_devices_count: u32 = physical_devices_buffer.len;
        _ = try instance.enumeratePhysicalDevices(&physical_devices_count, &physical_devices_buffer);

        // Look for the first GPU with its `device_type` equal to
        // `.discrete_gpu`.

        for (physical_devices_buffer[0..physical_devices_count]) |physical_device_candidate| {
            const physical_device_properties = instance.getPhysicalDeviceProperties(physical_device_candidate);
            if (physical_device_properties.device_type == .discrete_gpu) {
                break :blk physical_device_candidate;
            }
        }

        // Look for the first GPU.

        if (physical_devices_count > 0) {
            break :blk physical_devices_buffer[0];
        }

        return error.NoDevice;
    };

    const physical_device_memory_properties = instance.getPhysicalDeviceMemoryProperties(physical_device);

    var memory_types = try std.ArrayList(vk.MemoryType).initCapacity(allocator, vk.MAX_MEMORY_TYPES);
    errdefer memory_types.deinit(allocator);
    memory_types.appendSliceAssumeCapacity(physical_device_memory_properties.memory_types[0..physical_device_memory_properties.memory_type_count]);

    var memory_heaps = try std.ArrayList(vk.MemoryHeap).initCapacity(allocator, vk.MAX_MEMORY_HEAPS);
    errdefer memory_heaps.deinit(allocator);
    memory_heaps.appendSliceAssumeCapacity(physical_device_memory_properties.memory_heaps[0..physical_device_memory_properties.memory_heap_count]);

    // --- CREATE SURFACE ------------------------------------------------------

    var surface: vk.SurfaceKHR = undefined;
    if (maybe_window) |window| try glfw.createWindowSurface(instance_handle, window, &vk_allocator.interface, &surface);
    errdefer {
        if (maybe_window != null) instance.destroySurfaceKHR(surface, &vk_allocator.interface);
    }

    // --- ALLOCATE QUEUES -----------------------------------------------------

    const max_queues = 4;

    var queue_create_info_buffer: [max_queues]vk.DeviceQueueCreateInfo = undefined;
    var queue_create_info: std.ArrayList(vk.DeviceQueueCreateInfo) = .initBuffer(&queue_create_info_buffer);

    const queue_allocations: QueueAllocations = blk: {
        var queue_allocator: QueueAllocator = .init(instance, physical_device);
        const queue_families_properties = queue_allocator.queueFamiliesProperties();

        const graphics_queue_family = findGraphicsQueueFamily(queue_families_properties) orelse return error.NoGraphicsQueue;
        const compute_queue_family = findComputeQueueFamily(queue_families_properties) orelse return error.NoComputeQueue;
        const transfer_queue_family = findTransferQueueFamily(queue_families_properties) orelse return error.NoTransferQueue;
        const presentation_queue_family = if (maybe_window != null) (try findPresentationQueueFamily(queue_families_properties, instance, physical_device, surface) orelse return error.NoPresentationQueue) else undefined;

        const graphics_queue_index = queue_allocator.allocateQueueIndex(graphics_queue_family);
        const compute_queue_index = queue_allocator.allocateQueueIndex(compute_queue_family);
        const transfer_queue_index = queue_allocator.allocateQueueIndex(transfer_queue_family);
        const presentation_queue_index = if (maybe_window != null) queue_allocator.allocateQueueIndex(presentation_queue_family) else undefined;

        const queue_priorities = [_]f32{1.0} ** max_queues;

        for (queue_allocator.queueInUse(), 0..) |count, queue_family| {
            if (count == 0) continue;

            queue_create_info.appendBounded(.{
                .queue_family_index = @intCast(queue_family),
                .queue_count = count,
                .p_queue_priorities = &queue_priorities,
            }) catch unreachable;
        }

        break :blk .{
            .graphics_queue = .{ .family = graphics_queue_family, .index = graphics_queue_index },
            .compute_queue = .{ .family = compute_queue_family, .index = compute_queue_index },
            .transfer_queue = .{ .family = transfer_queue_family, .index = transfer_queue_index },
            .presentation_queue = if (maybe_window != null) .{ .family = presentation_queue_family, .index = presentation_queue_index } else undefined,
        };
    };

    // --- CREATE (LOGICAL) DEVICE, ITS WRAPPER AND PROXY ----------------------

    const device_handle = blk: {
        var enabled_extensions_buffer: [16][*:0]const u8 = undefined;
        var enabled_extensions: std.ArrayList([*:0]const u8) = .initBuffer(&enabled_extensions_buffer);

        if (maybe_window != null) {
            try enabled_extensions.appendBounded(vk.extensions.khr_swapchain.name);
        }

        break :blk try instance.createDevice(physical_device, &.{
            .queue_create_info_count = @intCast(queue_create_info.items.len),
            .p_queue_create_infos = queue_create_info.items.ptr,
            .enabled_extension_count = @intCast(enabled_extensions.items.len),
            .pp_enabled_extension_names = enabled_extensions.items.ptr,
        }, &vk_allocator.interface);
    };

    const device_wrapper_ptr = try allocator.create(vk.DeviceWrapper);
    errdefer allocator.destroy(device_wrapper_ptr);

    device_wrapper_ptr.* = .load(device_handle, instance.wrapper.dispatch.vkGetDeviceProcAddr.?);
    const device = vk.DeviceProxy.init(device_handle, device_wrapper_ptr);
    errdefer device.destroyDevice(vk_allocator.interface);

    // -------------------------------------------------------------------------

    return .{
        .mode = if (maybe_window) |window| .{ .surface = .{
            .window = window,
            .surface = surface,
            .presentation_queue = .initAllocation(device, queue_allocations.presentation_queue),
        } } else .{ .headless = {} },
        .vk_allocator = vk_allocator,

        .base = base,
        .instance = instance,
        .device = device,

        .debug_utils_messenger = debug_utils_messenger,
        .physical_device = physical_device,
        .memory_types = memory_types,
        .memory_heaps = memory_heaps,

        .graphics_queue = .initAllocation(device, queue_allocations.graphics_queue),
        .compute_queue = .initAllocation(device, queue_allocations.compute_queue),
        .transfer_queue = .initAllocation(device, queue_allocations.transfer_queue),
    };
}

pub fn deinit(self: *Engine) void {
    const allocator = self.vk_allocator.allocator;

    self.device.destroyDevice(&self.vk_allocator.interface);
    allocator.destroy(self.device.wrapper);

    if (std.meta.activeTag(self.mode) == .surface) {
        self.instance.destroySurfaceKHR(self.mode.surface.surface, &self.vk_allocator.interface);
    }

    self.memory_heaps.deinit(allocator);
    self.memory_types.deinit(allocator);

    if (debug) self.instance.destroyDebugUtilsMessengerEXT(self.debug_utils_messenger, &self.vk_allocator.interface);

    self.instance.destroyInstance(&self.vk_allocator.interface);
    allocator.destroy(self.instance.wrapper);

    self.vk_allocator.deinit();

    self.* = undefined;
}

pub fn allocate(self: *const Engine, memory_requirements: vk.MemoryRequirements, memory_property_flags: vk.MemoryPropertyFlags) !vk.DeviceMemory {
    for (self.memory_types.items, 0..) |memory_type, i| {
        const is_type = memory_requirements.memory_type_bits & (@as(u32, 1) << @truncate(i)) != 0;
        const has_flags = memory_type.property_flags.contains(memory_property_flags);
        if (is_type and has_flags) {
            return try self.device.allocateMemory(&.{
                .allocation_size = memory_requirements.size,
                .memory_type_index = @truncate(i),
            });
        }
    }

    return error.NoSuitableMemoryType;
}

fn debugUtilsMessengerCallback(
    severity: vk.DebugUtilsMessageSeverityFlagsEXT,
    message_type: vk.DebugUtilsMessageTypeFlagsEXT,
    maybe_callback_data: ?*const vk.DebugUtilsMessengerCallbackDataEXT,
    _: ?*anyopaque,
) callconv(vk.vulkan_call_conv) vk.Bool32 {
    const log = std.log.scoped(.vulkan);

    const message = if (maybe_callback_data) |callback_data| callback_data.p_message orelse return .false else return .false;
    const type_name = if (message_type.general_bit_ext) "general" else if (message_type.validation_bit_ext) "validation" else if (message_type.performance_bit_ext) "performance" else if (message_type.device_address_binding_bit_ext) "device_address_binding" else "unknown";

    const format = "[{s}] {s}";
    const args = .{ type_name, message };

    if (severity.error_bit_ext) {
        log.err(format, args);
    } else if (severity.warning_bit_ext) {
        log.warn(format, args);
    } else if (severity.info_bit_ext) {
        log.info(format, args);
    } else if (severity.verbose_bit_ext) {
        log.debug(format, args);
    }

    return .false;
}

fn findGraphicsQueueFamily(queue_families_properties: []const vk.QueueFamilyProperties) ?u32 {
    // Look for the first family that has the `graphics_bit` set.

    for (queue_families_properties, 0..) |queue_family_properties, i| {
        const graphics_bit = queue_family_properties.queue_flags.graphics_bit;
        if (graphics_bit) {
            return @intCast(i);
        }
    }

    return null;
}

fn findComputeQueueFamily(queue_families_properties: []const vk.QueueFamilyProperties) ?u32 {
    // Look for the first family that has the `compute_bit` set, but not
    // `graphics_bit`.

    for (queue_families_properties, 0..) |queue_family_properties, i| {
        const compute_bit = queue_family_properties.queue_flags.compute_bit;
        const graphics_bit = queue_family_properties.queue_flags.graphics_bit;
        if (compute_bit and !graphics_bit) {
            return @intCast(i);
        }
    }

    // Look for the first family that has the `compute_bit` set.

    for (queue_families_properties, 0..) |queue_family_properties, i| {
        const compute_bit = queue_family_properties.queue_flags.compute_bit;
        if (compute_bit) {
            return @intCast(i);
        }
    }

    return null;
}

fn findTransferQueueFamily(queue_families_properties: []const vk.QueueFamilyProperties) ?u32 {
    // Look for the first family that has the `transfer_bit` set, but neither
    // `graphics_bit` or `compute_bit`.

    for (queue_families_properties, 0..) |queue_family_properties, i| {
        const transfer_bit = queue_family_properties.queue_flags.transfer_bit;
        const graphics_bit = queue_family_properties.queue_flags.graphics_bit;
        const compute_bit = queue_family_properties.queue_flags.compute_bit;
        if (transfer_bit and !graphics_bit and !compute_bit) {
            return @intCast(i);
        }
    }

    // Look for the first family that has either `transfer_bit`, `graphics_bit`
    // or `compute_bit` set.

    // NOTE `graphics_bit` or `compute_bit` imply `transfer_bit`, even if not
    // set explicitly.

    for (queue_families_properties, 0..) |queue_family_properties, i| {
        const transfer_bit = queue_family_properties.queue_flags.transfer_bit;
        const graphics_bit = queue_family_properties.queue_flags.graphics_bit;
        const compute_bit = queue_family_properties.queue_flags.compute_bit;
        if (transfer_bit or graphics_bit or compute_bit) {
            return @intCast(i);
        }
    }

    return null;
}

fn findPresentationQueueFamily(queue_families_properties: []const vk.QueueFamilyProperties, instance: vk.InstanceProxy, physical_device: vk.PhysicalDevice, surface: vk.SurfaceKHR) !?u32 {
    // Look for the first family that supports presentation to a given surface.

    for (queue_families_properties, 0..) |_, i| {
        if (try instance.getPhysicalDeviceSurfaceSupportKHR(physical_device, @intCast(i), surface) == .true) {
            return @intCast(i);
        }
    }

    return null;
}