shdrlib 0.1.2

//! Buffer creation helpers for EX tier
//!
//! This module provides ergonomic functions for creating common buffer types
//! with sensible defaults while exposing all configuration options.

use crate::core;
use crate::ex::RuntimeError;
use ash::vk;
use std::sync::Arc;

/// Buffer usage patterns
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BufferUsage {
    /// Vertex buffer (GPU read, CPU write once)
    Vertex,
    /// Index buffer (GPU read, CPU write once)
    Index,
    /// Uniform buffer (GPU read, CPU write frequently)
    Uniform,
    /// Storage buffer (GPU read/write, CPU read)
    Storage,
    /// Staging buffer (CPU write, GPU read once)
    Staging,
    /// Storage buffer (host-visible, for CPU-GPU data transfer)
    StorageHostVisible,
}

impl BufferUsage {
    /// Get Vulkan buffer usage flags
    #[inline]
    pub fn flags(self) -> vk::BufferUsageFlags {
        match self {
            BufferUsage::Vertex => {
                vk::BufferUsageFlags::VERTEX_BUFFER | vk::BufferUsageFlags::TRANSFER_DST
            }
            BufferUsage::Index => {
                vk::BufferUsageFlags::INDEX_BUFFER | vk::BufferUsageFlags::TRANSFER_DST
            }
            BufferUsage::Uniform => vk::BufferUsageFlags::UNIFORM_BUFFER,
            BufferUsage::Storage => vk::BufferUsageFlags::STORAGE_BUFFER,
            BufferUsage::Staging => vk::BufferUsageFlags::TRANSFER_SRC,
            BufferUsage::StorageHostVisible => vk::BufferUsageFlags::STORAGE_BUFFER,
        }
    }

    /// Get Vulkan memory property flags
    #[inline]
    pub fn memory_properties(self) -> vk::MemoryPropertyFlags {
        match self {
            BufferUsage::Vertex | BufferUsage::Index => vk::MemoryPropertyFlags::DEVICE_LOCAL,
            BufferUsage::Uniform | BufferUsage::Staging => {
                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT
            }
            BufferUsage::Storage => vk::MemoryPropertyFlags::DEVICE_LOCAL,
            BufferUsage::StorageHostVisible => {
                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT
            }
        }
    }
}

/// Builder for creating buffers with custom configuration
pub struct BufferBuilder {
    size: vk::DeviceSize,
    usage: vk::BufferUsageFlags,
    memory_properties: vk::MemoryPropertyFlags,
    sharing_mode: vk::SharingMode,
}

impl BufferBuilder {
    /// Create a new buffer builder with required size
    pub fn new(size: vk::DeviceSize) -> Self {
        Self {
            size,
            usage: vk::BufferUsageFlags::empty(),
            memory_properties: vk::MemoryPropertyFlags::empty(),
            sharing_mode: vk::SharingMode::EXCLUSIVE,
        }
    }

    /// Set buffer usage flags
    pub fn usage(mut self, usage: vk::BufferUsageFlags) -> Self {
        self.usage = usage;
        self
    }

    /// Set memory property flags
    pub fn memory_properties(mut self, properties: vk::MemoryPropertyFlags) -> Self {
        self.memory_properties = properties;
        self
    }

    /// Set sharing mode
    pub fn sharing_mode(mut self, mode: vk::SharingMode) -> Self {
        self.sharing_mode = mode;
        self
    }

    /// Build the buffer
    ///
    /// # Errors
    ///
    /// Returns error if buffer or memory allocation fails.
    pub fn build(self, device: &Arc<core::Device>) -> Result<core::Buffer, RuntimeError> {
        let buffer_info = vk::BufferCreateInfo::default()
            .size(self.size)
            .usage(self.usage)
            .sharing_mode(self.sharing_mode);

        let buffer = unsafe { device.handle().create_buffer(&buffer_info, None) }
            .map_err(|e| RuntimeError::Other(format!("Buffer creation failed: {:?}", e)))?;

        // Allocate memory
        let mem_reqs = unsafe { device.handle().get_buffer_memory_requirements(buffer) };
        let mem_type_index = device
            .find_memory_type(mem_reqs.memory_type_bits, self.memory_properties)
            .ok_or_else(|| RuntimeError::Other("No suitable memory type found".to_string()))?;

        let alloc_info = vk::MemoryAllocateInfo::default()
            .allocation_size(mem_reqs.size)
            .memory_type_index(mem_type_index);

        let memory = unsafe { device.handle().allocate_memory(&alloc_info, None) }
            .map_err(|e| RuntimeError::Other(format!("Memory allocation failed: {:?}", e)))?;

        // Bind memory
        unsafe { device.handle().bind_buffer_memory(buffer, memory, 0) }
            .map_err(|e| RuntimeError::Other(format!("Memory bind failed: {:?}", e)))?;

        Ok(core::Buffer::from_raw(buffer, memory, self.size))
    }
}

/// Create a vertex buffer
///
/// Creates a device-local buffer suitable for vertex data.
///
/// # Example
///
/// ```rust,no_run
/// # use shdrlib::ex::helpers::*;
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # let runtime = shdrlib::ex::RuntimeManager::new(Default::default())?;
/// let vertices: Vec<f32> = vec![0.0, 0.0, 1.0, 1.0];
/// let vertex_buffer = create_vertex_buffer(&runtime.device(), &vertices)?;
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn create_vertex_buffer<T: Copy>(
    device: &Arc<core::Device>,
    data: &[T],
) -> Result<core::Buffer, RuntimeError> {
    let size = std::mem::size_of_val(data) as vk::DeviceSize;

    BufferBuilder::new(size)
        .usage(BufferUsage::Vertex.flags())
        .memory_properties(BufferUsage::Vertex.memory_properties())
        .build(device)
}

/// Create an index buffer
///
/// Creates a device-local buffer suitable for index data.
#[inline]
pub fn create_index_buffer<T: Copy>(
    device: &Arc<core::Device>,
    data: &[T],
) -> Result<core::Buffer, RuntimeError> {
    let size = std::mem::size_of_val(data) as vk::DeviceSize;

    BufferBuilder::new(size)
        .usage(BufferUsage::Index.flags())
        .memory_properties(BufferUsage::Index.memory_properties())
        .build(device)
}

/// Create a uniform buffer
///
/// Creates a host-visible buffer suitable for uniform data (updated frequently).
///
/// # Example
///
/// ```rust,no_run
/// # use shdrlib::ex::helpers::*;
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # let runtime = shdrlib::ex::RuntimeManager::new(Default::default())?;
/// #[repr(C)]
/// struct UniformData {
///     mvp: [[f32; 4]; 4],
/// }
///
/// let uniform_buffer = create_uniform_buffer::<UniformData>(&runtime.device())?;
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn create_uniform_buffer<T>(device: &Arc<core::Device>) -> Result<core::Buffer, RuntimeError> {
    let size = std::mem::size_of::<T>() as vk::DeviceSize;

    BufferBuilder::new(size)
        .usage(BufferUsage::Uniform.flags())
        .memory_properties(BufferUsage::Uniform.memory_properties())
        .build(device)
}

/// Create a storage buffer
///
/// Creates a device-local buffer suitable for compute shader storage.
#[inline]
pub fn create_storage_buffer(
    device: &Arc<core::Device>,
    size: vk::DeviceSize,
) -> Result<core::Buffer, RuntimeError> {
    BufferBuilder::new(size)
        .usage(BufferUsage::StorageHostVisible.flags())
        .memory_properties(BufferUsage::StorageHostVisible.memory_properties())
        .build(device)
}

/// Create a staging buffer
///
/// Creates a host-visible buffer for transferring data to GPU.
pub fn create_staging_buffer<T: Copy>(
    device: &Arc<core::Device>,
    data: &[T],
) -> Result<core::Buffer, RuntimeError> {
    let size = std::mem::size_of_val(data) as vk::DeviceSize;

    BufferBuilder::new(size)
        .usage(BufferUsage::Staging.flags())
        .memory_properties(BufferUsage::Staging.memory_properties())
        .build(device)
}

/// Map buffer memory and write data
///
/// # Safety
///
/// The buffer must have been created with HOST_VISIBLE memory.
///
/// # Errors
///
/// Returns error if memory mapping fails.
pub fn write_buffer<T: Copy>(
    device: &Arc<core::Device>,
    buffer: &core::Buffer,
    data: &[T],
) -> Result<(), RuntimeError> {
    let size = std::mem::size_of_val(data) as vk::DeviceSize;

    unsafe {
        let ptr = device
            .handle()
            .map_memory(buffer.memory(), 0, size, vk::MemoryMapFlags::empty())
            .map_err(|e| RuntimeError::Other(format!("Memory map failed: {:?}", e)))?;

        std::ptr::copy_nonoverlapping(data.as_ptr() as *const u8, ptr as *mut u8, size as usize);

        device.handle().unmap_memory(buffer.memory());
    }

    Ok(())
}

/// Copy data from one buffer to another using a command buffer
///
/// # Errors
///
/// Returns error if command buffer operations fail.
pub fn copy_buffer(
    device: &Arc<core::Device>,
    command_buffer: &core::CommandBuffer,
    src: &core::Buffer,
    dst: &core::Buffer,
    size: vk::DeviceSize,
) -> Result<(), RuntimeError> {
    let region = vk::BufferCopy::default()
        .src_offset(0)
        .dst_offset(0)
        .size(size);

    unsafe {
        device.handle().cmd_copy_buffer(
            command_buffer.handle(),
            src.handle(),
            dst.handle(),
            &[region],
        );
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ex::{RuntimeConfig, RuntimeManager};

    fn test_runtime() -> RuntimeManager {
        RuntimeManager::new(RuntimeConfig {
            enable_validation: false,
            ..Default::default()
        })
        .unwrap()
    }

    #[test]
    fn test_buffer_usage_flags() {
        assert!(BufferUsage::Vertex
            .flags()
            .contains(vk::BufferUsageFlags::VERTEX_BUFFER));
        assert!(BufferUsage::Index
            .flags()
            .contains(vk::BufferUsageFlags::INDEX_BUFFER));
        assert!(BufferUsage::Uniform
            .flags()
            .contains(vk::BufferUsageFlags::UNIFORM_BUFFER));
        assert!(BufferUsage::Storage
            .flags()
            .contains(vk::BufferUsageFlags::STORAGE_BUFFER));
    }

    #[test]
    fn test_create_vertex_buffer() {
        let runtime = test_runtime();
        let vertices: Vec<f32> = vec![0.0, 0.5, -0.5, -0.5, 0.5, -0.5];
        let result = create_vertex_buffer(&runtime.device(), &vertices);
        assert!(result.is_ok());
    }

    #[test]
    fn test_create_uniform_buffer() {
        let runtime = test_runtime();

        #[repr(C)]
        struct UniformData {
            mvp: [[f32; 4]; 4],
        }

        let result = create_uniform_buffer::<UniformData>(&runtime.device());
        assert!(result.is_ok());
    }

    #[test]
    fn test_create_storage_buffer() {
        let runtime = test_runtime();
        let result = create_storage_buffer(&runtime.device(), 1024);
        assert!(result.is_ok());
    }

    #[test]
    fn test_buffer_builder() {
        let runtime = test_runtime();
        let buffer = BufferBuilder::new(256)
            .usage(vk::BufferUsageFlags::UNIFORM_BUFFER)
            .memory_properties(
                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
            )
            .build(&runtime.device());

        assert!(buffer.is_ok());
        let buffer = buffer.unwrap();
        assert_eq!(buffer.size(), 256);
    }
}