use crate::prelude::*;

use anyhow::Result;

use vma_rs::prelude::*;

use std::marker::PhantomData;
use std::sync::Arc;

#[derive(Debug, Clone, PartialEq, Hash, Eq)]
pub enum MemoryUsage {
    GpuOnly,
    CpuOnly,
    CpuToGpu,
    GpuToCpu,
}

impl MemoryUsage {
    pub fn into_vma(usage: Option<Self>) -> VmaMemoryUsage {
        match usage {
            Some(usage) => usage.into(),
            None => VMA_MEMORY_USAGE_UNKNOWN,
        }
    }
}

impl Into<VmaMemoryUsage> for MemoryUsage {
    fn into(self) -> VmaMemoryUsage {
        match self {
            Self::GpuOnly => VMA_MEMORY_USAGE_GPU_ONLY,
            Self::CpuOnly => VMA_MEMORY_USAGE_CPU_ONLY,
            Self::CpuToGpu => VMA_MEMORY_USAGE_CPU_TO_GPU,
            Self::GpuToCpu => VMA_MEMORY_USAGE_GPU_TO_CPU,
        }
    }
}

#[derive(Debug)]
pub struct Memory<T> {
    device: Arc<Device>,

    allocation: Allocation,

    data_type: PhantomData<T>,
}

impl<T> Memory<T> {
    pub(crate) fn forced_requirements(
        device: &Arc<Device>,
        memory_requirements: VkMemoryRequirements,
        buffer: VkBuffer,
        memory_usage: VmaMemoryUsage,
    ) -> Result<Arc<Memory<T>>> {
        let mut memory = Self::create_and_bind(device, memory_requirements, memory_usage, ())?;

        if let Some(mut_memory) = Arc::get_mut(&mut memory) {
            mut_memory.allocation.bind_buffer_memory(buffer)?;
        }

        Ok(memory)
    }

    pub(crate) fn buffer_memory(
        device: &Arc<Device>,
        buffer: VkBuffer,
        memory_usage: VmaMemoryUsage,
    ) -> Result<Arc<Memory<T>>> {
        let memory_requirements = device.buffer_memory_requirements(buffer);

        Self::create_and_bind(device, memory_requirements, memory_usage, buffer)
    }

    pub(crate) fn image_memory(
        device: &Arc<Device>,
        image: VkImage,
        memory_usage: VmaMemoryUsage,
    ) -> Result<Arc<Memory<T>>> {
        let memory_requirements = device.image_memory_requirements(image);

        Self::create_and_bind(device, memory_requirements, memory_usage, image)
    }

    pub(crate) fn vk_handle(&self) -> VkDeviceMemory {
        self.allocation.device_memory()
    }
}

trait MemoryBinder<T, K> {
    fn create_and_bind(
        device: &Arc<Device>,
        memory_requirements: VkMemoryRequirements,
        memory_usage: VmaMemoryUsage,
        argument: T,
    ) -> Result<Arc<Memory<K>>>;
}

impl<K> MemoryBinder<(), K> for Memory<K> {
    fn create_and_bind(
        device: &Arc<Device>,
        memory_requirements: VkMemoryRequirements,
        memory_usage: VmaMemoryUsage,
        _: (),
    ) -> Result<Arc<Memory<K>>> {
        let allocation = device
            .allocator()
            .allocate()
            .set_usage(memory_usage)
            .set_memory_type_bits(memory_requirements.memoryTypeBits.into())
            .build(&memory_requirements)?;

        Ok(Arc::new(Memory {
            device: device.clone(),

            allocation,

            data_type: PhantomData,
        }))
    }
}

impl<K> MemoryBinder<VkImage, K> for Memory<K> {
    fn create_and_bind(
        device: &Arc<Device>,
        memory_requirements: VkMemoryRequirements,
        memory_usage: VmaMemoryUsage,
        image: VkImage,
    ) -> Result<Arc<Memory<K>>> {
        let allocation = device
            .allocator()
            .allocate()
            .set_usage(memory_usage)
            .set_memory_type_bits(memory_requirements.memoryTypeBits.into())
            .build(image)?;

        Ok(Arc::new(Memory {
            device: device.clone(),

            allocation,

            data_type: PhantomData,
        }))
    }
}

impl<K> MemoryBinder<VkBuffer, K> for Memory<K> {
    fn create_and_bind(
        device: &Arc<Device>,
        memory_requirements: VkMemoryRequirements,
        memory_usage: VmaMemoryUsage,
        buffer: VkBuffer,
    ) -> Result<Arc<Memory<K>>> {
        let allocation = device
            .allocator()
            .allocate()
            .set_usage(memory_usage)
            .set_memory_type_bits(memory_requirements.memoryTypeBits.into())
            .build(buffer)?;

        Ok(Arc::new(Memory {
            device: device.clone(),

            allocation,

            data_type: PhantomData,
        }))
    }
}

impl<T> VulkanDevice for Memory<T> {
    fn device(&self) -> &Arc<Device> {
        &self.device
    }
}

impl<T: Clone> Memory<T> {
    pub fn map(&self, length: VkDeviceSize) -> Result<VkMappedMemory<'_, T>> {
        self.allocation.map(length)
    }
}