shdrlib 0.1.5

//! Image creation helpers for EX tier
//!
//! This module provides ergonomic functions for creating images, textures, and render targets
//! with sensible defaults while exposing all configuration options.

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

/// Image usage patterns
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImageUsage {
    /// Render target (color attachment)
    RenderTarget,
    /// Depth/stencil attachment
    DepthStencil,
    /// Texture (sampled in shaders)
    Texture,
    /// Storage image (read/write in compute)
    Storage,
}

impl ImageUsage {
    /// Get Vulkan image usage flags
    #[inline]
    pub fn flags(self) -> vk::ImageUsageFlags {
        match self {
            ImageUsage::RenderTarget => {
                vk::ImageUsageFlags::COLOR_ATTACHMENT
                    | vk::ImageUsageFlags::SAMPLED
                    | vk::ImageUsageFlags::TRANSFER_DST
            }
            ImageUsage::DepthStencil => {
                vk::ImageUsageFlags::DEPTH_STENCIL_ATTACHMENT | vk::ImageUsageFlags::SAMPLED
            }
            ImageUsage::Texture => {
                vk::ImageUsageFlags::SAMPLED
                    | vk::ImageUsageFlags::TRANSFER_DST
                    | vk::ImageUsageFlags::TRANSFER_SRC
            }
            ImageUsage::Storage => vk::ImageUsageFlags::STORAGE | vk::ImageUsageFlags::TRANSFER_DST,
        }
    }

    /// Get initial image layout
    #[inline]
    pub fn initial_layout(self) -> vk::ImageLayout {
        match self {
            ImageUsage::RenderTarget => vk::ImageLayout::UNDEFINED,
            ImageUsage::DepthStencil => vk::ImageLayout::UNDEFINED,
            ImageUsage::Texture => vk::ImageLayout::UNDEFINED,
            ImageUsage::Storage => vk::ImageLayout::GENERAL,
        }
    }

    /// Get optimal final layout
    #[inline]
    pub fn optimal_layout(self) -> vk::ImageLayout {
        match self {
            ImageUsage::RenderTarget => vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL,
            ImageUsage::DepthStencil => vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
            ImageUsage::Texture => vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
            ImageUsage::Storage => vk::ImageLayout::GENERAL,
        }
    }
}

/// Builder for creating images with custom configuration
pub struct ImageBuilder {
    width: u32,
    height: u32,
    format: vk::Format,
    usage: vk::ImageUsageFlags,
    tiling: vk::ImageTiling,
    memory_properties: vk::MemoryPropertyFlags,
    mip_levels: u32,
    array_layers: u32,
    samples: vk::SampleCountFlags,
    image_type: vk::ImageType,
}

impl ImageBuilder {
    /// Create a new 2D image builder
    pub fn new_2d(width: u32, height: u32, format: vk::Format) -> Self {
        Self {
            width,
            height,
            format,
            usage: vk::ImageUsageFlags::empty(),
            tiling: vk::ImageTiling::OPTIMAL,
            memory_properties: vk::MemoryPropertyFlags::DEVICE_LOCAL,
            mip_levels: 1,
            array_layers: 1,
            samples: vk::SampleCountFlags::TYPE_1,
            image_type: vk::ImageType::TYPE_2D,
        }
    }

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

    /// Set image tiling
    pub fn tiling(mut self, tiling: vk::ImageTiling) -> Self {
        self.tiling = tiling;
        self
    }

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

    /// Set mip levels
    pub fn mip_levels(mut self, levels: u32) -> Self {
        self.mip_levels = levels;
        self
    }

    /// Set array layers
    pub fn array_layers(mut self, layers: u32) -> Self {
        self.array_layers = layers;
        self
    }

    /// Set sample count
    pub fn samples(mut self, samples: vk::SampleCountFlags) -> Self {
        self.samples = samples;
        self
    }

    /// Build the image
    ///
    /// # Errors
    ///
    /// Returns error if image or memory allocation fails.
    pub fn build(self, device: &Arc<core::Device>) -> Result<core::Image, RuntimeError> {
        let image_info = vk::ImageCreateInfo::default()
            .image_type(self.image_type)
            .format(self.format)
            .extent(vk::Extent3D {
                width: self.width,
                height: self.height,
                depth: 1,
            })
            .mip_levels(self.mip_levels)
            .array_layers(self.array_layers)
            .samples(self.samples)
            .tiling(self.tiling)
            .usage(self.usage)
            .sharing_mode(vk::SharingMode::EXCLUSIVE)
            .initial_layout(vk::ImageLayout::UNDEFINED);

        let image = unsafe { device.handle().create_image(&image_info, None) }
            .map_err(|e| RuntimeError::Other(format!("Image creation failed: {:?}", e)))?;

        // Allocate memory
        let mem_reqs = unsafe { device.handle().get_image_memory_requirements(image) };
        let mem_type_index = device
            .find_memory_type(mem_reqs.memory_type_bits, self.memory_properties)
            .ok_or_else(|| {
                // Cleanup image on error
                unsafe { device.handle().destroy_image(image, None) };
                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| {
                // Cleanup image on error
                unsafe { device.handle().destroy_image(image, None) };
                RuntimeError::Other(format!("Memory allocation failed: {:?}", e))
            })?;

        // Bind memory
        unsafe { device.handle().bind_image_memory(image, memory, 0) }
            .map_err(|e| {
                // Cleanup image and memory on error
                unsafe {
                    device.handle().free_memory(memory, None);
                    device.handle().destroy_image(image, None);
                }
                RuntimeError::Other(format!("Memory bind failed: {:?}", e))
            })?;

        // Create image view
        let view_info = vk::ImageViewCreateInfo::default()
            .image(image)
            .view_type(vk::ImageViewType::TYPE_2D)
            .format(self.format)
            .subresource_range(vk::ImageSubresourceRange {
                aspect_mask: if is_depth_format(self.format) {
                    vk::ImageAspectFlags::DEPTH
                } else {
                    vk::ImageAspectFlags::COLOR
                },
                base_mip_level: 0,
                level_count: self.mip_levels,
                base_array_layer: 0,
                layer_count: self.array_layers,
            });

        let view = unsafe { device.handle().create_image_view(&view_info, None) }
            .map_err(|e| {
                // Cleanup image and memory on error
                unsafe {
                    device.handle().free_memory(memory, None);
                    device.handle().destroy_image(image, None);
                }
                RuntimeError::Other(format!("Image view creation failed: {:?}", e))
            })?;

        Ok(core::Image::from_raw(
            image,
            memory,
            view,
            self.format,
            vk::Extent3D {
                width: self.width,
                height: self.height,
                depth: 1,
            },
        ))
    }
}

/// Create a render target (color attachment)
///
/// Creates a device-local image suitable for rendering.
///
/// # 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 render_target = create_render_target(
///     &runtime.device(),
///     800,
///     600,
///     ash::vk::Format::R8G8B8A8_UNORM,
/// )?;
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn create_render_target(
    device: &Arc<core::Device>,
    width: u32,
    height: u32,
    format: vk::Format,
) -> Result<crate::ex::ExImage, RuntimeError> {
    let core_image = ImageBuilder::new_2d(width, height, format)
        .usage(ImageUsage::RenderTarget.flags())
        .build(device)?;
    Ok(crate::ex::ExImage::new(Arc::clone(device), core_image))
}

/// Create a depth/stencil attachment
///
/// Creates a device-local depth buffer with automatic cleanup.
#[inline]
pub fn create_depth_stencil(
    device: &Arc<core::Device>,
    width: u32,
    height: u32,
    format: vk::Format,
) -> Result<crate::ex::ExImage, RuntimeError> {
    let core_image = ImageBuilder::new_2d(width, height, format)
        .usage(ImageUsage::DepthStencil.flags())
        .build(device)?;
    Ok(crate::ex::ExImage::new(Arc::clone(device), core_image))
}

/// Create a texture (sampled image)
///
/// Creates a device-local texture for sampling in shaders with automatic cleanup.
#[inline]
pub fn create_texture(
    device: &Arc<core::Device>,
    width: u32,
    height: u32,
    format: vk::Format,
) -> Result<crate::ex::ExImage, RuntimeError> {
    let core_image = ImageBuilder::new_2d(width, height, format)
        .usage(ImageUsage::Texture.flags())
        .build(device)?;
    Ok(crate::ex::ExImage::new(Arc::clone(device), core_image))
}

/// Create a storage image (for compute shaders)
///
/// Creates a device-local storage image for read/write in compute with automatic cleanup.
#[inline]
pub fn create_storage_image(
    device: &Arc<core::Device>,
    width: u32,
    height: u32,
    format: vk::Format,
) -> Result<crate::ex::ExImage, RuntimeError> {
    let core_image = ImageBuilder::new_2d(width, height, format)
        .usage(ImageUsage::Storage.flags())
        .build(device)?;
    Ok(crate::ex::ExImage::new(Arc::clone(device), core_image))
}

/// Transition image layout using a command buffer
///
/// # Errors
///
/// Returns error if command buffer operations fail.
pub fn transition_image_layout(
    device: &Arc<core::Device>,
    command_buffer: &core::CommandBuffer,
    image: vk::Image,
    old_layout: vk::ImageLayout,
    new_layout: vk::ImageLayout,
    aspect_mask: vk::ImageAspectFlags,
) -> Result<(), RuntimeError> {
    let barrier = vk::ImageMemoryBarrier::default()
        .old_layout(old_layout)
        .new_layout(new_layout)
        .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
        .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
        .image(image)
        .subresource_range(vk::ImageSubresourceRange {
            aspect_mask,
            base_mip_level: 0,
            level_count: 1,
            base_array_layer: 0,
            layer_count: 1,
        })
        .src_access_mask(get_access_mask(old_layout))
        .dst_access_mask(get_access_mask(new_layout));

    let (src_stage, dst_stage) = get_pipeline_stages(old_layout, new_layout);

    unsafe {
        device.handle().cmd_pipeline_barrier(
            command_buffer.handle(),
            src_stage,
            dst_stage,
            vk::DependencyFlags::empty(),
            &[],
            &[],
            &[barrier],
        );
    }

    Ok(())
}

/// Helper to determine if format is a depth format
#[inline]
fn is_depth_format(format: vk::Format) -> bool {
    matches!(
        format,
        vk::Format::D16_UNORM
            | vk::Format::D32_SFLOAT
            | vk::Format::D16_UNORM_S8_UINT
            | vk::Format::D24_UNORM_S8_UINT
            | vk::Format::D32_SFLOAT_S8_UINT
    )
}

/// Get access mask for image layout
fn get_access_mask(layout: vk::ImageLayout) -> vk::AccessFlags {
    match layout {
        vk::ImageLayout::UNDEFINED => vk::AccessFlags::empty(),
        vk::ImageLayout::TRANSFER_DST_OPTIMAL => vk::AccessFlags::TRANSFER_WRITE,
        vk::ImageLayout::TRANSFER_SRC_OPTIMAL => vk::AccessFlags::TRANSFER_READ,
        vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL => vk::AccessFlags::SHADER_READ,
        vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL => {
            vk::AccessFlags::COLOR_ATTACHMENT_READ | vk::AccessFlags::COLOR_ATTACHMENT_WRITE
        }
        vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL => {
            vk::AccessFlags::DEPTH_STENCIL_ATTACHMENT_READ
                | vk::AccessFlags::DEPTH_STENCIL_ATTACHMENT_WRITE
        }
        vk::ImageLayout::GENERAL => vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE,
        _ => vk::AccessFlags::empty(),
    }
}

/// Get pipeline stages for layout transition
fn get_pipeline_stages(
    old_layout: vk::ImageLayout,
    new_layout: vk::ImageLayout,
) -> (vk::PipelineStageFlags, vk::PipelineStageFlags) {
    let src_stage = match old_layout {
        vk::ImageLayout::UNDEFINED => vk::PipelineStageFlags::TOP_OF_PIPE,
        vk::ImageLayout::TRANSFER_DST_OPTIMAL | vk::ImageLayout::TRANSFER_SRC_OPTIMAL => {
            vk::PipelineStageFlags::TRANSFER
        }
        vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL => vk::PipelineStageFlags::FRAGMENT_SHADER,
        vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL => {
            vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
        }
        vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL => {
            vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS
        }
        _ => vk::PipelineStageFlags::ALL_COMMANDS,
    };

    let dst_stage = match new_layout {
        vk::ImageLayout::TRANSFER_DST_OPTIMAL | vk::ImageLayout::TRANSFER_SRC_OPTIMAL => {
            vk::PipelineStageFlags::TRANSFER
        }
        vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL => vk::PipelineStageFlags::FRAGMENT_SHADER,
        vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL => {
            vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
        }
        vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL => {
            vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS
        }
        _ => vk::PipelineStageFlags::ALL_COMMANDS,
    };

    (src_stage, dst_stage)
}

#[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_image_usage_flags() {
        assert!(ImageUsage::RenderTarget
            .flags()
            .contains(vk::ImageUsageFlags::COLOR_ATTACHMENT));
        assert!(ImageUsage::DepthStencil
            .flags()
            .contains(vk::ImageUsageFlags::DEPTH_STENCIL_ATTACHMENT));
        assert!(ImageUsage::Texture
            .flags()
            .contains(vk::ImageUsageFlags::SAMPLED));
        assert!(ImageUsage::Storage
            .flags()
            .contains(vk::ImageUsageFlags::STORAGE));
    }

    #[test]
    fn test_create_render_target() {
        let runtime = test_runtime();
        let result = create_render_target(&runtime.device(), 800, 600, vk::Format::R8G8B8A8_UNORM);
        assert!(result.is_ok());
    }

    #[test]
    fn test_create_depth_stencil() {
        let runtime = test_runtime();
        let result = create_depth_stencil(&runtime.device(), 800, 600, vk::Format::D32_SFLOAT);
        assert!(result.is_ok());
    }

    #[test]
    fn test_create_texture() {
        let runtime = test_runtime();
        let result = create_texture(&runtime.device(), 256, 256, vk::Format::R8G8B8A8_UNORM);
        assert!(result.is_ok());
    }

    #[test]
    fn test_image_builder() {
        let runtime = test_runtime();
        let image = ImageBuilder::new_2d(512, 512, vk::Format::R8G8B8A8_UNORM)
            .usage(vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST)
            .mip_levels(4)
            .build(&runtime.device());

        assert!(image.is_ok());
    }

    #[test]
    fn test_depth_format_detection() {
        assert!(is_depth_format(vk::Format::D32_SFLOAT));
        assert!(is_depth_format(vk::Format::D16_UNORM));
        assert!(!is_depth_format(vk::Format::R8G8B8A8_UNORM));
    }
}