const Self = @This();

const std = @import("std");
const c = @import("c.zig");
const vk = @import("vk");
const builtin = @import("builtin");

const USE_DEBUG_LAYERS = switch (builtin.mode) {
    .ReleaseSafe, .Debug => true,
    .ReleaseSmall, .ReleaseFast => false,
};

const MAX_DEVICES = 16;
const MAX_DEVICE_EXTENSIONS = 512;
const MAX_INSTANCE_EXTENSIONS = 64;
const MAX_LAYERS = 512;

vkb: BaseDispatch,
vki: InstanceDispatch,
vkd: DeviceDispatch,

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

queues: struct {
    graphics: vk.Queue,
    present: vk.Queue,
},

surface: vk.SurfaceKHR,

messenger: if (USE_DEBUG_LAYERS) vk.DebugUtilsMessengerEXT else void,

fn Enumeration(comptime T: type, comptime cap: u32) type {
    return struct {
        buf: [cap]T = undefined,
        len: u32 = 0,

        const FULL: @This() = .{ .len = cap };
        const EMPTY: @This() = .{ .len = 0 };

        pub fn slice(self: anytype) switch (@TypeOf(&self.buf)) {
            *[cap]T => []T,
            *const [cap]T => []const T,
            else => unreachable,
        } {
            return self.buf[0..self.len];
        }

        pub fn appendSlice(self: *@This(), source: []const T) !void {
            if (self.len + source.len > cap) return error.Overflow;
            @memcpy(self.buf[self.len..][0..source.len], source);
            self.len += @intCast(source.len);
        }

        pub fn append(self: *@This(), val: T) !void {
            if (self.len + 1 > cap) return error.Overflow;
            self.buf[self.len] = val;
            self.len += 1;
        }
    };
}

pub fn init(window: *c.GLFWwindow) !Self {
    var self: Self = undefined;

    self.vkb = try BaseDispatch.load(&c.glfwGetInstanceProcAddress);
    const vkb = self.vkb;

    var req_exts = Enumeration([*:0]const u8, MAX_INSTANCE_EXTENSIONS).EMPTY;
    var req_layers = Enumeration([*:0]const u8, MAX_LAYERS).EMPTY;
    var req_dev_exts = Enumeration([*:0]const u8, MAX_DEVICE_EXTENSIONS).EMPTY;

    try req_dev_exts.append("VK_KHR_swapchain");

    if (USE_DEBUG_LAYERS) {
        try req_layers.append("VK_LAYER_KHRONOS_validation");
        try req_exts.append("VK_EXT_debug_utils");
    }

    {
        var glfw_ext_count: u32 = 0;
        const glfw_exts: [*][*:0]const u8 = @ptrCast(c.glfwGetRequiredInstanceExtensions(&glfw_ext_count));
        try req_exts.appendSlice(glfw_exts[0..glfw_ext_count]);
    }

    std.log.debug("requesting extensions: {s}", .{req_exts.slice()});
    std.log.debug("requesting layers: {s}", .{req_layers.slice()});
    std.log.debug("requesting device extensions: {s}", .{req_dev_exts.slice()});

    var available_exts = Enumeration(vk.ExtensionProperties, MAX_INSTANCE_EXTENSIONS).FULL;
    _ = try vkb.enumerateInstanceExtensionProperties(
        null,
        &available_exts.len,
        &available_exts.buf,
    );

    var available_layers = Enumeration(vk.LayerProperties, MAX_LAYERS).FULL;
    _ = try vkb.enumerateInstanceLayerProperties(
        &available_layers.len,
        &available_layers.buf,
    );

    for (req_exts.slice()) |name| {
        const required_name = std.mem.sliceTo(name, 0);
        for (available_exts.slice()) |prop| {
            const available_name = std.mem.sliceTo(&prop.extension_name, 0);
            if (std.mem.eql(u8, required_name, available_name)) break;
        } else {
            return error.ExtensionNotPresent;
        }
    }

    for (req_layers.slice()) |name| {
        const required_name = std.mem.sliceTo(name, 0);
        for (available_layers.slice()) |prop| {
            const available_name = std.mem.sliceTo(&prop.layer_name, 0);
            if (std.mem.eql(u8, required_name, available_name)) break;
        } else {
            return error.LayerNotPresent;
        }
    }

    const debug_create_info = vk.DebugUtilsMessengerCreateInfoEXT{
        .message_severity = vk.DebugUtilsMessageSeverityFlagsEXT{
            .verbose_bit_ext = false,
            .warning_bit_ext = true,
            .error_bit_ext = true,
            .info_bit_ext = false,
        },
        .message_type = vk.DebugUtilsMessageTypeFlagsEXT{
            .general_bit_ext = true,
            .validation_bit_ext = true,
            .performance_bit_ext = true,
            .device_address_binding_bit_ext = false,
        },
        .pfn_user_callback = &debug_callback,
        .p_user_data = null,
    };

    const app_info = vk.ApplicationInfo{
        .p_application_name = "Hello World",
        .application_version = vk.makeApiVersion(0, 0, 0, 0),
        .p_engine_name = "No Engine",
        .engine_version = vk.makeApiVersion(0, 0, 0, 0),
        .api_version = vk.API_VERSION_1_3,
    };

    const instance_create_info = vk.InstanceCreateInfo{
        .p_application_info = &app_info,
        .enabled_extension_count = req_exts.len,
        .pp_enabled_extension_names = &req_exts.buf,
        .enabled_layer_count = req_layers.len,
        .pp_enabled_layer_names = &req_layers.buf,
        .p_next = if (USE_DEBUG_LAYERS) &debug_create_info else null,
    };

    self.instance = try vkb.createInstance(&instance_create_info, null);
    self.vki = try InstanceDispatch.load(self.instance, vkb.dispatch.vkGetInstanceProcAddr);
    const vki = self.vki;
    errdefer vki.destroyInstance(self.instance, null);

    if (USE_DEBUG_LAYERS) self.messenger = try vki.createDebugUtilsMessengerEXT(
        self.instance,
        &debug_create_info,
        null,
    );
    errdefer if (USE_DEBUG_LAYERS) vki.destroyDebugUtilsMessengerEXT(
        self.instance,
        self.messenger,
        null,
    );

    switch (c.glfwCreateWindowSurface(
        self.instance,
        window,
        null,
        &self.surface,
    )) {
        .success => {},
        else => |e| {
            std.log.err("{}", .{e});
            return error.Unknown;
        },
    }
    errdefer vki.destroySurfaceKHR(self.instance, self.surface, null);

    var devices = Enumeration(vk.PhysicalDevice, MAX_DEVICES).FULL;
    _ = try vki.enumeratePhysicalDevices(
        self.instance,
        &devices.len,
        &devices.buf,
    );

    // todo some ranking strategy to find the most-suitable device
    const Selection = struct {
        device: vk.PhysicalDevice,
        props: vk.PhysicalDeviceProperties,
        feats: vk.PhysicalDeviceFeatures,
    };
    const selected: Selection = find_device: for (devices.slice()) |device| {
        const props = vki.getPhysicalDeviceProperties(device);
        const feats = vki.getPhysicalDeviceFeatures(device);

        if (props.device_type != vk.PhysicalDeviceType.discrete_gpu) continue;

        // if (feats.geometry_shader == vk.FALSE) continue;

        var available_dev_exts = Enumeration(vk.ExtensionProperties, MAX_DEVICE_EXTENSIONS).FULL;
        _ = try vki.enumerateDeviceExtensionProperties(
            device,
            null,
            &available_dev_exts.len,
            &available_dev_exts.buf,
        );

        for (req_dev_exts.slice()) |name| {
            const required_name = std.mem.sliceTo(name, 0);
            for (available_dev_exts.slice()) |prop| {
                const available_name = std.mem.sliceTo(&prop.extension_name, 0);
                if (std.mem.eql(u8, required_name, available_name)) break;
            } else {
                std.log.warn("cannot find {s}\n", .{required_name});
                continue :find_device;
            }
        }

        break .{
            .device = device,
            .props = props,
            .feats = feats,
        };
    } else {
        return error.NoSuitablePhysicalDevice;
    };

    var queue_families = Enumeration(vk.QueueFamilyProperties, 64).FULL;
    vki.getPhysicalDeviceQueueFamilyProperties(
        selected.device,
        &queue_families.len,
        &queue_families.buf,
    );

    // todo this should be incorporated with physical device selection/ranking.
    const Indices = struct {
        graphics: u32,
        present: u32,
    };
    const indices: Indices = find_index: {
        var graphics: ?u32 = null;
        var present: ?u32 = null;

        for (queue_families.slice(), 0..) |prop, idx| {
            if (graphics == null and prop.queue_flags.graphics_bit) {
                graphics = @intCast(idx);
                // continue;  // forces distinct queue families
            }

            if (present == null) {
                const present_support = try vki.getPhysicalDeviceSurfaceSupportKHR(
                    selected.device,
                    @intCast(idx),
                    self.surface,
                ) == vk.TRUE;
                if (present_support) {
                    present = @intCast(idx);
                }
            }

            if (graphics != null and present != null) {
                break :find_index .{
                    .graphics = graphics.?,
                    .present = present.?,
                };
            }
        }

        return error.IncompatibleDeviceQueues;
    };

    const gp_priorities = [_]f32{ 1.0, 1.0 };

    var queue_create_infos = Enumeration(vk.DeviceQueueCreateInfo, 2).EMPTY;

    // queue info family indices must be unique. so if the graphics and present queues are the same, create two queues
    // in the same family. otherwise create queues in separate families. there should probably be some general way to
    // group and unpack the queues, but I'm not bothering with that for now until I restructure this monolithic function
    // in general.
    if (indices.graphics == indices.present) {
        const gp_slice = gp_priorities[0..2];
        try queue_create_infos.append(.{
            .queue_family_index = indices.graphics,
            .queue_count = @intCast(gp_slice.len),
            .p_queue_priorities = gp_slice.ptr,
        });
    } else {
        const g_slice = gp_priorities[0..1];
        const p_slice = gp_priorities[1..2];
        try queue_create_infos.append(.{
            .queue_family_index = indices.graphics,
            .queue_count = @intCast(g_slice.len),
            .p_queue_priorities = g_slice.ptr,
        });
        try queue_create_infos.append(.{
            .queue_family_index = indices.present,
            .queue_count = @intCast(p_slice.len),
            .p_queue_priorities = p_slice.ptr,
        });
    }

    const device_create_info = vk.DeviceCreateInfo{
        .queue_create_info_count = queue_create_infos.len,
        .p_queue_create_infos = &queue_create_infos.buf,
        .p_enabled_features = &selected.feats,
        .enabled_extension_count = req_dev_exts.len,
        .pp_enabled_extension_names = &req_dev_exts.buf,
        .enabled_layer_count = req_layers.len,
        .pp_enabled_layer_names = &req_layers.buf,
    };

    self.device = try vki.createDevice(
        selected.device,
        &device_create_info,
        null,
    );
    self.vkd = try DeviceDispatch.load(self.device, vki.dispatch.vkGetDeviceProcAddr);
    const vkd = self.vkd;
    errdefer vkd.destroyDevice(self.device, null);

    if (indices.graphics == indices.present) {
        // two queues in the same family
        self.queues = .{
            .graphics = vkd.getDeviceQueue(self.device, indices.graphics, 0),
            .present = vkd.getDeviceQueue(self.device, indices.present, 1),
        };
    } else {
        // queues from different families
        self.queues = .{
            .graphics = vkd.getDeviceQueue(self.device, indices.graphics, 0),
            .present = vkd.getDeviceQueue(self.device, indices.present, 0),
        };
    }

    return self;
}

pub fn deinit(self: Self) void {
    self.vki.destroySurfaceKHR(self.instance, self.surface, null);
    self.vkd.destroyDevice(self.device, null);
    if (USE_DEBUG_LAYERS) self.vki.destroyDebugUtilsMessengerEXT(
        self.instance,
        self.messenger,
        null,
    );
    self.vki.destroyInstance(self.instance, null);
}

export fn debug_callback(
    message_severity: vk.DebugUtilsMessageSeverityFlagsEXT,
    message_type: vk.DebugUtilsMessageTypeFlagsEXT,
    p_callback_data: ?*const vk.DebugUtilsMessengerCallbackDataEXT,
    _: ?*anyopaque,
) callconv(.C) vk.Bool32 {
    if (p_callback_data == null) return vk.FALSE;
    if (p_callback_data.?.p_message == null) return vk.FALSE;
    const msg = p_callback_data.?.p_message.?;

    const scopes = .{
        "validation",
        "performance",
        "device_address_binding",
        "general",
    };

    const scope: []const u8 = inline for (scopes) |tag| {
        if (@field(message_type, tag ++ "_bit_ext")) {
            break tag;
        }
    } else {
        return vk.FALSE;
    };

    const levels = .{
        "error",
        "info",
        "warning",
        "verbose",
    };

    const level: []const u8 = inline for (levels) |tag| {
        if (@field(message_severity, tag ++ "_bit_ext")) {
            break tag;
        }
    } else {
        return vk.FALSE;
    };

    // ripped from std.log, but with my own levels and scope.
    const stderr = std.io.getStdErr().writer();
    var bw = std.io.bufferedWriter(stderr);
    const writer = bw.writer();

    std.debug.getStderrMutex().lock();
    defer std.debug.getStderrMutex().unlock();
    nosuspend {
        writer.print("vk-{s}({s}): {s}\n", .{ level, scope, msg }) catch return vk.FALSE;
        bw.flush() catch return vk.FALSE;
    }

    return vk.FALSE;
}

const BaseDispatch = vk.BaseWrapper(.{
    .createInstance = true,
    .getInstanceProcAddr = true,
    .enumerateInstanceExtensionProperties = true,
    .enumerateInstanceLayerProperties = true,
});

const InstanceDispatch = vk.InstanceWrapper(.{
    .destroyInstance = true,
    .createDebugUtilsMessengerEXT = USE_DEBUG_LAYERS,
    .destroyDebugUtilsMessengerEXT = USE_DEBUG_LAYERS,
    .submitDebugUtilsMessageEXT = USE_DEBUG_LAYERS,
    .enumeratePhysicalDevices = true,
    .getPhysicalDeviceProperties = true,
    .getPhysicalDeviceFeatures = true,
    .getPhysicalDeviceQueueFamilyProperties = true,
    .createDevice = true,
    .getDeviceProcAddr = true,
    .destroySurfaceKHR = true,
    .getPhysicalDeviceSurfaceSupportKHR = true,
    .enumerateDeviceExtensionProperties = true,
});

const DeviceDispatch = vk.DeviceWrapper(.{
    .destroyDevice = true,
    .getDeviceQueue = true,
});