use erupt::vk;
use std::sync::Arc;

pub struct Allocator {
    pub allocator: vma::Allocator,
}

#[derive(Clone, Copy)] // FIXME: Fix this
pub struct AllocatedImage {
    pub image: vk::Image,
    pub allocation: vma::Allocation,
}

#[derive(Clone, Copy)]
pub struct AllocatedBuffer {
    pub buffer: vk::Buffer,
    pub allocation: vma::Allocation,
}

impl Allocator {
    #[inline]
    pub fn new(
        instance: Arc<erupt::InstanceLoader>,
        physical_device: vk::PhysicalDevice,
        device: Arc<erupt::DeviceLoader>,
    ) -> Self {
        let allocator_info = vma::AllocatorCreateInfo {
            physical_device,
            device,
            instance,
            flags: vma::AllocatorCreateFlags::empty(),
            frame_in_use_count: 2,
            preferred_large_heap_block_size: 0,
            heap_size_limits: None,
        };

        Self {
            allocator: vma::Allocator::new(&allocator_info).unwrap(),
        }
    }

    #[inline]
    pub fn create_buffer(
        &self,
        size: u64,
        usage: vk::BufferUsageFlags,
        memory_usage: vma::MemoryUsage,
    ) -> AllocatedBuffer {
        let buffer_info = vk::BufferCreateInfoBuilder::new().size(size).usage(usage);

        let allocation_info = vma::AllocationCreateInfo {
            usage: memory_usage,
            ..Default::default()
        };

        let allocated_buffer = self
            .allocator
            .create_buffer(&buffer_info, &allocation_info)
            .unwrap();

        AllocatedBuffer {
            buffer: allocated_buffer.0,
            allocation: allocated_buffer.1,
        }
    }

    #[inline]
    pub fn create_image(
        &self,
        image_info: vk::ImageCreateInfo,
        memory_usage: vma::MemoryUsage,
    ) -> AllocatedImage {
        let allocation_info = vma::AllocationCreateInfo {
            usage: memory_usage,
            ..Default::default()
        };

        let allocated_image = self
            .allocator
            .create_image(&image_info, &allocation_info)
            .unwrap();

        AllocatedImage {
            image: allocated_image.0,
            allocation: allocated_image.1,
        }
    }
}
