use std::{
    borrow::Cow,
    collections::{BTreeMap, BTreeSet, HashMap},
    ffi::{CStr, CString},
    ops::Deref,
    sync::Arc,
};

use ash::{
    khr,
    prelude::VkResult,
    vk::{self, Handle},
};
use raw_window_handle::RawDisplayHandle;
use tinyvec::{array_vec, ArrayVec};

use crate::{
    make_extention_properties, sync, Error, ExtendsDeviceProperties2Debug, Instance,
    PhysicalDevice, PhysicalDeviceFeatures, PhysicalDeviceProperties, Queue, Result, VkNameList,
};

#[derive(Debug, Default)]
pub struct DeviceQueueFamilies {
    pub(crate) families: Vec<(u32, u32)>,
    pub(crate) graphics: (u32, u32),
    pub(crate) present: (u32, u32),
    pub(crate) async_compute: (u32, u32),
    pub(crate) transfer: (u32, u32),
}

impl DeviceQueueFamilies {
    pub fn swapchain_family_indices(&self) -> ArrayVec<[u32; 2]> {
        let mut indices = array_vec!([u32; 2] => self.graphics.0);

        if self.present.0 != self.graphics.0 {
            indices.push(self.present.0);
        }

        indices
    }
    pub fn graphics_familty(&self) -> u32 {
        self.graphics.0
    }
    pub fn present_familty(&self) -> u32 {
        self.present.0
    }
    pub fn async_compute_familty(&self) -> u32 {
        self.async_compute.0
    }
    pub fn transfer_familty(&self) -> u32 {
        self.transfer.0
    }

    pub fn family_indices(&self, flags: QueueFlags) -> ArrayVec<[u32; 4]> {
        let mut indices = array_vec!([u32; 4]);
        if flags.contains(QueueFlags::GRAPHICS) {
            indices.push(self.graphics_familty());
        }
        if flags.contains(QueueFlags::PRESENT) {
            indices.push(self.present_familty());
        }
        if flags.contains(QueueFlags::ASYNC_COMPUTE) {
            indices.push(self.async_compute_familty());
        }
        if flags.contains(QueueFlags::TRANSFER) {
            indices.push(self.transfer_familty());
        }

        let unique_len = indices.partition_dedup().0.len();
        indices.drain(unique_len..);

        indices
    }
}

bitflags::bitflags! {
    #[repr(transparent)]
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct QueueFlags: u32 {
        const GRAPHICS = 1 << 0;
        const PRESENT = 1 << 1;
        const ASYNC_COMPUTE = 1 << 2;
        const TRANSFER = 1 << 3;

        const NONE = 0;
        const PRESENT_GRAPHICS = 1 << 0 | 1 << 1;
    }
}

#[repr(transparent)]
struct DeviceWrapper(ash::Device);

impl Deref for DeviceWrapper {
    type Target = ash::Device;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Drop for DeviceWrapper {
    fn drop(&mut self) {
        unsafe {
            _ = self.0.device_wait_idle();
            self.0.destroy_device(None);
        }
    }
}

#[allow(unused)]
pub struct DeviceInner {
    alloc: vk_mem::Allocator,
    device: DeviceWrapper,
    physical: PhysicalDevice,
    instance: Arc<Instance>,
    swapchain: khr::swapchain::Device,
    debug_utils: ash::ext::debug_utils::Device,
    allocated_queues: BTreeMap<(u32, u32), Queue>,
    // these are resident in allocated_queues, and may in fact be clones of each
    // other, for ease of access
    main_queue: Queue,
    compute_queue: Queue,
    transfer_queue: Queue,
    present_queue: Queue,
    sync_threadpool: sync::SyncThreadpool,
    features: crate::PhysicalDeviceFeatures,
}

impl core::fmt::Debug for DeviceInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DeviceInner")
            .field("device", &self.device.handle())
            .finish()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Extension<'a> {
    pub name: &'a str,
    pub version: u32,
}

impl<'a> std::hash::Hash for Extension<'a> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
    }
}

#[derive(Debug)]
pub struct DeviceDesc<'a> {
    pub app_name: Option<&'a str>,
    pub app_version: u32,
    pub layers: &'a [&'a CStr],
    pub layer_settings: &'a [vk::LayerSettingEXT<'a>],
    pub instance_extensions: &'a [Extension<'a>],
    pub display_handle: Option<RawDisplayHandle>,
    pub features: crate::PhysicalDeviceFeatures,
}

const VALIDATION_LAYER_NAME: &'static core::ffi::CStr = c"VK_LAYER_KHRONOS_validation";
const DEBUG_LAYERS: [&'static core::ffi::CStr; 1] = [VALIDATION_LAYER_NAME];
impl DeviceDesc<'_> {
    fn debug_layer_settings() -> &'static [vk::LayerSettingEXT<'static>; 3] {
        static SETTINGS: std::sync::LazyLock<[vk::LayerSettingEXT; 3]> =
            std::sync::LazyLock::new(|| {
                [
                    vk::LayerSettingEXT::default()
                        .layer_name(VALIDATION_LAYER_NAME)
                        .setting_name(c"VK_KHRONOS_VALIDATION_VALIDATE_BEST_PRACTICES")
                        .ty(vk::LayerSettingTypeEXT::BOOL32)
                        .values(&[1]),
                    vk::LayerSettingEXT::default()
                        .layer_name(VALIDATION_LAYER_NAME)
                        .setting_name(c"VK_KHRONOS_VALIDATION_VALIDATE_BEST_PRACTICES_AMD")
                        .ty(vk::LayerSettingTypeEXT::BOOL32)
                        .values(&[1]),
                    vk::LayerSettingEXT::default()
                        .layer_name(VALIDATION_LAYER_NAME)
                        .setting_name(c"VK_KHRONOS_VALIDATION_VALIDATE_SYNC")
                        .ty(vk::LayerSettingTypeEXT::BOOL32)
                        .values(&[1]),
                ]
            });
        &SETTINGS
    }
}

impl<'a> Default for DeviceDesc<'a> {
    fn default() -> Self {
        Self {
            app_name: Default::default(),
            app_version: Default::default(),
            #[cfg(debug_assertions)]
            layers: &DEBUG_LAYERS,
            #[cfg(debug_assertions)]
            layer_settings: Self::debug_layer_settings(),
            #[cfg(not(debug_assertions))]
            layers: &[],
            #[cfg(not(debug_assertions))]
            layer_settings: &[],
            instance_extensions: Default::default(),
            display_handle: Default::default(),
            features: Default::default(),
        }
    }
}

struct DeviceBuilder;

impl DeviceBuilder {
    fn queue_family_supports_presentation(
        instance: &Instance,
        pdev: vk::PhysicalDevice,
        queue_family: u32,
        display_handle: RawDisplayHandle,
    ) -> bool {
        unsafe {
            match display_handle {
                RawDisplayHandle::Xlib(display) => {
                    let surface =
                        ash::khr::xlib_surface::Instance::new(&instance.entry, &instance.instance);
                    surface.get_physical_device_xlib_presentation_support(
                        pdev,
                        queue_family,
                        display.display.unwrap().as_ptr() as _,
                        display.screen as _,
                    )
                    //todo!("xlib")
                }
                RawDisplayHandle::Xcb(_xcb_display_handle) => todo!("xcb"),
                RawDisplayHandle::Wayland(wayland_display_handle) => {
                    let surface = ash::khr::wayland_surface::Instance::new(
                        &instance.entry,
                        &instance.instance,
                    );
                    surface.get_physical_device_wayland_presentation_support(
                        pdev,
                        queue_family,
                        wayland_display_handle.display.cast().as_mut(),
                    )
                }
                RawDisplayHandle::Drm(_) => {
                    todo!()
                }
                RawDisplayHandle::Windows(_) => {
                    ash::khr::win32_surface::Instance::new(&instance.entry, &instance.instance)
                        .get_physical_device_win32_presentation_support(pdev, queue_family)
                }
                _ => panic!("unsupported platform"),
            }
        }
    }

    fn select_pdev_queue_families(
        instance: &Instance,
        display_handle: Option<RawDisplayHandle>,
        pdev: vk::PhysicalDevice,
    ) -> DeviceQueueFamilies {
        let queue_families = unsafe {
            instance
                .instance
                .get_physical_device_queue_family_properties(pdev)
        };

        struct QueueFamily {
            num_queues: u32,
            is_present: bool,
            is_compute: bool,
            is_graphics: bool,
            is_transfer: bool,
        }

        impl QueueFamily {
            #[allow(dead_code)]
            fn is_graphics_and_compute(&self) -> bool {
                self.is_compute && self.is_graphics
            }
        }

        struct QueueFamilies(Vec<QueueFamily>);
        impl QueueFamilies {
            fn find_first<F>(&mut self, mut pred: F) -> Option<u32>
            where
                F: FnMut(&QueueFamily) -> bool,
            {
                if let Some((q, family)) = self
                    .0
                    .iter_mut()
                    .enumerate()
                    .filter(|(_, family)| family.num_queues > 0)
                    .find(|(_, family)| pred(family))
                {
                    family.num_queues -= 1;
                    Some(q as u32)
                } else {
                    None
                }
            }

            fn find_best<F>(&mut self, mut pred: F) -> Option<u32>
            where
                F: FnMut(&QueueFamily) -> Option<u32>,
            {
                let (_, q, family) = self
                    .0
                    .iter_mut()
                    .enumerate()
                    .filter_map(|(i, family)| {
                        if family.num_queues == 0 {
                            return None;
                        }
                        pred(family).map(|score| (score, i, family))
                    })
                    .max_by_key(|(score, _, _)| *score)?;
                family.num_queues -= 1;
                Some(q as u32)
            }
        }

        let mut queue_families = QueueFamilies(
            queue_families
                .iter()
                .enumerate()
                .map(|(i, family)| {
                    let q = i as u32;
                    let is_graphics = family.queue_flags.contains(vk::QueueFlags::GRAPHICS);
                    let is_compute = family.queue_flags.contains(vk::QueueFlags::COMPUTE);
                    let is_transfer = family.queue_flags.contains(vk::QueueFlags::TRANSFER)
                        || is_compute
                        || is_graphics;
                    let is_present = display_handle
                        .map(|display_handle| {
                            Self::queue_family_supports_presentation(
                                instance,
                                pdev,
                                q,
                                display_handle,
                            )
                        })
                        .unwrap_or(false);
                    QueueFamily {
                        num_queues: family.queue_count,
                        is_compute,
                        is_graphics,
                        is_present,
                        is_transfer,
                    }
                })
                .collect::<Vec<_>>(),
        );

        let graphics = queue_families
            .find_best(|family| {
                if !family.is_graphics {
                    return None;
                }
                // a queue with Graphics+Compute is guaranteed to exist
                Some(family.is_compute as u32 * 2 + family.is_present as u32)
            })
            .unwrap();

        // find present queue first because it is rather more important than a secondary compute queue
        let present =
            if !queue_families.0.get(graphics as usize).unwrap().is_present {
                queue_families.find_first(|family| family.is_present)
            } else {
                None
            }.or({
                if display_handle.is_none() {
                    // in this case the graphics queue will be used by default
                    tracing::info!("no present queue available, using graphics queue as fallback for headless_surface");
                    Some(graphics)
                } else {
                    tracing::warn!("no present queue available, this is unexpected!");
                    None}
            });

        let async_compute = queue_families.find_first(|family| family.is_compute);
        let transfer = queue_families.find_first(|family| family.is_transfer);

        let mut unique_families = BTreeMap::<u32, u32>::new();

        let mut helper = |family: u32| {
            use std::collections::btree_map::Entry;
            let index = match unique_families.entry(family) {
                Entry::Vacant(vacant_entry) => {
                    vacant_entry.insert(1);
                    0
                }
                Entry::Occupied(mut occupied_entry) => {
                    let idx = occupied_entry.get_mut();
                    *idx += 1;
                    *idx - 1
                }
            };

            (family, index)
        };

        let graphics = helper(graphics);
        let async_compute = async_compute.map(|f| helper(f)).unwrap_or(graphics);
        let transfer = transfer.map(|f| helper(f)).unwrap_or(async_compute);
        let present = present.map(|f| helper(f)).unwrap_or(graphics);

        let families = unique_families
            .into_iter()
            .filter(|&(_family, count)| count > 0)
            .collect::<Vec<_>>();

        // family of each queue, of which one is allocated for each queue, with
        // graphics being the fallback queue for compute and transfer, and
        // present possibly being `None`, in which case it is Graphics
        let queues = DeviceQueueFamilies {
            families,
            graphics,
            async_compute,
            transfer,
            present,
        };

        queues
    }

    fn choose_physical_device(
        instance: &Instance,
        display_handle: Option<RawDisplayHandle>,
        requirements: &PhysicalDeviceFeatures,
        extra_properties: Vec<Box<dyn ExtendsDeviceProperties2Debug>>,
    ) -> Result<PhysicalDevice> {
        let pdevs = unsafe { instance.instance.enumerate_physical_devices()? };

        let (pdev, properties) = pdevs
            .into_iter()
            .map(|pdev| {
                let mut props = PhysicalDeviceProperties::default().extra_properties(
                    extra_properties
                        .iter()
                        .map(|b| dyn_clone::clone_box(&**b))
                        .collect::<Vec<_>>(),
                );
                props.query(&instance.instance, pdev);

                (pdev, props)
            })
            // filter devices which dont support the version of Vulkan we are requesting
            .filter(|(_, props)| props.base.api_version >= requirements.version)
            // filter devices which don't support the device extensions we
            // are requesting
            // TODO: figure out a way to fall back to some
            // device which doesn't support all of the extensions.
            .filter(|(pdev, _)| {
                let query_features =
                    PhysicalDeviceFeatures::query(&instance.instance, *pdev).unwrap();

                requirements.compatible_with(&query_features)
            })
            .max_by_key(|(_, props)| {
                let score = match props.base.device_type {
                    vk::PhysicalDeviceType::DISCRETE_GPU => 5,
                    vk::PhysicalDeviceType::INTEGRATED_GPU => 4,
                    vk::PhysicalDeviceType::VIRTUAL_GPU => 3,
                    vk::PhysicalDeviceType::CPU => 2,
                    vk::PhysicalDeviceType::OTHER => 1,
                    _ => unreachable!(),
                };

                // score based on limits or other properties

                score
            })
            .ok_or(Error::NoPhysicalDevice)?;

        Ok(PhysicalDevice {
            queue_families: Self::select_pdev_queue_families(instance, display_handle, pdev),
            pdev,
            properties,
        })
    }

    fn get_available_extensions(
        entry: &ash::Entry,
        layers: &[&CStr],
    ) -> Result<Vec<ash::vk::ExtensionProperties>> {
        unsafe {
            let extensions = core::iter::once(entry.enumerate_instance_extension_properties(None))
                .chain(
                    layers
                        .iter()
                        .map(|&layer| entry.enumerate_instance_extension_properties(Some(layer))),
                )
                .filter_map(|result| result.ok())
                .flatten()
                .collect::<Vec<ash::vk::ExtensionProperties>>();

            Ok(extensions)
        }
    }

    /// returns a tuple of supported/enabled extensions and unsupported/requested extensions
    fn get_extensions<'a>(
        entry: &ash::Entry,
        layers: &[&'a CStr],
        extensions: impl Iterator<Item = Extension<'a>> + 'a,
        display_handle: Option<RawDisplayHandle>,
    ) -> Result<(Vec<Extension<'a>>, Vec<Extension<'a>>)> {
        unsafe {
            let available_extensions = Self::get_available_extensions(entry, layers)?;

            let available_extension_names = available_extensions
                .iter()
                .filter_map(|ext| {
                    Some((
                        ext.extension_name_as_c_str().ok()?.to_str().ok()?,
                        ext.spec_version,
                    ))
                })
                .collect::<HashMap<_, _>>();

            let mut out_extensions = Vec::new();
            let mut unsupported_extensions = Vec::new();

            for extension in extensions {
                if let Some(available_version) = available_extension_names.get(&extension.name)
                    && *available_version >= extension.version
                {
                    out_extensions.push(extension);
                } else {
                    unsupported_extensions.push(extension);
                }
            }

            let required_extension_names = display_handle
                .map(|display_handle| ash_window::enumerate_required_extensions(display_handle))
                .unwrap_or(Ok(&[]))?;

            for &extension in required_extension_names {
                let extension = core::ffi::CStr::from_ptr(extension);
                let extension = Extension {
                    name: extension.to_str()?,
                    version: 0,
                };

                if let Some(available_version) = available_extension_names.get(&extension.name)
                    && *available_version >= extension.version
                {
                    out_extensions.push(extension);
                } else {
                    unsupported_extensions.push(extension);
                }
            }

            Ok((out_extensions, unsupported_extensions))
        }
    }

    fn get_layers<'a>(
        entry: &ash::Entry,
        wants_layers: impl Iterator<Item = &'a CStr> + 'a,
    ) -> core::result::Result<Vec<&'a CStr>, (Vec<&'a CStr>, Vec<&'a CStr>)> {
        unsafe {
            let wants_layers = wants_layers.collect::<Vec<_>>();

            let available_layers = entry
                .enumerate_instance_layer_properties()
                .map_err(|_| (Vec::<&'a CStr>::new(), wants_layers.clone()))?;
            let available_layer_names = available_layers
                .iter()
                .map(|layer| layer.layer_name_as_c_str())
                .collect::<core::result::Result<BTreeSet<_>, _>>()
                .map_err(|_| (Vec::<&'a CStr>::new(), wants_layers.clone()))?;

            let mut out_layers = Vec::new();
            let mut unsupported_layers = Vec::new();

            for layer in wants_layers {
                if available_layer_names.contains(&layer) {
                    out_layers.push(layer);
                } else {
                    unsupported_layers.push(layer);
                }
            }

            if !unsupported_layers.is_empty() {
                Err((out_layers, unsupported_layers))
            } else {
                Ok(out_layers)
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct Device(Arc<DeviceInner>);
pub type WeakDevice = std::sync::Weak<DeviceInner>;

impl Device {
    pub fn new_from_default_desc(
        display_handle: Option<RawDisplayHandle>,
        with_instance_extensions: &[Extension<'_>],
    ) -> crate::Result<Self> {
        Self::new_from_desc(DeviceDesc {
            app_name: Some("Vidya"),
            app_version: vk::make_api_version(0, 0, 1, 0),
            display_handle,
            features: crate::PhysicalDeviceFeatures::default()
                .version(vk::make_api_version(0, 1, 3, 0))
                .features10(
                    vk::PhysicalDeviceFeatures::default()
                        .sampler_anisotropy(true)
                        .fill_mode_non_solid(true)
                        .multi_draw_indirect(true),
                )
                .features11(
                    vk::PhysicalDeviceVulkan11Features::default().shader_draw_parameters(true),
                )
                .features12(
                    vk::PhysicalDeviceVulkan12Features::default()
                        .shader_int8(true)
                        .runtime_descriptor_array(true)
                        .descriptor_binding_partially_bound(true)
                        .shader_sampled_image_array_non_uniform_indexing(true)
                        .descriptor_binding_sampled_image_update_after_bind(true)
                        .storage_buffer8_bit_access(true),
                )
                .with_extension(
                    make_extention_properties(
                        ash::ext::mesh_shader::NAME,
                        ash::ext::mesh_shader::SPEC_VERSION,
                    ),
                    vk::PhysicalDeviceMeshShaderFeaturesEXT::default()
                        .mesh_shader(true)
                        .task_shader(true),
                )
                .with_extension(
                    make_extention_properties(
                        ash::ext::index_type_uint8::NAME,
                        ash::ext::index_type_uint8::SPEC_VERSION,
                    ),
                    vk::PhysicalDeviceIndexTypeUint8FeaturesEXT::default().index_type_uint8(true),
                )
                .with_extensions2([make_extention_properties(
                    khr::spirv_1_4::NAME,
                    khr::spirv_1_4::SPEC_VERSION,
                )]),
            instance_extensions: with_instance_extensions,
            ..Default::default()
        })
    }
    pub fn new_from_desc(desc: DeviceDesc) -> crate::Result<Self> {
        tracing::debug!("creating new device with: {desc:#?}");
        let entry = unsafe { ash::Entry::load()? };

        let app_name = desc
            .app_name
            .and_then(|name| CString::new(name).ok())
            .unwrap_or(c"ShooterGame".to_owned());

        let app_info = vk::ApplicationInfo::default()
            .api_version(desc.features.version)
            .application_name(&app_name)
            .application_version(desc.app_version)
            .engine_name(c"VidyaEngine")
            .engine_version(vk::make_api_version(0, 0, 1, 0));

        let mut validation_info =
            vk::LayerSettingsCreateInfoEXT::default().settings(&desc.layer_settings);

        cfg_if::cfg_if! {
            if #[cfg(debug_assertions)] {
                let debug_layers = [Extension {
                    name: "VK_EXT_layer_settings",
                    version: 2,
                }];
            } else {
                let debug_layers = [];
            }
        };

        let extra_instance_extensions = [Extension {
            name: "VK_EXT_debug_utils",
            version: 1,
        }]
        .into_iter()
        .chain(debug_layers.into_iter());

        let layers = DeviceBuilder::get_layers(&entry, desc.layers.into_iter().cloned()).unwrap();

        let (extensions, unsupported_extensions) = DeviceBuilder::get_extensions(
            &entry,
            &layers,
            desc.instance_extensions
                .into_iter()
                .cloned()
                .chain(extra_instance_extensions),
            desc.display_handle,
        )?;

        if !unsupported_extensions.is_empty() {
            tracing::error!(
                "extensions were requested but not supported by instance: {:?}",
                unsupported_extensions
            );
        }

        let layers = VkNameList::from_strs(&layers);
        let extensions = crate::util::CStringList::from_iter(extensions.iter().map(|ext| ext.name));

        let create_info = vk::InstanceCreateInfo::default()
            .application_info(&app_info)
            .enabled_extension_names(&extensions.strings)
            .enabled_layer_names(&layers.names)
            .push_next(&mut validation_info);

        tracing::debug!(
            "Creating instance:\napp_info: {app_info:#?}\ncreate_info: {create_info:#?}"
        );

        let instance = unsafe { entry.create_instance(&create_info, None)? };

        let debug_info = vk::DebugUtilsMessengerCreateInfoEXT::default()
            .message_severity(
                vk::DebugUtilsMessageSeverityFlagsEXT::ERROR
                    | vk::DebugUtilsMessageSeverityFlagsEXT::WARNING
                    | vk::DebugUtilsMessageSeverityFlagsEXT::INFO,
            )
            .message_type(
                vk::DebugUtilsMessageTypeFlagsEXT::GENERAL
                    | vk::DebugUtilsMessageTypeFlagsEXT::VALIDATION
                    | vk::DebugUtilsMessageTypeFlagsEXT::PERFORMANCE,
            )
            .pfn_user_callback(Some(crate::debug::debug_callback));

        let debug_utils_instance = ash::ext::debug_utils::Instance::new(&entry, &instance);
        let debug_utils_messenger =
            unsafe { debug_utils_instance.create_debug_utils_messenger(&debug_info, None)? };

        let surface_instance = ash::khr::surface::Instance::new(&entry, &instance);

        let instance = Arc::new(Instance {
            instance,
            debug_utils: debug_utils_instance,
            debug_utils_messenger,
            surface: surface_instance,
            entry,
        });

        let mut features = desc.features.with_extension2(make_extention_properties(
            khr::swapchain::NAME,
            khr::swapchain::SPEC_VERSION,
        ));

        //these are required for the renderpass
        let features13 = features.physical_features_13.get_or_insert_default();
        features13.synchronization2 = vk::TRUE;
        features13.dynamic_rendering = vk::TRUE;
        features13.maintenance4 = vk::TRUE;

        // Consider this: switching physical device in game?
        // anything above this point is device agnostic, everything below would have to be recreated
        // additionally, pdev would have to be derived from a device and not a scoring function.

        let pdev = DeviceBuilder::choose_physical_device(
            &instance,
            desc.display_handle,
            &features,
            vec![Box::new(
                vk::PhysicalDeviceMeshShaderPropertiesEXT::default(),
            )],
        )?;

        tracing::trace!("pdev: {pdev:?}");
        let device = Device::new(instance.clone(), pdev, features)?;

        Ok(device)
    }

    pub fn new(
        instance: Arc<Instance>,
        physical: PhysicalDevice,
        mut features: crate::PhysicalDeviceFeatures,
    ) -> VkResult<Self> {
        // we have 4 queues at most: graphics, compute, transfer, present
        let priorities = [1.0f32; 4];

        let queue_infos = physical
            .queue_families
            .families
            .iter()
            .map(|&(family, queues)| {
                vk::DeviceQueueCreateInfo::default()
                    .queue_family_index(family)
                    .queue_priorities(&priorities[..queues as usize])
            })
            .collect::<Vec<_>>();

        let extensions = features
            .device_extensions
            .iter()
            .map(|ext| ext.extension_name.as_ptr())
            .collect::<Vec<_>>();

        let mut features2 = features.features2();
        let device_info = vk::DeviceCreateInfo::default()
            .queue_create_infos(&queue_infos)
            .enabled_extension_names(&extensions)
            .push_next(&mut features2);

        let device = unsafe {
            let device = instance
                .instance
                .create_device(physical.pdev, &device_info, None)?;

            let allocated_queues = queue_infos
                .iter()
                .flat_map(|info| {
                    (0..info.queue_count).map(|i| {
                        (
                            (info.queue_family_index, i),
                            Queue::new(&device, info.queue_family_index, i),
                        )
                    })
                })
                .collect::<BTreeMap<_, _>>();

            let get_queue =
                |(family, index)| allocated_queues.get(&(family, index)).cloned().unwrap();

            let main_queue = get_queue(physical.queue_families.graphics);
            let present_queue = get_queue(physical.queue_families.present);
            let compute_queue = get_queue(physical.queue_families.async_compute);
            let transfer_queue = get_queue(physical.queue_families.transfer);

            let alloc_info =
                vk_mem::AllocatorCreateInfo::new(&instance.instance, &device, physical.pdev);

            let alloc = vk_mem::Allocator::new(alloc_info)?;

            DeviceInner {
                device: DeviceWrapper(device.clone()),
                physical,
                swapchain: khr::swapchain::Device::new(&instance.instance, &device),
                debug_utils: ash::ext::debug_utils::Device::new(&instance.instance, &device),
                instance,
                alloc,
                allocated_queues,
                main_queue,
                present_queue,
                compute_queue,
                transfer_queue,
                features,
                sync_threadpool: sync::SyncThreadpool::new(),
            }
        };

        Ok(Self(Arc::new(device)))
    }
    pub fn sync_threadpool(&self) -> &sync::SyncThreadpool {
        &self.0.sync_threadpool
    }
    pub fn weak(&self) -> WeakDevice {
        Arc::downgrade(&self.0)
    }
    pub fn alloc(&self) -> &vk_mem::Allocator {
        &self.0.alloc
    }
    pub fn dev(&self) -> &ash::Device {
        &self.0.device
    }
    pub fn instance(&self) -> &Arc<Instance> {
        &self.0.instance
    }
    pub fn swapchain(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
    pub fn debug_utils(&self) -> &ash::ext::debug_utils::Device {
        &self.0.debug_utils
    }
    pub fn queue_families(&self) -> &DeviceQueueFamilies {
        &self.0.physical.queue_families
    }
    pub fn phy(&self) -> vk::PhysicalDevice {
        self.0.physical.pdev
    }
    pub fn features(&self) -> &crate::PhysicalDeviceFeatures {
        &self.0.features
    }
    pub fn physical_device(&self) -> &PhysicalDevice {
        &self.0.physical
    }
    pub fn graphics_queue(&self) -> &Queue {
        &self.0.main_queue
    }
    pub fn transfer_queue(&self) -> &Queue {
        &self.0.transfer_queue
    }
    pub fn present_queue(&self) -> &Queue {
        &self.0.present_queue
    }

    pub unsafe fn lock_queues(&self) {
        // this is obviously awful, allocating for this
        self.0
            .allocated_queues
            .values()
            .for_each(|q| core::mem::forget(q.lock()));
    }

    pub unsafe fn unlock_queues(&self) {
        self.0
            .allocated_queues
            .values()
            .for_each(|q| unsafe { q.0.force_unlock() });
    }

    pub fn wait_queue_idle(&self, queue: &Queue) -> VkResult<()> {
        tracing::warn!("locking queue {queue:?} and waiting for idle");

        queue.with_locked(|q| unsafe { self.dev().queue_wait_idle(q) })?;

        tracing::warn!("finished waiting: unlocking queue {queue:?}.");
        Ok(())
    }

    pub fn wait_idle(&self) -> VkResult<()> {
        tracing::warn!("locking all queues and waiting for device to idle");
        unsafe {
            self.lock_queues();
            self.dev().device_wait_idle()?;
            self.unlock_queues();
        }
        tracing::warn!("finished waiting: unlocking all queues.");
        Ok(())
    }

    pub fn debug_name_object<T: vk::Handle>(&self, handle: T, name: &str) -> VkResult<()> {
        let name = std::ffi::CString::new(name.as_bytes()).unwrap_or(c"invalid name".to_owned());
        unsafe {
            self.debug_utils().set_debug_utils_object_name(
                &vk::DebugUtilsObjectNameInfoEXT::default()
                    .object_handle(handle)
                    .object_name(&name),
            )?;
        }
        Ok(())
    }
}

impl AsRef<khr::swapchain::Device> for Device {
    fn as_ref(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
}

impl AsRef<ash::Device> for Device {
    fn as_ref(&self) -> &ash::Device {
        &self.0.device
    }
}

#[allow(dead_code)]
pub struct DeviceAndQueues {
    pub(crate) device: Device,
    pub(crate) main_queue: Queue,
    pub(crate) compute_queue: Queue,
    pub(crate) transfer_queue: Queue,
    pub(crate) present_queue: Queue,
}

impl AsRef<ash::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::Device {
        &self.device.as_ref()
    }
}
impl AsRef<ash::khr::swapchain::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::khr::swapchain::Device {
        &self.device.as_ref()
    }
}

#[derive(Clone)]
pub struct DeviceOwnedDebugObject<T> {
    pub(crate) device: Device,
    pub(crate) object: T,
    #[cfg(debug_assertions)]
    name: Option<Cow<'static, str>>,
}

impl<T: Eq> Eq for DeviceOwnedDebugObject<T> {}
impl<T: PartialEq> PartialEq for DeviceOwnedDebugObject<T> {
    fn eq(&self, other: &Self) -> bool {
        std::sync::Arc::ptr_eq(&self.device.0, &other.device.0) && self.object == other.object
    }
}

impl<T: std::fmt::Debug + vk::Handle + Copy> std::fmt::Debug for DeviceOwnedDebugObject<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut fmt = f.debug_struct(core::any::type_name::<T>());

        fmt.field_with("device", |f| {
            write!(f, "0x{:x}", self.device.0.device.handle().as_raw())
        })
        .field_with("handle", |f| write!(f, "0x{:x}", &self.object.as_raw()));

        #[cfg(debug_assertions)]
        {
            fmt.field("name", &self.name);
        }

        fmt.finish()
    }
}

impl<T> DeviceOwnedDebugObject<T> {
    pub fn new(
        device: crate::Device,
        object: T,
        name: Option<Cow<'static, str>>,
    ) -> ash::prelude::VkResult<Self>
    where
        T: vk::Handle + Copy,
    {
        if let Some(name) = name.as_ref() {
            device.debug_name_object(object, name)?;
        }

        Ok(Self {
            device,
            object,
            #[cfg(debug_assertions)]
            name,
        })
    }

    pub fn dev(&self) -> &crate::Device {
        &self.device
    }
    pub fn handle(&self) -> T
    where
        T: Copy,
    {
        self.object
    }
}

pub trait DeviceOwned<T> {
    fn device(&self) -> &Device;
    fn handle(&self) -> T;
}

#[macro_export]
macro_rules! define_device_owned_handle {
    ($(#[$attr:meta])*
        $ty_vis:vis $ty:ident($handle:ty) {
        $($(#[$field_attr:meta])* $field_vis:vis $field_name:ident : $field_ty:ty),*
        $(,)?
    } $(=> |$this:ident| $dtor:stmt)?) => {
        $(#[$attr])*
            $ty_vis struct $ty {
                inner: crate::device::DeviceOwnedDebugObject<$handle>,
                $(
                    $(#[$field_attr])*
                    $field_vis $field_name: $field_ty,
                )*
            }

        impl crate::device::DeviceOwned<$handle> for $ty {
            fn device(&self) -> &crate::device::Device {
                self.inner.dev()
            }
            fn handle(&self) -> $handle {
                self.inner.handle()
            }
        }

        impl $ty {
            fn construct(
                device: crate::device::Device,
                handle: $handle,
                name: Option<::std::borrow::Cow<'static, str>>,
                $($field_name: $field_ty,)*
            ) -> ::ash::prelude::VkResult<Self> {
                Ok(Self {
                    inner: crate::device::DeviceOwnedDebugObject::new(
                        device,
                        handle,
                        name,
                    )?,
                    $($field_name,)*
                })
            }
        }

        $(
            impl Drop for $ty {
                fn drop(&mut self) {
                    #[allow(unused_mut)]
                    let mut $this = self;
                    $dtor
                }
            }
        )?
    };
}
