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

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

use crate::Device;

pub struct Buffer {
    device: Device,
    buffer: vk::Buffer,
    allocation: vk_mem::Allocation,
    usage: vk::BufferUsageFlags,
    size: u64,
}

impl Drop for Buffer {
    fn drop(&mut self) {
        unsafe {
            self.device
                .alloc()
                .destroy_buffer(self.buffer, &mut self.allocation);
        }
    }
}

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,
    ) -> 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()
                },
            )?
        };

        let buffer = Self {
            device,
            buffer,
            allocation,
            usage,
            size: size as u64,
        };

        Ok(Arc::new(buffer))
    }

    pub fn map(&mut self) -> VkResult<MappedBuffer<'_>> {
        let bytes = unsafe {
            let data = self.device.alloc().map_memory(&mut self.allocation)?;
            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.buffer
    }
}

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

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

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
    }
}