shdrlib 0.1.5

//! ShaderManager - Manages shader compilation, reflection, and pipeline creation
//!
//! The ShaderManager handles shader compilation from GLSL source, SPIR-V reflection,
//! and pipeline construction. It owns all shaders and pipelines with correct lifetime
//! management.
//!
//! ## Example
//!
//! ```rust,no_run
//! use shdrlib::ex::*;
//! use shdrlib::core::ShaderStage;
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! # let mut runtime = RuntimeManager::new(RuntimeConfig::default())?;
//! let device = runtime.device();
//! let mut shaders = ShaderManager::new(device)?;
//!
//! // Add shaders
//! let vert_id = shaders.add_shader(VERTEX_GLSL, ShaderStage::Vertex, "triangle_vert")?;
//! let frag_id = shaders.add_shader(FRAGMENT_GLSL, ShaderStage::Fragment, "triangle_frag")?;
//!
//! // Build pipeline
//! let pipeline_id = shaders.build_pipeline(
//!     PipelineBuilder::new()
//!         .vertex_shader(shaders.get_shader(vert_id)?.handle(), "main")
//!         .fragment_shader(shaders.get_shader(frag_id)?.handle(), "main")
//!         .color_attachment_formats(vec![ash::vk::Format::R8G8B8A8_UNORM]),
//!     "triangle_pipeline"
//! )?;
//!
//! // Use pipeline
//! let pipeline = shaders.get_pipeline(pipeline_id)?;
//! # Ok(())
//! # }
//! # const VERTEX_GLSL: &str = "#version 450\nvoid main() {}";
//! # const FRAGMENT_GLSL: &str = "#version 450\nvoid main() {}";
//! ```

use crate::core::{self, Device, Pipeline, Shader, ShaderReflection, ShaderStage};
use crate::ex::errors::{PipelineId, ShaderId, ShaderManagerError};
use crate::ex::PipelineBuilder;
use std::sync::Arc;

/// Entry for a compiled shader with metadata
struct ShaderEntry {
    shader: Shader,
    reflection: Option<ShaderReflection>,
    name: String,
    stage: ShaderStage,
}

/// Entry for a created pipeline with metadata
struct PipelineEntry {
    pipeline: Pipeline,
    name: String,
}

/// ShaderManager - Manages shaders and pipelines with correct lifetime management
///
/// Owns all shaders and pipelines, ensuring correct destruction order.
/// Device is Arc'd from RuntimeManager for safe sharing.
pub struct ShaderManager {
    device: Arc<Device>,
    shaders: Vec<ShaderEntry>,
    pipelines: Vec<PipelineEntry>,
}

impl ShaderManager {
    /// Create a new ShaderManager
    ///
    /// # Arguments
    ///
    /// * `device` - Shared device reference from RuntimeManager
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::ex::*;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let mut runtime = RuntimeManager::new(RuntimeConfig::default())?;
    /// let device = runtime.device();
    /// let shaders = ShaderManager::new(device)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(device: Arc<Device>) -> Result<Self, ShaderManagerError> {
        Ok(Self {
            device,
            shaders: Vec::new(),
            pipelines: Vec::new(),
        })
    }

    /// Add a shader from GLSL source code
    ///
    /// Compiles GLSL to SPIR-V, performs reflection, and creates shader module.
    ///
    /// # Arguments
    ///
    /// * `source` - GLSL source code
    /// * `stage` - Shader stage (Vertex, Fragment, Compute, etc.)
    /// * `name` - Debug name for the shader
    ///
    /// # Returns
    ///
    /// Returns a `ShaderId` that can be used to retrieve the shader later.
    ///
    /// # Errors
    ///
    /// Returns error if GLSL compilation or shader module creation fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::ex::*;
    /// use shdrlib::core::ShaderStage;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let mut runtime = RuntimeManager::new(RuntimeConfig::default())?;
    /// # let device = runtime.device();
    /// let mut shaders = ShaderManager::new(device)?;
    ///
    /// let vert_id = shaders.add_shader(
    ///     r#"#version 450
    ///     layout(location = 0) out vec3 color;
    ///     void main() {
    ///         gl_Position = vec4(0.0, 0.0, 0.0, 1.0);
    ///         color = vec3(1.0, 0.0, 0.0);
    ///     }"#,
    ///     ShaderStage::Vertex,
    ///     "my_vertex_shader"
    /// )?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn add_shader(
        &mut self,
        source: &str,
        stage: ShaderStage,
        name: &str,
    ) -> Result<ShaderId, ShaderManagerError> {
        // Compile GLSL to shader module
        let shader = Shader::from_glsl(&self.device, source, stage, "main")?;

        // Perform reflection (optional - some shaders may not need it)
        let reflection = Self::try_reflect_shader(source, stage);

        // Store shader
        let id = ShaderId(self.shaders.len());
        self.shaders.push(ShaderEntry {
            shader,
            reflection,
            name: name.to_string(),
            stage,
        });

        Ok(id)
    }

    /// Add a shader from pre-compiled SPIR-V bytecode
    ///
    /// # Arguments
    ///
    /// * `spirv` - SPIR-V bytecode (u32 words)
    /// * `stage` - Shader stage
    /// * `name` - Debug name for the shader
    ///
    /// # Returns
    ///
    /// Returns a `ShaderId` that can be used to retrieve the shader later.
    ///
    /// # Errors
    ///
    /// Returns error if shader module creation fails.
    pub fn add_shader_spirv(
        &mut self,
        spirv: &[u32],
        stage: ShaderStage,
        name: &str,
    ) -> Result<ShaderId, ShaderManagerError> {
        let shader = Shader::from_spirv(&self.device, spirv, stage, "main")?;

        // Try to reflect from SPIR-V
        let reflection = ShaderReflection::from_spirv(spirv).ok();

        let id = ShaderId(self.shaders.len());
        self.shaders.push(ShaderEntry {
            shader,
            reflection,
            name: name.to_string(),
            stage,
        });

        Ok(id)
    }

    /// Get a shader by ID
    ///
    /// # Errors
    ///
    /// Returns error if the shader ID is invalid.
    #[inline]
    pub fn get_shader(&self, id: ShaderId) -> Result<&Shader, ShaderManagerError> {
        self.shaders
            .get(id.0)
            .map(|entry| &entry.shader)
            .ok_or(ShaderManagerError::InvalidShaderId(id))
    }

    /// Get shader reflection data by ID
    ///
    /// # Errors
    ///
    /// Returns error if the shader ID is invalid.
    #[inline]
    pub fn shader_reflection(
        &self,
        id: ShaderId,
    ) -> Result<Option<&ShaderReflection>, ShaderManagerError> {
        self.shaders
            .get(id.0)
            .map(|entry| entry.reflection.as_ref())
            .ok_or(ShaderManagerError::InvalidShaderId(id))
    }

    /// Get shader name by ID
    ///
    /// # Errors
    ///
    /// Returns error if the shader ID is invalid.
    pub fn shader_name(&self, id: ShaderId) -> Result<&str, ShaderManagerError> {
        self.shaders
            .get(id.0)
            .map(|entry| entry.name.as_str())
            .ok_or(ShaderManagerError::InvalidShaderId(id))
    }

    /// Get shader stage by ID
    ///
    /// # Errors
    ///
    /// Returns error if the shader ID is invalid.
    pub fn shader_stage(&self, id: ShaderId) -> Result<ShaderStage, ShaderManagerError> {
        self.shaders
            .get(id.0)
            .map(|entry| entry.stage)
            .ok_or(ShaderManagerError::InvalidShaderId(id))
    }

    /// Build a pipeline using PipelineBuilder
    ///
    /// # Arguments
    ///
    /// * `builder` - Configured PipelineBuilder
    /// * `name` - Debug name for the pipeline
    ///
    /// # Returns
    ///
    /// Returns a `PipelineId` that can be used to retrieve the pipeline later.
    ///
    /// # Errors
    ///
    /// Returns error if pipeline creation fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::ex::*;
    /// use shdrlib::core::ShaderStage;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let mut runtime = RuntimeManager::new(RuntimeConfig::default())?;
    /// # let device = runtime.device();
    /// # let mut shaders = ShaderManager::new(device)?;
    /// # const VERT_SRC: &str = "#version 450\nvoid main() {}";
    /// # const FRAG_SRC: &str = "#version 450\nvoid main() {}";
    /// let vert_id = shaders.add_shader(VERT_SRC, ShaderStage::Vertex, "v")?;
    /// let frag_id = shaders.add_shader(FRAG_SRC, ShaderStage::Fragment, "f")?;
    ///
    /// let pipeline_id = shaders.build_pipeline(
    ///     PipelineBuilder::new()
    ///         .vertex_shader(shaders.get_shader(vert_id)?.handle(), "main")
    ///         .fragment_shader(shaders.get_shader(frag_id)?.handle(), "main")
    ///         .color_attachment_formats(vec![ash::vk::Format::R8G8B8A8_UNORM]),
    ///     "my_pipeline"
    /// )?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn build_pipeline(
        &mut self,
        builder: PipelineBuilder,
        name: &str,
    ) -> Result<PipelineId, ShaderManagerError> {
        let pipeline = builder.build_graphics(&self.device)?;

        let id = PipelineId(self.pipelines.len());
        self.pipelines.push(PipelineEntry {
            pipeline,
            name: name.to_string(),
        });

        Ok(id)
    }

    /// Get a pipeline by ID
    ///
    /// # Errors
    ///
    /// Returns error if the pipeline ID is invalid.
    #[inline]
    pub fn get_pipeline(&self, id: PipelineId) -> Result<&Pipeline, ShaderManagerError> {
        self.pipelines
            .get(id.0)
            .map(|entry| &entry.pipeline)
            .ok_or(ShaderManagerError::InvalidPipelineId(id))
    }

    /// Get pipeline name by ID
    ///
    /// # Errors
    ///
    /// Returns error if the pipeline ID is invalid.
    #[inline]
    pub fn pipeline_name(&self, id: PipelineId) -> Result<&str, ShaderManagerError> {
        self.pipelines
            .get(id.0)
            .map(|entry| entry.name.as_str())
            .ok_or(ShaderManagerError::InvalidPipelineId(id))
    }

    /// Get shared device reference
    #[inline]
    pub fn device(&self) -> Arc<Device> {
        Arc::clone(&self.device)
    }

    /// Get number of shaders stored
    #[inline]
    pub fn shader_count(&self) -> usize {
        self.shaders.len()
    }

    /// Get number of pipelines stored
    #[inline]
    pub fn pipeline_count(&self) -> usize {
        self.pipelines.len()
    }

    /// Try to perform shader reflection (doesn't fail if reflection not possible)
    fn try_reflect_shader(source: &str, stage: ShaderStage) -> Option<ShaderReflection> {
        // Compile to SPIR-V for reflection
        match core::ShaderCompiler::compile_glsl(source, stage) {
            Ok(spirv) => ShaderReflection::from_spirv(&spirv).ok(),
            Err(_) => None,
        }
    }
}

impl Drop for ShaderManager {
    fn drop(&mut self) {
        // Pipelines must be destroyed before shaders
        // Manually destroy all resources, then clear the vectors to prevent
        // automatic drop from running on already-destroyed objects

        // Destroy all pipelines first (they reference shaders)
        for entry in &mut self.pipelines {
            entry.pipeline.destroy(&self.device);
        }
        // Clear to prevent double-drop
        self.pipelines.clear();

        // Destroy all shaders second
        for entry in &mut self.shaders {
            entry.shader.destroy(&self.device);
        }
        // Clear to prevent double-drop
        self.shaders.clear();

        // Device Arc will drop last automatically
    }
}

impl std::fmt::Debug for ShaderManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ShaderManager")
            .field("shader_count", &self.shaders.len())
            .field("pipeline_count", &self.pipelines.len())
            .field(
                "shaders",
                &self.shaders.iter().map(|s| &s.name).collect::<Vec<_>>(),
            )
            .field(
                "pipelines",
                &self.pipelines.iter().map(|p| &p.name).collect::<Vec<_>>(),
            )
            .finish()
    }
}

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

    fn test_config() -> RuntimeConfig {
        RuntimeConfig {
            enable_validation: false,
            ..Default::default()
        }
    }

    #[test]
    fn test_shader_manager_creation() {
        let runtime = RuntimeManager::new(test_config()).unwrap();
        let device = runtime.device();
        let shaders = ShaderManager::new(device);
        assert!(shaders.is_ok());
    }

    #[test]
    fn test_add_shader() {
        let runtime = RuntimeManager::new(test_config()).unwrap();
        let device = runtime.device();
        let mut shaders = ShaderManager::new(device).unwrap();

        const SIMPLE_VERT: &str = r#"
            #version 450
            void main() {
                gl_Position = vec4(0.0, 0.0, 0.0, 1.0);
            }
        "#;

        let result = shaders.add_shader(SIMPLE_VERT, ShaderStage::Vertex, "test_vert");
        assert!(result.is_ok());

        let id = result.unwrap();
        assert_eq!(shaders.shader_count(), 1);
        assert!(shaders.get_shader(id).is_ok());
        assert_eq!(shaders.shader_name(id).unwrap(), "test_vert");
    }

    #[test]
    fn test_invalid_shader_id() {
        let runtime = RuntimeManager::new(test_config()).unwrap();
        let device = runtime.device();
        let shaders = ShaderManager::new(device).unwrap();

        let invalid_id = ShaderId(999);
        assert!(shaders.get_shader(invalid_id).is_err());
    }

    #[test]
    fn test_multiple_shaders() {
        let runtime = RuntimeManager::new(test_config()).unwrap();
        let device = runtime.device();
        let mut shaders = ShaderManager::new(device).unwrap();

        const VERT: &str = "#version 450\nvoid main() { gl_Position = vec4(0.0); }";
        const FRAG: &str =
            "#version 450\nlayout(location=0) out vec4 color; void main() { color = vec4(1.0); }";

        let vert_id = shaders
            .add_shader(VERT, ShaderStage::Vertex, "vert")
            .unwrap();
        let frag_id = shaders
            .add_shader(FRAG, ShaderStage::Fragment, "frag")
            .unwrap();

        assert_eq!(shaders.shader_count(), 2);
        assert!(shaders.get_shader(vert_id).is_ok());
        assert!(shaders.get_shader(frag_id).is_ok());
        assert_eq!(shaders.shader_stage(vert_id).unwrap(), ShaderStage::Vertex);
        assert_eq!(
            shaders.shader_stage(frag_id).unwrap(),
            ShaderStage::Fragment
        );
    }
}