use std::alloc::Layout;

pub trait GpuAllocator: Send + Sync {
    /// Allocate `size` bytes with given alignment.
    /// Returns a host-mapped pointer.
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);
}

/// CPU fallback — standard system allocator.
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 backend using CudaAllocator
#[cfg(feature = "cuda")]
pub mod cuda {
    use super::GpuAllocator;
    use std::alloc::Layout;

    pub struct CudaAllocator;

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

    impl GpuAllocator for CudaAllocator {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            use cust::memory::cuda_malloc_unified;
            use cust_raw::driver_sys::*;

            let size = layout.size().max(layout.align());
            if size == 0 {
                return layout.align() as *mut u8;
            }

            let mut ctx: CUcontext = std::ptr::null_mut();
            cuCtxGetCurrent(&mut ctx);
            if ctx.is_null() {
                let mut primary: CUcontext = std::ptr::null_mut();
                cuDevicePrimaryCtxRetain(&mut primary, 0);
                cuCtxSetCurrent(primary);
            }

            let mut unified_ptr =
                unsafe { cuda_malloc_unified::<u8>(size).expect("cuda_malloc_unified failed") };
            let raw = unified_ptr.as_raw_mut();
            std::mem::forget(unified_ptr);
            raw
        }

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

/// Vulkan backend (gpu-allocator for now, there might be a better solution)
#[cfg(feature = "vulkan")]
pub mod vulkan {
    use super::GpuAllocator;
    use ash::vk;
    use gpu_allocator::MemoryLocation;
    use gpu_allocator::vulkan::{
        Allocation, AllocationCreateDesc, AllocationScheme, Allocator, AllocatorCreateDesc,
    };
    use parking_lot::Mutex;
    use std::alloc::Layout;
    use std::collections::HashMap;
    use std::sync::OnceLock;

    // So, having a static allocator seems like a terrible idea
    // But I cant find a way to get a functioning generic Arena constructor
    // That might not even be a necessity, since rust-gpu/rust-cuda might actually handle that
    // differently
    static VK_ALLOCATOR: OnceLock<VulkanAllocatorInner> = OnceLock::new();

    struct VulkanAllocatorInner {
        state: Mutex<VulkanState>,
    }

    struct VulkanState {
        allocator: Allocator,
        allocations: HashMap<usize, Allocation>,
    }

    pub fn init_vulkan(
        instance: &ash::Instance,
        device: &ash::Device,
        physical_device: vk::PhysicalDevice,
    ) {
        VK_ALLOCATOR.get_or_init(|| {
            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");

            VulkanAllocatorInner {
                state: Mutex::new(VulkanState {
                    allocator,
                    allocations: HashMap::new(),
                }),
            }
        });
    }

    fn inner() -> &'static VulkanAllocatorInner {
        VK_ALLOCATOR
            .get()
            .expect("Vulkan not initialized — call init_vulkan() before Arena::default()")
    }

    impl Default for VulkanAllocator {
        fn default() -> Self {
            let _ = inner();
            Self
        }
    }

    pub struct VulkanAllocator;

    // impl VulkanAllocator {
    //     pub fn new(allocator: Allocator) -> Self {
    //         Self {
    //             allocator: Mutex::new(allocator),
    //             allocations: Mutex::new(HashMap::new()),
    //         }
    //     }
    // }

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

            let inner = inner();
            let mut state = inner.state.lock();
            let allocation = state
                .allocator
                .allocate(&AllocationCreateDesc {
                    name: "arena",
                    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;

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

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