shdrlib 0.1.5

//! Pipeline module - Graphics and compute pipelines
//!
//! **WARNING**: CORE objects do NOT enforce lifetime dependencies.
//! Pipelines must not outlive the device.

use crate::core::Device;
use ash::vk;
use thiserror::Error;

/// Pipeline creation error
#[derive(Debug, Error)]
pub enum PipelineError {
    #[error("Pipeline creation failed: {0}")]
    CreationFailed(vk::Result),

    #[error("Pipeline layout creation failed: {0}")]
    LayoutCreationFailed(vk::Result),

    #[error("Missing required field: {0}")]
    MissingField(&'static str),

    #[error("Invalid pipeline configuration: {0}")]
    InvalidConfiguration(String),
}

/// Pipeline layout wrapper
///
/// Describes the interface between pipeline stages and resources.
pub struct PipelineLayout {
    layout: vk::PipelineLayout,
}

impl PipelineLayout {
    /// Create a new pipeline layout
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the layout on
    /// * `descriptor_set_layouts` - Descriptor set layouts
    /// * `push_constant_ranges` - Push constant ranges
    ///
    /// # Errors
    ///
    /// Returns `PipelineError::LayoutCreationFailed` if creation fails.
    pub fn new(
        device: &Device,
        descriptor_set_layouts: &[vk::DescriptorSetLayout],
        push_constant_ranges: &[vk::PushConstantRange],
    ) -> Result<Self, PipelineError> {
        let layout_info = vk::PipelineLayoutCreateInfo {
            set_layout_count: descriptor_set_layouts.len() as u32,
            p_set_layouts: descriptor_set_layouts.as_ptr(),
            push_constant_range_count: push_constant_ranges.len() as u32,
            p_push_constant_ranges: push_constant_ranges.as_ptr(),
            ..Default::default()
        };

        // SAFETY: device is valid, layout_info is properly initialized
        let layout = unsafe {
            device
                .handle()
                .create_pipeline_layout(&layout_info, None)
                .map_err(PipelineError::LayoutCreationFailed)?
        };

        Ok(Self { layout })
    }

    /// Get the raw Vulkan pipeline layout handle
    #[inline]
    pub fn handle(&self) -> vk::PipelineLayout {
        self.layout
    }

    /// Destroy the layout manually
    ///
    /// # Safety
    ///
    /// The device must be the same device used to create this layout.
    /// After calling destroy(), the layout handle is invalidated.
    pub fn destroy(&self, device: &Device) {
        // SAFETY: device and layout are valid
        unsafe {
            device.handle().destroy_pipeline_layout(self.layout, None);
        }
    }
}

impl Drop for PipelineLayout {
    fn drop(&mut self) {
        if self.layout != vk::PipelineLayout::null() {
            eprintln!(
                "WARNING: PipelineLayout dropped without calling .destroy() - potential memory leak"
            );
        }
    }
}

/// Graphics or compute pipeline wrapper
///
/// Owns its pipeline layout for correct drop order.
pub struct Pipeline {
    pipeline: vk::Pipeline,
    layout: PipelineLayout,
    bind_point: vk::PipelineBindPoint,
}

impl Pipeline {
    /// Create Pipeline from raw handle (for EX tier usage)
    ///
    /// # Safety
    ///
    /// The handle must be valid and the bind point must match the actual pipeline.
    pub(crate) fn from_handle(
        pipeline: vk::Pipeline,
        layout: PipelineLayout,
        bind_point: vk::PipelineBindPoint,
    ) -> Self {
        Self {
            pipeline,
            layout,
            bind_point,
        }
    }

    /// Get the raw Vulkan pipeline handle
    #[inline]
    pub fn handle(&self) -> vk::Pipeline {
        self.pipeline
    }

    /// Get the pipeline layout
    #[inline]
    pub fn layout(&self) -> &PipelineLayout {
        &self.layout
    }

    /// Get the pipeline bind point
    #[inline]
    pub fn bind_point(&self) -> vk::PipelineBindPoint {
        self.bind_point
    }

    /// Destroy the pipeline manually
    ///
    /// # Safety
    ///
    /// The device must be the same device used to create this pipeline.
    /// After calling destroy(), the pipeline handle is invalidated.
    pub fn destroy(&self, device: &Device) {
        // SAFETY: device and pipeline are valid
        unsafe {
            device.handle().destroy_pipeline(self.pipeline, None);
        }
        self.layout.destroy(device);
    }
}

impl Drop for Pipeline {
    fn drop(&mut self) {
        if self.pipeline != vk::Pipeline::null() {
            eprintln!(
                "WARNING: Pipeline dropped without calling .destroy() - potential memory leak"
            );
        }
    }
}

/// Shader stage information for pipeline creation
#[derive(Clone)]
pub struct ShaderStageInfo {
    pub stage: vk::ShaderStageFlags,
    pub module: vk::ShaderModule,
    pub entry_point: String,
}

/// Graphics pipeline builder
///
/// Provides a fluent API for configuring graphics pipelines with sensible defaults.
///
/// # Example
///
/// ```rust,no_run
/// use shdrlib::core::{PipelineBuilder, Device};
/// use ash::vk;
///
/// # fn example(
/// #     device: &Device,
/// #     vertex_module: vk::ShaderModule,
/// #     fragment_module: vk::ShaderModule,
/// #     descriptor_layout: vk::DescriptorSetLayout,
/// # ) -> Result<(), shdrlib::core::PipelineError> {
/// let pipeline = PipelineBuilder::new()
///     .vertex_shader(vertex_module, "main")
///     .fragment_shader(fragment_module, "main")
///     .viewport(vk::Viewport {
///         x: 0.0,
///         y: 0.0,
///         width: 800.0,
///         height: 600.0,
///         min_depth: 0.0,
///         max_depth: 1.0,
///     })
///     .scissor(vk::Rect2D {
///         offset: vk::Offset2D { x: 0, y: 0 },
///         extent: vk::Extent2D { width: 800, height: 600 },
///     })
///     .build_graphics(device, &[descriptor_layout], &[])?;
/// # Ok(())
/// # }
/// ```
pub struct PipelineBuilder {
    // Shader stages
    vertex_shader: Option<ShaderStageInfo>,
    fragment_shader: Option<ShaderStageInfo>,
    geometry_shader: Option<ShaderStageInfo>,
    tessellation_control_shader: Option<ShaderStageInfo>,
    tessellation_evaluation_shader: Option<ShaderStageInfo>,

    // Vertex input
    vertex_binding_descriptions: Vec<vk::VertexInputBindingDescription>,
    vertex_attribute_descriptions: Vec<vk::VertexInputAttributeDescription>,

    // Input assembly
    topology: vk::PrimitiveTopology,
    primitive_restart_enable: bool,

    // Viewport/Scissor
    viewports: Vec<vk::Viewport>,
    scissors: Vec<vk::Rect2D>,

    // Rasterization
    polygon_mode: vk::PolygonMode,
    cull_mode: vk::CullModeFlags,
    front_face: vk::FrontFace,
    depth_bias_enable: bool,
    line_width: f32,

    // Multisample
    sample_count: vk::SampleCountFlags,

    // Depth/Stencil
    depth_test_enable: bool,
    depth_write_enable: bool,
    depth_compare_op: vk::CompareOp,
    stencil_test_enable: bool,

    // Color blend
    color_blend_attachments: Vec<vk::PipelineColorBlendAttachmentState>,
    blend_constants: [f32; 4],

    // Dynamic state
    dynamic_states: Vec<vk::DynamicState>,

    // Render pass (for traditional rendering)
    render_pass: Option<vk::RenderPass>,
    subpass: u32,

    // Dynamic rendering (Vulkan 1.3+)
    color_attachment_formats: Vec<vk::Format>,
    depth_attachment_format: Option<vk::Format>,
    stencil_attachment_format: Option<vk::Format>,
}

impl Default for PipelineBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl PipelineBuilder {
    /// Create a new pipeline builder with default values
    pub fn new() -> Self {
        Self {
            vertex_shader: None,
            fragment_shader: None,
            geometry_shader: None,
            tessellation_control_shader: None,
            tessellation_evaluation_shader: None,
            vertex_binding_descriptions: Vec::new(),
            vertex_attribute_descriptions: Vec::new(),
            topology: vk::PrimitiveTopology::TRIANGLE_LIST,
            primitive_restart_enable: false,
            viewports: Vec::new(),
            scissors: Vec::new(),
            polygon_mode: vk::PolygonMode::FILL,
            cull_mode: vk::CullModeFlags::BACK,
            front_face: vk::FrontFace::COUNTER_CLOCKWISE,
            depth_bias_enable: false,
            line_width: 1.0,
            sample_count: vk::SampleCountFlags::TYPE_1,
            depth_test_enable: true,
            depth_write_enable: true,
            depth_compare_op: vk::CompareOp::LESS,
            stencil_test_enable: false,
            color_blend_attachments: vec![vk::PipelineColorBlendAttachmentState {
                blend_enable: vk::FALSE,
                color_write_mask: vk::ColorComponentFlags::RGBA,
                ..Default::default()
            }],
            blend_constants: [0.0; 4],
            dynamic_states: Vec::new(),
            render_pass: None,
            subpass: 0,
            color_attachment_formats: Vec::new(),
            depth_attachment_format: None,
            stencil_attachment_format: None,
        }
    }

    /// Set vertex shader stage
    pub fn vertex_shader(mut self, module: vk::ShaderModule, entry_point: &str) -> Self {
        self.vertex_shader = Some(ShaderStageInfo {
            stage: vk::ShaderStageFlags::VERTEX,
            module,
            entry_point: entry_point.to_string(),
        });
        self
    }

    /// Set fragment shader stage
    pub fn fragment_shader(mut self, module: vk::ShaderModule, entry_point: &str) -> Self {
        self.fragment_shader = Some(ShaderStageInfo {
            stage: vk::ShaderStageFlags::FRAGMENT,
            module,
            entry_point: entry_point.to_string(),
        });
        self
    }

    /// Set geometry shader stage
    pub fn geometry_shader(mut self, module: vk::ShaderModule, entry_point: &str) -> Self {
        self.geometry_shader = Some(ShaderStageInfo {
            stage: vk::ShaderStageFlags::GEOMETRY,
            module,
            entry_point: entry_point.to_string(),
        });
        self
    }

    /// Set vertex input bindings
    pub fn vertex_bindings(mut self, bindings: Vec<vk::VertexInputBindingDescription>) -> Self {
        self.vertex_binding_descriptions = bindings;
        self
    }

    /// Set vertex input attributes
    pub fn vertex_attributes(
        mut self,
        attributes: Vec<vk::VertexInputAttributeDescription>,
    ) -> Self {
        self.vertex_attribute_descriptions = attributes;
        self
    }

    /// Set primitive topology
    pub fn topology(mut self, topology: vk::PrimitiveTopology) -> Self {
        self.topology = topology;
        self
    }

    /// Set viewport
    pub fn viewport(mut self, viewport: vk::Viewport) -> Self {
        self.viewports = vec![viewport];
        self
    }

    /// Set multiple viewports
    pub fn viewports(mut self, viewports: Vec<vk::Viewport>) -> Self {
        self.viewports = viewports;
        self
    }

    /// Set scissor rectangle
    pub fn scissor(mut self, scissor: vk::Rect2D) -> Self {
        self.scissors = vec![scissor];
        self
    }

    /// Set multiple scissors
    pub fn scissors(mut self, scissors: Vec<vk::Rect2D>) -> Self {
        self.scissors = scissors;
        self
    }

    /// Set polygon mode
    pub fn polygon_mode(mut self, mode: vk::PolygonMode) -> Self {
        self.polygon_mode = mode;
        self
    }

    /// Set cull mode
    pub fn cull_mode(mut self, mode: vk::CullModeFlags) -> Self {
        self.cull_mode = mode;
        self
    }

    /// Set front face winding
    pub fn front_face(mut self, front_face: vk::FrontFace) -> Self {
        self.front_face = front_face;
        self
    }

    /// Set line width
    pub fn line_width(mut self, width: f32) -> Self {
        self.line_width = width;
        self
    }

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

    /// Enable/disable depth testing
    pub fn depth_test(mut self, enable: bool) -> Self {
        self.depth_test_enable = enable;
        self
    }

    /// Enable/disable depth writes
    pub fn depth_write(mut self, enable: bool) -> Self {
        self.depth_write_enable = enable;
        self
    }

    /// Set depth compare operation
    pub fn depth_compare_op(mut self, op: vk::CompareOp) -> Self {
        self.depth_compare_op = op;
        self
    }

    /// Set color blend attachments
    pub fn color_blend_attachments(
        mut self,
        attachments: Vec<vk::PipelineColorBlendAttachmentState>,
    ) -> Self {
        self.color_blend_attachments = attachments;
        self
    }

    /// Set dynamic states
    pub fn dynamic_states(mut self, states: Vec<vk::DynamicState>) -> Self {
        self.dynamic_states = states;
        self
    }

    /// Set render pass for traditional rendering
    pub fn render_pass(mut self, render_pass: vk::RenderPass, subpass: u32) -> Self {
        self.render_pass = Some(render_pass);
        self.subpass = subpass;
        self
    }

    /// Set color attachment formats for dynamic rendering (Vulkan 1.3+)
    pub fn color_attachment_formats(mut self, formats: Vec<vk::Format>) -> Self {
        self.color_attachment_formats = formats;
        self
    }

    /// Set depth attachment format for dynamic rendering
    pub fn depth_attachment_format(mut self, format: vk::Format) -> Self {
        self.depth_attachment_format = Some(format);
        self
    }

    /// Build a graphics pipeline
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the pipeline on
    /// * `descriptor_set_layouts` - Descriptor set layouts for the pipeline
    /// * `push_constant_ranges` - Push constant ranges
    ///
    /// # Errors
    ///
    /// Returns error if vertex shader is missing or pipeline creation fails.
    pub fn build_graphics(
        self,
        device: &Device,
        descriptor_set_layouts: &[vk::DescriptorSetLayout],
        push_constant_ranges: &[vk::PushConstantRange],
    ) -> Result<Pipeline, PipelineError> {
        // Validate required fields
        let vertex_shader = self
            .vertex_shader
            .ok_or(PipelineError::MissingField("vertex_shader"))?;

        // Build shader stages
        let mut shader_stages = Vec::new();

        // Collect all shader infos that are present
        let shader_infos: Vec<&ShaderStageInfo> = vec![
            Some(&vertex_shader),
            self.fragment_shader.as_ref(),
            self.geometry_shader.as_ref(),
            self.tessellation_control_shader.as_ref(),
            self.tessellation_evaluation_shader.as_ref(),
        ]
        .into_iter()
        .flatten()
        .collect();

        // Create entry point CStrings
        let entry_points: Vec<_> = shader_infos
            .iter()
            .map(|s| std::ffi::CString::new(s.entry_point.clone()).unwrap())
            .collect();

        // Build shader stage create infos
        for (idx, info) in shader_infos.iter().enumerate() {
            shader_stages.push(vk::PipelineShaderStageCreateInfo {
                stage: info.stage,
                module: info.module,
                p_name: entry_points[idx].as_ptr(),
                ..Default::default()
            });
        }

        // Vertex input state
        let vertex_input_state = vk::PipelineVertexInputStateCreateInfo {
            vertex_binding_description_count: self.vertex_binding_descriptions.len() as u32,
            p_vertex_binding_descriptions: self.vertex_binding_descriptions.as_ptr(),
            vertex_attribute_description_count: self.vertex_attribute_descriptions.len() as u32,
            p_vertex_attribute_descriptions: self.vertex_attribute_descriptions.as_ptr(),
            ..Default::default()
        };

        // Input assembly state
        let input_assembly_state = vk::PipelineInputAssemblyStateCreateInfo {
            topology: self.topology,
            primitive_restart_enable: if self.primitive_restart_enable {
                vk::TRUE
            } else {
                vk::FALSE
            },
            ..Default::default()
        };

        // Viewport state
        let viewport_state = vk::PipelineViewportStateCreateInfo {
            viewport_count: if self.viewports.is_empty() {
                1
            } else {
                self.viewports.len() as u32
            },
            p_viewports: if self.viewports.is_empty() {
                std::ptr::null()
            } else {
                self.viewports.as_ptr()
            },
            scissor_count: if self.scissors.is_empty() {
                1
            } else {
                self.scissors.len() as u32
            },
            p_scissors: if self.scissors.is_empty() {
                std::ptr::null()
            } else {
                self.scissors.as_ptr()
            },
            ..Default::default()
        };

        // Rasterization state
        let rasterization_state = vk::PipelineRasterizationStateCreateInfo {
            polygon_mode: self.polygon_mode,
            cull_mode: self.cull_mode,
            front_face: self.front_face,
            depth_bias_enable: if self.depth_bias_enable {
                vk::TRUE
            } else {
                vk::FALSE
            },
            line_width: self.line_width,
            ..Default::default()
        };

        // Multisample state
        let multisample_state = vk::PipelineMultisampleStateCreateInfo {
            rasterization_samples: self.sample_count,
            ..Default::default()
        };

        // Depth/stencil state
        let depth_stencil_state = vk::PipelineDepthStencilStateCreateInfo {
            depth_test_enable: if self.depth_test_enable {
                vk::TRUE
            } else {
                vk::FALSE
            },
            depth_write_enable: if self.depth_write_enable {
                vk::TRUE
            } else {
                vk::FALSE
            },
            depth_compare_op: self.depth_compare_op,
            stencil_test_enable: if self.stencil_test_enable {
                vk::TRUE
            } else {
                vk::FALSE
            },
            ..Default::default()
        };

        // Color blend state
        let color_blend_state = vk::PipelineColorBlendStateCreateInfo {
            attachment_count: self.color_blend_attachments.len() as u32,
            p_attachments: self.color_blend_attachments.as_ptr(),
            blend_constants: self.blend_constants,
            ..Default::default()
        };

        // Dynamic state
        let dynamic_state = if !self.dynamic_states.is_empty() {
            Some(vk::PipelineDynamicStateCreateInfo {
                dynamic_state_count: self.dynamic_states.len() as u32,
                p_dynamic_states: self.dynamic_states.as_ptr(),
                ..Default::default()
            })
        } else {
            None
        };

        // Create pipeline layout
        let layout = PipelineLayout::new(device, descriptor_set_layouts, push_constant_ranges)?;

        // Create pipeline
        let mut pipeline_info = vk::GraphicsPipelineCreateInfo {
            stage_count: shader_stages.len() as u32,
            p_stages: shader_stages.as_ptr(),
            p_vertex_input_state: &vertex_input_state,
            p_input_assembly_state: &input_assembly_state,
            p_viewport_state: &viewport_state,
            p_rasterization_state: &rasterization_state,
            p_multisample_state: &multisample_state,
            p_depth_stencil_state: &depth_stencil_state,
            p_color_blend_state: &color_blend_state,
            p_dynamic_state: dynamic_state
                .as_ref()
                .map(|s| s as *const _)
                .unwrap_or(std::ptr::null()),
            layout: layout.handle(),
            render_pass: self.render_pass.unwrap_or(vk::RenderPass::null()),
            subpass: self.subpass,
            ..Default::default()
        };

        // Add dynamic rendering info if using Vulkan 1.3+
        let mut rendering_info = if !self.color_attachment_formats.is_empty() {
            Some(vk::PipelineRenderingCreateInfo {
                color_attachment_count: self.color_attachment_formats.len() as u32,
                p_color_attachment_formats: self.color_attachment_formats.as_ptr(),
                depth_attachment_format: self
                    .depth_attachment_format
                    .unwrap_or(vk::Format::UNDEFINED),
                stencil_attachment_format: self
                    .stencil_attachment_format
                    .unwrap_or(vk::Format::UNDEFINED),
                ..Default::default()
            })
        } else {
            None
        };

        if let Some(ref mut info) = rendering_info {
            pipeline_info.p_next = info as *const _ as *const std::ffi::c_void;
        }

        // SAFETY: device and all create info structs are valid
        let pipelines = unsafe {
            device
                .handle()
                .create_graphics_pipelines(vk::PipelineCache::null(), &[pipeline_info], None)
                .map_err(|e| PipelineError::CreationFailed(e.1))?
        };

        Ok(Pipeline {
            pipeline: pipelines[0],
            layout,
            bind_point: vk::PipelineBindPoint::GRAPHICS,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{Device, DeviceCreateInfo, Instance, InstanceCreateInfo, QueueCreateInfo};

    fn create_test_device() -> (Instance, Device) {
        let instance = Instance::new(InstanceCreateInfo {
            enable_validation: false,
            ..Default::default()
        })
        .unwrap();
        let physical_devices = instance.enumerate_physical_devices().unwrap();
        let physical_device = physical_devices[0];

        let graphics_family = unsafe {
            instance
                .get_physical_device_queue_family_properties(physical_device)
                .iter()
                .enumerate()
                .find(|(_, qf)| qf.queue_flags.contains(vk::QueueFlags::GRAPHICS))
                .map(|(i, _)| i as u32)
                .unwrap()
        };

        let device = Device::new(
            &instance,
            physical_device,
            DeviceCreateInfo {
                queue_create_infos: vec![QueueCreateInfo {
                    queue_family_index: graphics_family,
                    queue_count: 1,
                    queue_priorities: vec![1.0],
                }],
                ..Default::default()
            },
        )
        .unwrap();

        (instance, device)
    }

    #[test]
    fn test_pipeline_layout_creation() {
        let (_instance, device) = create_test_device();

        let layout = PipelineLayout::new(&device, &[], &[]).unwrap();
        assert_ne!(layout.handle(), vk::PipelineLayout::null());
        layout.destroy(&device);
    }

    #[test]
    fn test_pipeline_builder_defaults() {
        let builder = PipelineBuilder::new();
        assert_eq!(builder.topology, vk::PrimitiveTopology::TRIANGLE_LIST);
        assert_eq!(builder.polygon_mode, vk::PolygonMode::FILL);
        assert_eq!(builder.cull_mode, vk::CullModeFlags::BACK);
        assert_eq!(builder.front_face, vk::FrontFace::COUNTER_CLOCKWISE);
        assert!(builder.depth_test_enable);
        assert!(builder.depth_write_enable);
    }

    #[test]
    fn test_pipeline_builder_fluent_api() {
        let builder = PipelineBuilder::new()
            .topology(vk::PrimitiveTopology::LINE_LIST)
            .cull_mode(vk::CullModeFlags::NONE)
            .depth_test(false);

        assert_eq!(builder.topology, vk::PrimitiveTopology::LINE_LIST);
        assert_eq!(builder.cull_mode, vk::CullModeFlags::NONE);
        assert!(!builder.depth_test_enable);
    }
}