monolith

use crate::Error;

use crate::device::Shared;

use crate::format::{SUPPORTED, vk_format};

use crate::image::{ImageInner, Memory};

use ash::vk;

use smithay::backend::allocator::dmabuf::{Dmabuf, DmabufFlags};

use smithay::backend::allocator::format::FormatSet;

use smithay::backend::allocator::{Allocator, Buffer, Format, Fourcc, Modifier};

use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd};

use std::sync::Arc;

/// Which (format, modifier) pairs the device can import for sampling and

/// render into, queried once at startup.

#[derive(Debug, Default)]

pub(crate) struct DmabufCaps {

    pub sample: FormatSet,

    pub render: FormatSet,

}

pub(crate) const RENDER_USAGE: vk::ImageUsageFlags = vk::ImageUsageFlags::from_raw(

    vk::ImageUsageFlags::COLOR_ATTACHMENT.as_raw()

        | vk::ImageUsageFlags::SAMPLED.as_raw()

        | vk::ImageUsageFlags::TRANSFER_SRC.as_raw(),

);

pub(crate) const SAMPLE_USAGE: vk::ImageUsageFlags = vk::ImageUsageFlags::SAMPLED;

impl DmabufCaps {

    pub fn query(shared: &Shared) -> Self {

        if shared.device.dmabuf.is_none() {

            return Self::default();

        }

        let (mut sample, mut render) = (Vec::new(), Vec::new());

        for fourcc in SUPPORTED {

            let Some(fmt) = vk_format(fourcc) else {

                continue;

            };

            for props in modifier_properties(shared, fmt.format) {

                if props.drm_format_modifier_plane_count != 1 {

                    continue;

                }

                let features = props.drm_format_modifier_tiling_features;

                let format = Format {

                    code: fourcc,

                    modifier: Modifier::from(props.drm_format_modifier),

                };

                let supports =

                    |usage| external_support(shared, fmt.format, props.drm_format_modifier, usage);

                if features.contains(vk::FormatFeatureFlags::SAMPLED_IMAGE)

                    && supports(SAMPLE_USAGE).importable

                {

                    sample.push(format);

                }

                let renderable = vk::FormatFeatureFlags::COLOR_ATTACHMENT

                    | vk::FormatFeatureFlags::COLOR_ATTACHMENT_BLEND;

                let s = supports(RENDER_USAGE);

                if features.contains(renderable) && s.importable && s.exportable {

                    render.push(format);

                }

            }

        }

        Self {

            sample: sample.into_iter().collect(),

            render: render.into_iter().collect(),

        }

    }

}

fn modifier_properties(

    shared: &Shared,

    format: vk::Format,

) -> Vec<vk::DrmFormatModifierPropertiesEXT> {

    let dev = &shared.device;

    let mut list = vk::DrmFormatModifierPropertiesListEXT::default();

    let mut props = vk::FormatProperties2::default().push_next(&mut list);

    unsafe {

        dev.instance

            .get_physical_device_format_properties2(dev.physical, format, &mut props)

    };

    let count = list.drm_format_modifier_count as usize;

    let mut out = vec![vk::DrmFormatModifierPropertiesEXT::default(); count];

    let mut list =

        vk::DrmFormatModifierPropertiesListEXT::default().drm_format_modifier_properties(&mut out);

    let mut props = vk::FormatProperties2::default().push_next(&mut list);

    unsafe {

        dev.instance

            .get_physical_device_format_properties2(dev.physical, format, &mut props)

    };

    out

}

struct External {

    importable: bool,

    exportable: bool,

}

fn external_support(

    shared: &Shared,

    format: vk::Format,

    modifier: u64,

    usage: vk::ImageUsageFlags,

) -> External {

    let dev = &shared.device;

    let mut modifier_info = vk::PhysicalDeviceImageDrmFormatModifierInfoEXT::default()

        .drm_format_modifier(modifier)

        .sharing_mode(vk::SharingMode::EXCLUSIVE);

    let mut external_info = vk::PhysicalDeviceExternalImageFormatInfo::default()

        .handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT);

    let info = vk::PhysicalDeviceImageFormatInfo2::default()

        .format(format)

        .ty(vk::ImageType::TYPE_2D)

        .tiling(vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT)

        .usage(usage)

        .push_next(&mut external_info)

        .push_next(&mut modifier_info);

    let mut external_props = vk::ExternalImageFormatProperties::default();

    let mut props = vk::ImageFormatProperties2::default().push_next(&mut external_props);

    let ok = unsafe {

        dev.instance

            .get_physical_device_image_format_properties2(dev.physical, &info, &mut props)

    }

    .is_ok();

    let features = external_props

        .external_memory_properties

        .external_memory_features;

    External {

        importable: ok && features.contains(vk::ExternalMemoryFeatureFlags::IMPORTABLE),

        exportable: ok && features.contains(vk::ExternalMemoryFeatureFlags::EXPORTABLE),

    }

}

fn memory_type(shared: &Shared, bits: u32, prefer: vk::MemoryPropertyFlags) -> Option<u32> {

    let props = unsafe {

        shared

            .device

            .instance

            .get_physical_device_memory_properties(shared.device.physical)

    };

    let types = &props.memory_types[..props.memory_type_count as usize];

    let allowed = |i: usize| bits & (1 << i) != 0;

    (0..types.len())

        .find(|&i| allowed(i) && types[i].property_flags.contains(prefer))

        .or_else(|| (0..types.len()).find(|&i| allowed(i)))

        .map(|i| i as u32)

}

fn image_info<'a>(

    fourcc: Fourcc,

    size: (u32, u32),

    usage: vk::ImageUsageFlags,

) -> Result<vk::ImageCreateInfo<'a>, Error> {

    let fmt = vk_format(fourcc).ok_or(Error::UnsupportedFormat(fourcc))?;

    Ok(vk::ImageCreateInfo::default()

        .image_type(vk::ImageType::TYPE_2D)

        .format(fmt.format)

        .extent(vk::Extent3D {

            width: size.0,

            height: size.1,

            depth: 1,

        })

        .mip_levels(1)

        .array_layers(1)

        .samples(vk::SampleCountFlags::TYPE_1)

        .tiling(vk::ImageTiling::DRM_FORMAT_MODIFIER_EXT)

        .usage(usage)

        .sharing_mode(vk::SharingMode::EXCLUSIVE)

        .initial_layout(vk::ImageLayout::UNDEFINED))

}

/// Imports a single-plane dmabuf as an image, without copying.

pub(crate) fn import(

    shared: &Arc<Shared>,

    caps: &DmabufCaps,

    dmabuf: &Dmabuf,

    usage: vk::ImageUsageFlags,

) -> Result<Arc<ImageInner>, Error> {

    let ext = shared

        .device

        .dmabuf

        .as_ref()

        .ok_or(Error::DmabufUnsupported)?;

    let format = dmabuf.format();

    let set = if usage.contains(vk::ImageUsageFlags::COLOR_ATTACHMENT) {

        &caps.render

    } else {

        &caps.sample

    };

    if dmabuf.num_planes() != 1 || !set.contains(&format) {

        return Err(Error::UnsupportedDmabuf(format));

    }

    let (Some(fd), Some(offset), Some(stride)) = (

        dmabuf.handles().next(),

        dmabuf.offsets().next(),

        dmabuf.strides().next(),

    ) else {

        return Err(Error::UnsupportedDmabuf(format));

    };

    let size = (dmabuf.width(), dmabuf.height());

    let d = &shared.device.device;

    let layouts = [vk::SubresourceLayout {

        offset: u64::from(offset),

        size: 0,

        row_pitch: u64::from(stride),

        array_pitch: 0,

        depth_pitch: 0,

    }];

    let mut explicit = vk::ImageDrmFormatModifierExplicitCreateInfoEXT::default()

        .drm_format_modifier(u64::from(format.modifier))

        .plane_layouts(&layouts);

    let mut external = vk::ExternalMemoryImageCreateInfo::default()

        .handle_types(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT);

    let info = image_info(format.code, size, usage)?

        .push_next(&mut external)

        .push_next(&mut explicit);

    let image = unsafe { d.create_image(&info, None)? };

    let destroy = || unsafe { d.destroy_image(image, None) };

    let owned = match fd.try_clone_to_owned() {

        Ok(o) => o,

        Err(_) => {

            destroy();

            return Err(Error::UnsupportedDmabuf(format));

        }

    };

    let mut fd_props = vk::MemoryFdPropertiesKHR::default();

    if let Err(e) = unsafe {

        ext.memory_fd.get_memory_fd_properties(

            vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT,

            owned.as_raw_fd(),

            &mut fd_props,

        )

    } {

        destroy();

        return Err(e.into());

    }

    let requirements = unsafe { d.get_image_memory_requirements(image) };

    let Some(type_index) = memory_type(

        shared,

        requirements.memory_type_bits & fd_props.memory_type_bits,

        vk::MemoryPropertyFlags::empty(),

    ) else {

        destroy();

        return Err(Error::UnsupportedDmabuf(format));

    };

    let raw = owned.into_raw_fd();

    let mut import = vk::ImportMemoryFdInfoKHR::default()

        .handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT)

        .fd(raw);

    let mut dedicated = vk::MemoryDedicatedAllocateInfo::default().image(image);

    let alloc = vk::MemoryAllocateInfo::default()

        .allocation_size(requirements.size)

        .memory_type_index(type_index)

        .push_next(&mut import)

        .push_next(&mut dedicated);

    let memory = match unsafe { d.allocate_memory(&alloc, None) } {

        Ok(m) => m,

        Err(e) => {

            // SAFETY: a failed import leaves the fd with us, per

            // VK_KHR_external_memory_fd.

            drop(unsafe { OwnedFd::from_raw_fd(raw) });

            destroy();

            return Err(e.into());

        }

    };

    ImageInner::finish(

        shared,

        image,

        Memory::Imported(memory),

        memory,

        0,

        format.code,

        size.0,

        size.1,

        usage,

        true,

    )

}

/// Allocates a dmabuf the device can render into, picking among `modifiers`.

pub(crate) fn export(

    shared: &Arc<Shared>,

    caps: &DmabufCaps,

    fourcc: Fourcc,

    size: (u32, u32),

    modifiers: &[Modifier],

) -> Result<Dmabuf, Error> {

    let ext = shared

        .device

        .dmabuf

        .as_ref()

        .ok_or(Error::DmabufUnsupported)?;

    let allowed: Vec<u64> = modifiers

        .iter()

        .filter(|&&modifier| {

            caps.render.contains(&Format {

                code: fourcc,

                modifier,

            })

        })

        .map(|&m| u64::from(m))

        .collect();

    if allowed.is_empty() {

        return Err(Error::UnsupportedFormat(fourcc));

    }

    let d = &shared.device.device;

    let mut list =

        vk::ImageDrmFormatModifierListCreateInfoEXT::default().drm_format_modifiers(&allowed);

    let mut external = vk::ExternalMemoryImageCreateInfo::default()

        .handle_types(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT);

    let info = image_info(fourcc, size, RENDER_USAGE)?

        .push_next(&mut external)

        .push_next(&mut list);

    let image = unsafe { d.create_image(&info, None)? };

    let result = (|| -> Result<Dmabuf, Error> {

        let requirements = unsafe { d.get_image_memory_requirements(image) };

        let type_index = memory_type(

            shared,

            requirements.memory_type_bits,

            vk::MemoryPropertyFlags::DEVICE_LOCAL,

        )

        .ok_or(Error::UnsupportedFormat(fourcc))?;

        let mut export = vk::ExportMemoryAllocateInfo::default()

            .handle_types(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT);

        let mut dedicated = vk::MemoryDedicatedAllocateInfo::default().image(image);

        let alloc = vk::MemoryAllocateInfo::default()

            .allocation_size(requirements.size)

            .memory_type_index(type_index)

            .push_next(&mut export)

            .push_next(&mut dedicated);

        let memory = unsafe { d.allocate_memory(&alloc, None)? };

        let exported = (|| -> Result<Dmabuf, Error> {

            unsafe { d.bind_image_memory(image, memory, 0)? };

            let raw = unsafe {

                ext.memory_fd.get_memory_fd(

                    &vk::MemoryGetFdInfoKHR::default()

                        .memory(memory)

                        .handle_type(vk::ExternalMemoryHandleTypeFlags::DMA_BUF_EXT),

                )?

            };

            // SAFETY: vkGetMemoryFdKHR hands the caller a new fd to own.

            let fd = unsafe { OwnedFd::from_raw_fd(raw) };

            let mut props = vk::ImageDrmFormatModifierPropertiesEXT::default();

            unsafe {

                ext.modifier

                    .get_image_drm_format_modifier_properties(image, &mut props)?

            };

            let layout = unsafe {

                d.get_image_subresource_layout(

                    image,

                    vk::ImageSubresource::default()

                        .aspect_mask(vk::ImageAspectFlags::MEMORY_PLANE_0_EXT),

                )

            };

            let mut builder = Dmabuf::builder(

                (size.0 as i32, size.1 as i32),

                fourcc,

                Modifier::from(props.drm_format_modifier),

                DmabufFlags::empty(),

            );

            builder.add_plane(fd, layout.offset as u32, layout.row_pitch as u32);

            builder.build().ok_or(Error::UnsupportedFormat(fourcc))

        })();

        unsafe { d.free_memory(memory, None) };

        exported

    })();

    unsafe { d.destroy_image(image, None) };

    result

}

/// Allocates render-capable dmabufs, for swapchains.

#[derive(Clone, Debug)]

pub struct VulkanAllocator {

    pub(crate) shared: Arc<Shared>,

    pub(crate) caps: Arc<DmabufCaps>,

}

impl Allocator for VulkanAllocator {

    type Buffer = Dmabuf;

    type Error = Error;

    fn create_buffer(

        &mut self,

        width: u32,

        height: u32,

        fourcc: Fourcc,

        modifiers: &[Modifier],

    ) -> Result<Dmabuf, Error> {

        export(&self.shared, &self.caps, fourcc, (width, height), modifiers)

    }

}

impl std::fmt::Debug for Shared {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("Shared")

            .field("device", &self.device.name)

            .finish()

    }

}