use std::{
    ops::{Deref, DerefMut},
    sync::Arc,
};

use ash::{prelude::VkResult, vk};
use vk_mem::Alloc;

use crate::{define_device_owned_handle, device::DeviceOwned, Device};

define_device_owned_handle! {
    #[derive(Debug)]
    pub Buffer(vk::Buffer) {
        alloc: vk_mem::Allocation,
        usage: vk::BufferUsageFlags,
        size: u64,
    } => |this|  unsafe {
        this.device().clone().alloc().destroy_buffer(this.handle(), &mut this.alloc);
    }
}

impl Buffer {
    pub fn new(
        device: Device,
        size: usize,
        usage: vk::BufferUsageFlags,
        queue_families: &[u32],
        memory_usage: vk_mem::MemoryUsage,
        alloc_flags: vk_mem::AllocationCreateFlags,
        name: Option<std::borrow::Cow<'static, str>>,
    ) -> VkResult<Arc<Self>> {
        let sharing_mode = if queue_families.len() > 1 {
            vk::SharingMode::CONCURRENT
        } else {
            vk::SharingMode::EXCLUSIVE
        };

        let (buffer, allocation) = unsafe {
            device.alloc().create_buffer(
                &vk::BufferCreateInfo::default()
                    .size(size as u64)
                    .usage(usage)
                    .queue_family_indices(queue_families)
                    .sharing_mode(sharing_mode),
                &vk_mem::AllocationCreateInfo {
                    flags: alloc_flags,
                    usage: memory_usage,
                    ..Default::default()
                },
            )?
        };

        Ok(Arc::new(Self::construct(
            device,
            buffer,
            name,
            allocation,
            usage,
            size as u64,
        )?))
    }

    pub fn map_arc(self: &mut Arc<Self>) -> VkResult<MappedBuffer<'_>> {
        Arc::get_mut(self).map(Self::map).unwrap()
    }

    pub fn map(&mut self) -> VkResult<MappedBuffer<'_>> {
        let bytes = unsafe {
            let data = self.inner.dev().alloc().map_memory(&mut self.alloc)?;
            let slice = core::slice::from_raw_parts_mut(data, self.size as usize);

            slice
        };

        Ok(MappedBuffer { inner: self, bytes })
    }
    pub fn buffer(&self) -> vk::Buffer {
        self.handle()
    }
    pub fn len(&self) -> u64 {
        self.size
    }
}

pub struct MappedBuffer<'a> {
    bytes: &'a mut [u8],
    inner: &'a mut Buffer,
}

impl Drop for MappedBuffer<'_> {
    fn drop(&mut self) {
        unsafe {
            self.inner
                .inner
                .dev()
                .alloc()
                .unmap_memory(&mut self.inner.alloc);
        }
    }
}

impl Deref for MappedBuffer<'_> {
    type Target = [u8];

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

impl DerefMut for MappedBuffer<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.bytes
    }
}