monolith

use crate::Error;

use crate::device::Shared;

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

use ash::vk;

use gpu_allocator::MemoryLocation;

use gpu_allocator::vulkan::{Allocation, AllocationCreateDesc, AllocationScheme};

use smithay::backend::allocator::Fourcc;

use smithay::backend::renderer::{Texture, TextureMapping};

use std::fmt;

use std::marker::PhantomData;

use std::sync::{Arc, Mutex};

use std::time::{Duration, Instant};

pub(crate) enum Memory {

    Allocated(Allocation),

    /// Memory imported from a dmabuf, owned directly.

    Imported(vk::DeviceMemory),

}

pub(crate) struct ImageInner {

    pub shared: Arc<Shared>,

    pub image: vk::Image,

    pub view: vk::ImageView,

    pub memory: Mutex<Option<Memory>>,

    pub fourcc: Fourcc,

    pub format: vk::Format,

    pub width: u32,

    pub height: u32,

    pub layout: Mutex<vk::ImageLayout>,

    /// Shared with other devices or processes through a dmabuf. Every use

    /// acquires the image from the foreign queue family and releases it after.

    pub external: bool,

}

impl fmt::Debug for ImageInner {

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

        f.debug_struct("ImageInner")

            .field("image", &self.image)

            .field("fourcc", &self.fourcc)

            .field("size", &(self.width, self.height))

            .field("external", &self.external)

            .finish()

    }

}

impl Drop for ImageInner {

    fn drop(&mut self) {

        let d = &self.shared.device.device;

        unsafe {

            d.destroy_image_view(self.view, None);

            d.destroy_image(self.image, None);

        }

        match self.memory.lock().unwrap_or_else(|e| e.into_inner()).take() {

            Some(Memory::Allocated(a)) => {

                let _ = self

                    .shared

                    .allocator

                    .lock()

                    .unwrap_or_else(|e| e.into_inner())

                    .free(a);

            }

            Some(Memory::Imported(m)) => unsafe { d.free_memory(m, None) },

            None => {}

        }

    }

}

impl ImageInner {

    pub fn create(

        shared: &Arc<Shared>,

        fourcc: Fourcc,

        width: u32,

        height: u32,

        usage: vk::ImageUsageFlags,

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

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

        if width == 0 || height == 0 {

            return Err(Error::InvalidSize(width as i32, height as i32));

        }

        let d = &shared.device.device;

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

            .image_type(vk::ImageType::TYPE_2D)

            .format(fmt.format)

            .extent(vk::Extent3D {

                width,

                height,

                depth: 1,

            })

            .mip_levels(1)

            .array_layers(1)

            .samples(vk::SampleCountFlags::TYPE_1)

            .tiling(vk::ImageTiling::OPTIMAL)

            .usage(usage)

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

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

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

        let allocation = shared

            .allocator

            .lock()

            .unwrap_or_else(|e| e.into_inner())

            .allocate(&AllocationCreateDesc {

                name: "image",

                requirements,

                location: MemoryLocation::GpuOnly,

                linear: false,

                allocation_scheme: AllocationScheme::GpuAllocatorManaged,

            });

        let allocation = match allocation {

            Ok(a) => a,

            Err(e) => {

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

                return Err(e.into());

            }

        };

        let memory = Memory::Allocated(allocation);

        let (mem, offset) = match &memory {

            Memory::Allocated(a) => (unsafe { a.memory() }, a.offset()),

            Memory::Imported(m) => (*m, 0),

        };

        Self::finish(

            shared, image, memory, mem, offset, fourcc, width, height, false,

        )

    }

    /// Binds memory, creates the view and wraps everything up. Takes

    /// ownership of `image` and `memory`, releasing both on failure.

    #[expect(clippy::too_many_arguments)]

    pub fn finish(

        shared: &Arc<Shared>,

        image: vk::Image,

        memory: Memory,

        mem: vk::DeviceMemory,

        offset: u64,

        fourcc: Fourcc,

        width: u32,

        height: u32,

        external: bool,

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

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

        let d = &shared.device.device;

        let mut inner = Self {

            shared: shared.clone(),

            image,

            view: vk::ImageView::null(),

            memory: Mutex::new(Some(memory)),

            fourcc,

            format: fmt.format,

            width,

            height,

            layout: Mutex::new(if external {

                vk::ImageLayout::GENERAL

            } else {

                vk::ImageLayout::UNDEFINED

            }),

            external,

        };

        unsafe { d.bind_image_memory(image, mem, offset)? };

        let view_info = vk::ImageViewCreateInfo::default()

            .image(image)

            .view_type(vk::ImageViewType::TYPE_2D)

            .format(fmt.format)

            .components(components(fmt.opaque))

            .subresource_range(color_range());

        inner.view = unsafe { d.create_image_view(&view_info, None)? };

        Ok(Arc::new(inner))

    }

    pub fn layout(&self) -> vk::ImageLayout {

        *self.layout.lock().unwrap_or_else(|e| e.into_inner())

    }

    pub fn set_layout(&self, layout: vk::ImageLayout) {

        *self.layout.lock().unwrap_or_else(|e| e.into_inner()) = layout;

    }

}

pub(crate) fn color_range() -> vk::ImageSubresourceRange {

    vk::ImageSubresourceRange::default()

        .aspect_mask(vk::ImageAspectFlags::COLOR)

        .level_count(1)

        .layer_count(1)

}

/// A host-visible buffer for uploads and readbacks. Dropped ones go back to

/// a pool on `Shared` instead of being freed. Allocating and freeing large

/// host-visible memory takes a driver lock on NVIDIA that `vkQueueSubmit`

/// also waits on, so a staging buffer made per upload on the import worker

/// stalled the main thread's submits for 1 to 3 ms.

pub(crate) struct Staging {

    pub shared: Arc<Shared>,

    pub buffer: vk::Buffer,

    pub allocation: Option<Allocation>,

    size: usize,

    usage: vk::BufferUsageFlags,

    location: MemoryLocation,

}

pub(crate) struct Pooled {

    pub buffer: vk::Buffer,

    pub allocation: Allocation,

    size: usize,

    usage: vk::BufferUsageFlags,

    location: MemoryLocation,

    returned: Instant,

}

/// Pooled staging memory kept at most, past which dropped buffers are freed.

const POOL_BYTES: usize = 32 * 1024 * 1024;

/// How long a pooled buffer may go unused before it is freed.

const POOL_IDLE: Duration = Duration::from_secs(2);

/// Frees pooled staging buffers that have sat unused for [`POOL_IDLE`].

pub(crate) fn trim_staging(shared: &Shared) {

    let now = Instant::now();

    let expired: Vec<Pooled> = shared

        .staging_pool

        .lock()

        .unwrap_or_else(|e| e.into_inner())

        .extract_if(.., |p| now - p.returned >= POOL_IDLE)

        .collect();

    if expired.is_empty() {

        return;

    }

    let mut allocator = shared.allocator.lock().unwrap_or_else(|e| e.into_inner());

    for p in expired {

        unsafe { shared.device.device.destroy_buffer(p.buffer, None) };

        let _ = allocator.free(p.allocation);

    }

}

impl Staging {

    pub fn new(

        shared: &Arc<Shared>,

        size: usize,

        usage: vk::BufferUsageFlags,

        location: MemoryLocation,

    ) -> Result<Self, Error> {

        let reused = {

            let mut pool = shared

                .staging_pool

                .lock()

                .unwrap_or_else(|e| e.into_inner());

            let fits = |p: &Pooled| {

                p.usage == usage && p.location == location && p.size >= size && p.size <= size * 2

            };

            pool.iter().position(fits).map(|i| pool.swap_remove(i))

        };

        if let Some(p) = reused {

            return Ok(Self {

                shared: shared.clone(),

                buffer: p.buffer,

                allocation: Some(p.allocation),

                size: p.size,

                usage,

                location,

            });

        }

        let d = &shared.device.device;

        let buffer = unsafe {

            d.create_buffer(

                &vk::BufferCreateInfo::default()

                    .size(size as u64)

                    .usage(usage),

                None,

            )?

        };

        let requirements = unsafe { d.get_buffer_memory_requirements(buffer) };

        let allocation = shared

            .allocator

            .lock()

            .unwrap_or_else(|e| e.into_inner())

            .allocate(&AllocationCreateDesc {

                name: "staging",

                requirements,

                location,

                linear: true,

                allocation_scheme: AllocationScheme::GpuAllocatorManaged,

            });

        let allocation = match allocation {

            Ok(a) => a,

            Err(e) => {

                unsafe { d.destroy_buffer(buffer, None) };

                return Err(e.into());

            }

        };

        unsafe { d.bind_buffer_memory(buffer, allocation.memory(), allocation.offset())? };

        Ok(Self {

            shared: shared.clone(),

            buffer,

            allocation: Some(allocation),

            size,

            usage,

            location,

        })

    }

    pub fn bytes(&self) -> &[u8] {

        self.allocation

            .as_ref()

            .and_then(|a| a.mapped_slice())

            .unwrap_or(&[])

    }

    pub fn bytes_mut(&mut self) -> &mut [u8] {

        self.allocation

            .as_mut()

            .and_then(|a| a.mapped_slice_mut())

            .unwrap_or(&mut [])

    }

}

impl Drop for Staging {

    fn drop(&mut self) {

        let Some(allocation) = self.allocation.take() else {

            return;

        };

        let mut pool = self

            .shared

            .staging_pool

            .lock()

            .unwrap_or_else(|e| e.into_inner());

        let pooled: usize = pool.iter().map(|p| p.size).sum();

        if pooled + self.size <= POOL_BYTES {

            pool.push(Pooled {

                buffer: self.buffer,

                allocation,

                size: self.size,

                usage: self.usage,

                location: self.location,

                returned: Instant::now(),

            });

            return;

        }

        drop(pool);

        unsafe { self.shared.device.device.destroy_buffer(self.buffer, None) };

        let _ = self

            .shared

            .allocator

            .lock()

            .unwrap_or_else(|e| e.into_inner())

            .free(allocation);

    }

}

/// A sampled texture.

#[derive(Clone, Debug)]

pub struct VulkanTexture(pub(crate) Arc<ImageInner>);

/// An image the renderer can draw into, from [`smithay::backend::renderer::Offscreen`].

#[derive(Debug)]

pub struct VulkanImage(pub(crate) Arc<ImageInner>);

impl VulkanImage {

    /// The image's contents as a texture, for drawing it elsewhere.

    pub fn texture(&self) -> VulkanTexture {

        VulkanTexture(self.0.clone())

    }

}

/// A bound render target, borrowing whatever it was bound from.

#[derive(Debug)]

pub struct VulkanTarget<'a>(

    pub(crate) Arc<ImageInner>,

    pub(crate) PhantomData<&'a mut ()>,

);

/// Pixels read back from the GPU.

#[derive(Debug)]

pub struct VulkanMapping {

    pub(crate) data: Vec<u8>,

    pub(crate) width: u32,

    pub(crate) height: u32,

    pub(crate) fourcc: Fourcc,

}

macro_rules! texture_impl {

    ($t:ty, |$s:ident| $inner:expr) => {

        impl Texture for $t {

            fn width(&self) -> u32 {

                let $s = self;

                $inner.width

            }

            fn height(&self) -> u32 {

                let $s = self;

                $inner.height

            }

            fn format(&self) -> Option<Fourcc> {

                let $s = self;

                Some($inner.fourcc)

            }

        }

    };

}

texture_impl!(VulkanTexture, |s| s.0);

texture_impl!(VulkanImage, |s| s.0);

texture_impl!(VulkanTarget<'_>, |s| s.0);

texture_impl!(VulkanMapping, |s| s);

impl TextureMapping for VulkanMapping {

    fn flipped(&self) -> bool {

        false

    }

}