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

const vk = @import("vulkan");
const sdl = @import("sdl");
const img = @import("zstbi");

const validation = @import("./validation_layers.zig");
const Swapchain = @import("Swapchain.zig");
const QueueUtils = @import("queue_utils.zig");

const device_extensions = [_][*:0]const u8{vk.extensions.khr_swapchain.name};

pub 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 enable_validation_layers = builtin.mode == .Debug;
const validation_layers = [_][*:0]const u8{"VK_LAYER_KHRONOS_validation"};

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

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

// ---

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,

debug_utils: ?vk.DebugUtilsMessengerEXT,

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

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

    img.init(allocator);

    try self.createInstance();

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

    try self.createSurface();

    try self.getPhysicalDevice();
    try self.createLogicalDevice();

    return self;
}

pub fn deinit(self: *Self) void {
    if (enable_validation_layers) {
        self.instance.destroyDebugUtilsMessengerEXT(self.debug_utils.?, 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);

    img.deinit();
}

fn createInstance(self: *Self) !void {
    if (enable_validation_layers and !self.checkValidationLayersSupport()) {
        // TODO Better error
        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 = validation.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;
    }

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

    self.instance = Instance.init(instance_handle, vki);
}

fn createSurface(self: *Self) !void {
    self.surface = try sdl.vulkan.createSurface(self.window, self.instance.handle);
}

fn getPhysicalDevice(self: *Self) !void {
    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)) {
            self.physical_device = pdev;
            break;
        }
    } else {
        // TODO Obviously needs to be something else
        unreachable;
    }
}

fn createLogicalDevice(self: *Self) !void {
    const indices = try QueueUtils.getQueueFamilies(self.*, 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;

    // Device features
    const device_features: vk.PhysicalDeviceFeatures = .{
        .sampler_anisotropy = vk.TRUE, // Enable anisotropy
    };

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

    const device_handle = try self.instance.createDevice(self.physical_device, &device_create_info, null);

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

    self.device = Device.init(device_handle, vkd);

    const queues = try QueueUtils.getDeviceQueues(self.*);

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

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 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 (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 pdev_features = self.instance.getPhysicalDeviceFeatures(pdev);
    const queue_family_indices = QueueUtils.getQueueFamilies(self, pdev) catch return false;
    const extension_support = self.checkDeviceExtensions(pdev) catch return false;

    const swapchain_details = Swapchain.getSwapchainDetails(
        self.allocator,
        self.instance,
        pdev,
        self.surface,
    ) 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 and pdev_features.sampler_anisotropy == vk.TRUE;
}

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;
}
