use std::alloc::Layout;

pub trait GpuAllocator: Send + Sync + Clone {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);
}

/// CPU fallback
#[derive(Clone)]
pub struct SystemAllocator;

impl Default for SystemAllocator {
    fn default() -> Self {
        Self
    }
}

impl GpuAllocator for SystemAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if layout.size() == 0 {
            return layout.align() as *mut u8;
        }
        unsafe { std::alloc::alloc(layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if layout.size() > 0 {
            unsafe { std::alloc::dealloc(ptr, layout) };
        }
    }
}

/// CUDA unified memory. Still using CudaAllocator, might move over to
#[cfg(feature = "cuda")]
pub mod cuda {
    use super::GpuAllocator;
    use cust::memory::{cuda_free_unified, cuda_malloc_unified, UnifiedPointer};
    use std::alloc::Layout;

    #[derive(Clone)]
    pub struct CudaAllocator;

    impl Default for CudaAllocator {
        fn default() -> Self {
            Self
        }
    }

    impl GpuAllocator for CudaAllocator {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            let size = layout.size();
            if size == 0 {
                return layout.align() as *mut u8;
            }

            let ptr = cuda_malloc_unified::<u8>(size)
                .expect("cuda_malloc_unified failed — is a CUDA context current?");

            let raw = ptr.as_raw_mut();
            std::mem::forget(ptr); // Leak RAII wrapper; Arena owns the raw pointer
            raw
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            if layout.size() > 0 && !ptr.is_null() {
                let _ = cuda_free_unified(UnifiedPointer::wrap(ptr));
            }
        }
    }
}

#[cfg(feature = "vulkan")]
pub mod vulkan {
    use super::GpuAllocator;
    use ash::vk;
    use gpu_allocator::vulkan::{
        Allocation, AllocationCreateDesc, AllocationScheme, Allocator, AllocatorCreateDesc,
    };
    use gpu_allocator::MemoryLocation;
    use parking_lot::Mutex;
    use std::alloc::Layout;
    use std::collections::HashMap;
    use std::sync::Arc;

    #[derive(Clone)]
    pub struct VulkanAllocator {
        inner: Arc<Mutex<VulkanInner>>,
    }

    struct VulkanInner {
        device: ash::Device,
        allocator: Allocator,
        allocations: HashMap<usize, Allocation>,
    }

    impl VulkanAllocator {
        pub fn new(
            instance: &ash::Instance,
            device: ash::Device,
            physical_device: vk::PhysicalDevice,
        ) -> Self {
            let allocator = Allocator::new(&AllocatorCreateDesc {
                instance: instance.clone(),
                device: device.clone(),
                physical_device,
                debug_settings: Default::default(),
                buffer_device_address: false,
                allocation_sizes: Default::default(),
            })
            .expect("Failed to create Vulkan allocator");

            Self {
                inner: Arc::new(Mutex::new(VulkanInner {
                    device,
                    allocator,
                    allocations: HashMap::new(),
                })),
            }
        }
    }

    impl GpuAllocator for VulkanAllocator {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            let size = layout.size();
            if size == 0 {
                return layout.align() as *mut u8;
            }

            let mut inner = self.inner.lock();

            let allocation = inner
                .allocator
                .allocate(&AllocationCreateDesc {
                    name: "arena_chunk",
                    requirements: vk::MemoryRequirements {
                        size: size as u64,
                        alignment: layout.align() as u64,
                        memory_type_bits: u32::MAX,
                    },
                    location: MemoryLocation::CpuToGpu,
                    linear: true,
                    allocation_scheme: AllocationScheme::GpuAllocatorManaged,
                })
                .expect("Vulkan allocation failed");

            let ptr = allocation
                .mapped_ptr()
                .expect("Vulkan allocation not host-mapped")
                .as_ptr() as *mut u8;

            inner.allocations.insert(ptr as usize, allocation);
            ptr
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            if layout.size() == 0 || ptr.is_null() {
                return;
            }
            let mut inner = self.inner.lock();
            if let Some(allocation) = inner.allocations.remove(&(ptr as usize)) {
                inner
                    .allocator
                    .free(allocation)
                    .expect("Vulkan free failed");
            }
        }
    }
}
