shdrlib 0.1.5

//! Shader module - GLSL compilation and SPIR-V management
//!
//! **WARNING**: CORE objects do NOT enforce lifetime dependencies.
//! Shaders must not outlive the device.

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

/// Shader compilation and creation error
#[derive(Debug, Error)]
pub enum ShaderError {
    #[error("Shader compilation failed: {message}")]
    CompilationFailed { message: String, line: Option<u32> },

    #[error("Invalid SPIR-V bytecode")]
    InvalidSpirv,

    #[error("Shader module creation failed: {0}")]
    ModuleCreationFailed(vk::Result),

    #[error("Invalid entry point")]
    InvalidEntryPoint,

    #[error("Reflection failed: {0}")]
    ReflectionFailed(String),

    #[error("Unsupported shader stage")]
    UnsupportedStage,
}

/// Shader stage type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShaderStage {
    Vertex,
    Fragment,
    Compute,
    Geometry,
    TessellationControl,
    TessellationEvaluation,
}

impl ShaderStage {
    /// Convert to Vulkan shader stage flags
    pub fn to_vk_flags(self) -> vk::ShaderStageFlags {
        match self {
            ShaderStage::Vertex => vk::ShaderStageFlags::VERTEX,
            ShaderStage::Fragment => vk::ShaderStageFlags::FRAGMENT,
            ShaderStage::Compute => vk::ShaderStageFlags::COMPUTE,
            ShaderStage::Geometry => vk::ShaderStageFlags::GEOMETRY,
            ShaderStage::TessellationControl => vk::ShaderStageFlags::TESSELLATION_CONTROL,
            ShaderStage::TessellationEvaluation => vk::ShaderStageFlags::TESSELLATION_EVALUATION,
        }
    }

    /// Convert to naga shader stage
    fn to_naga_stage(self) -> naga::ShaderStage {
        match self {
            ShaderStage::Vertex => naga::ShaderStage::Vertex,
            ShaderStage::Fragment => naga::ShaderStage::Fragment,
            ShaderStage::Compute => naga::ShaderStage::Compute,
            _ => naga::ShaderStage::Vertex, // Naga doesn't have all stages
        }
    }
}

impl From<ShaderStage> for naga::ShaderStage {
    fn from(stage: ShaderStage) -> Self {
        stage.to_naga_stage()
    }
}

/// Shader module wrapper
///
/// Represents compiled SPIR-V shader code.
///
/// # Safety
///
/// This object does NOT enforce that the device outlives it.
pub struct Shader {
    module: vk::ShaderModule,
    stage: ShaderStage,
    entry_point: String,
}

impl Shader {
    /// Create a shader from SPIR-V bytecode
    ///
    /// # Arguments
    ///
    /// * \device\ - The device to create the shader on
    /// * \spirv\ - SPIR-V bytecode (u32 words)
    /// * \stage\ - Shader stage
    /// * \entry_point\ - Entry point function name
    pub fn from_spirv(
        device: &Device,
        spirv: &[u32],
        stage: ShaderStage,
        entry_point: &str,
    ) -> Result<Self, ShaderError> {
        let create_info = vk::ShaderModuleCreateInfo {
            code_size: spirv.len() * 4,
            p_code: spirv.as_ptr(),
            ..Default::default()
        };

        // SAFETY: device is valid, create_info properly initialized
        let module = unsafe {
            device
                .handle()
                .create_shader_module(&create_info, None)
                .map_err(ShaderError::ModuleCreationFailed)?
        };

        Ok(Self {
            module,
            stage,
            entry_point: entry_point.to_string(),
        })
    }

    /// Create a shader from GLSL source code
    ///
    /// Compiles GLSL to SPIR-V using naga.
    pub fn from_glsl(
        device: &Device,
        source: &str,
        stage: ShaderStage,
        entry_point: &str,
    ) -> Result<Self, ShaderError> {
        let spirv = ShaderCompiler::compile_glsl(source, stage)?;
        Self::from_spirv(device, &spirv, stage, entry_point)
    }

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

    /// Get the shader stage
    #[inline]
    pub fn stage(&self) -> ShaderStage {
        self.stage
    }

    /// Get the entry point name
    #[inline]
    pub fn entry_point(&self) -> &str {
        &self.entry_point
    }

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

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

/// GLSL to SPIR-V compiler using naga
pub struct ShaderCompiler;

impl ShaderCompiler {
    /// Compile GLSL source to SPIR-V bytecode
    pub fn compile_glsl(source: &str, stage: ShaderStage) -> Result<Vec<u32>, ShaderError> {
        use naga::back::spv;
        use naga::front::glsl;
        use naga::valid::{ValidationFlags, Validator};

        // Parse GLSL
        let mut parser = glsl::Frontend::default();
        let options = glsl::Options {
            stage: stage.to_naga_stage(),
            defines: Default::default(),
        };

        let module = parser.parse(&options, source).map_err(|errors| {
            let message = format!("{:?}", errors);
            ShaderError::CompilationFailed {
                message,
                line: None,
            }
        })?;

        // Validate module
        let mut validator = Validator::new(ValidationFlags::all(), Default::default());
        let info = validator
            .validate(&module)
            .map_err(|e| ShaderError::CompilationFailed {
                message: format!("Validation failed: {:?}", e),
                line: None,
            })?;

        // Generate SPIR-V
        let options = spv::Options {
            lang_version: (1, 0),
            flags: spv::WriterFlags::empty(),
            capabilities: None,
            bounds_check_policies: Default::default(),
            zero_initialize_workgroup_memory:
                naga::back::spv::ZeroInitializeWorkgroupMemoryMode::Polyfill,
            ..Default::default()
        };

        let mut words = Vec::new();
        let mut writer =
            spv::Writer::new(&options).map_err(|e| ShaderError::CompilationFailed {
                message: format!("SPIR-V writer creation failed: {:?}", e),
                line: None,
            })?;

        // Create empty pipeline options (shader not bound to specific pipeline)
        let pipeline_options = naga::back::spv::PipelineOptions {
            shader_stage: stage.into(),
            entry_point: "main".to_string(),
        };

        writer
            .write(&module, &info, Some(&pipeline_options), &None, &mut words)
            .map_err(|e| ShaderError::CompilationFailed {
                message: format!("SPIR-V generation failed: {:?}", e),
                line: None,
            })?;

        Ok(words)
    }
}

/// Shader reflection information
///
/// Extract metadata from SPIR-V bytecode.
pub struct ShaderReflection {
    descriptor_sets: Vec<DescriptorSetReflection>,
    push_constants: Vec<PushConstantReflection>,
}

#[derive(Debug, Clone)]
pub struct DescriptorSetReflection {
    pub set: u32,
    pub bindings: Vec<DescriptorBindingReflection>,
}

#[derive(Debug, Clone)]
pub struct DescriptorBindingReflection {
    pub binding: u32,
    pub descriptor_type: vk::DescriptorType,
    pub descriptor_count: u32,
    pub stage_flags: vk::ShaderStageFlags,
}

#[derive(Debug, Clone)]
pub struct PushConstantReflection {
    pub offset: u32,
    pub size: u32,
    pub stage_flags: vk::ShaderStageFlags,
}

impl ShaderReflection {
    /// Reflect SPIR-V bytecode to extract metadata
    pub fn from_spirv(spirv: &[u32]) -> Result<Self, ShaderError> {
        use spirv_reflect::types::ReflectDescriptorType;

        let module = spirv_reflect::ShaderModule::load_u32_data(spirv)
            .map_err(|e| ShaderError::ReflectionFailed(format!("{:?}", e)))?;

        // Extract descriptor sets
        let descriptor_sets = module
            .enumerate_descriptor_sets(None)
            .map_err(|e| ShaderError::ReflectionFailed(format!("{:?}", e)))?
            .into_iter()
            .map(|set| {
                let bindings = set
                    .bindings
                    .iter()
                    .map(|binding| {
                        let descriptor_type = match binding.descriptor_type {
                            ReflectDescriptorType::Sampler => vk::DescriptorType::SAMPLER,
                            ReflectDescriptorType::CombinedImageSampler => {
                                vk::DescriptorType::COMBINED_IMAGE_SAMPLER
                            }
                            ReflectDescriptorType::SampledImage => {
                                vk::DescriptorType::SAMPLED_IMAGE
                            }
                            ReflectDescriptorType::StorageImage => {
                                vk::DescriptorType::STORAGE_IMAGE
                            }
                            ReflectDescriptorType::UniformBuffer => {
                                vk::DescriptorType::UNIFORM_BUFFER
                            }
                            ReflectDescriptorType::StorageBuffer => {
                                vk::DescriptorType::STORAGE_BUFFER
                            }
                            ReflectDescriptorType::UniformBufferDynamic => {
                                vk::DescriptorType::UNIFORM_BUFFER_DYNAMIC
                            }
                            ReflectDescriptorType::StorageBufferDynamic => {
                                vk::DescriptorType::STORAGE_BUFFER_DYNAMIC
                            }
                            ReflectDescriptorType::InputAttachment => {
                                vk::DescriptorType::INPUT_ATTACHMENT
                            }
                            _ => vk::DescriptorType::UNIFORM_BUFFER,
                        };

                        DescriptorBindingReflection {
                            binding: binding.binding,
                            descriptor_type,
                            descriptor_count: binding.count,
                            stage_flags: vk::ShaderStageFlags::ALL,
                        }
                    })
                    .collect();

                DescriptorSetReflection {
                    set: set.set,
                    bindings,
                }
            })
            .collect();

        // Extract push constants
        let push_constants = module
            .enumerate_push_constant_blocks(None)
            .map_err(|e| ShaderError::ReflectionFailed(format!("{:?}", e)))?
            .into_iter()
            .map(|block| PushConstantReflection {
                offset: block.offset,
                size: block.size,
                stage_flags: vk::ShaderStageFlags::ALL,
            })
            .collect();

        Ok(Self {
            descriptor_sets,
            push_constants,
        })
    }

    /// Get descriptor set reflections
    pub fn descriptor_sets(&self) -> &[DescriptorSetReflection] {
        &self.descriptor_sets
    }

    /// Get push constant reflections
    pub fn push_constants(&self) -> &[PushConstantReflection] {
        &self.push_constants
    }
}

#[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_shader_stage_to_vk_flags() {
        assert_eq!(
            ShaderStage::Vertex.to_vk_flags(),
            vk::ShaderStageFlags::VERTEX
        );
        assert_eq!(
            ShaderStage::Fragment.to_vk_flags(),
            vk::ShaderStageFlags::FRAGMENT
        );
        assert_eq!(
            ShaderStage::Compute.to_vk_flags(),
            vk::ShaderStageFlags::COMPUTE
        );
    }

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

        // Minimal valid SPIR-V (just the header)
        let spirv = vec![
            0x07230203, // Magic number
            0x00010000, // Version 1.0
            0x00000000, // Generator
            0x00000001, // Bound
            0x00000000, // Schema
        ];

        let shader = Shader::from_spirv(&device, &spirv, ShaderStage::Vertex, "main").unwrap();

        assert_ne!(shader.handle(), vk::ShaderModule::null());
        assert_eq!(shader.stage(), ShaderStage::Vertex);
        assert_eq!(shader.entry_point(), "main");

        shader.destroy(&device);
    }

    #[test]
    fn test_glsl_compilation() {
        let glsl = r#"
            #version 450
            void main() {
                gl_Position = vec4(0.0, 0.0, 0.0, 1.0);
            }
        "#;

        let result = ShaderCompiler::compile_glsl(glsl, ShaderStage::Vertex);
        // Note: This might fail with naga's GLSL parser, which is limited
        // The test validates the API exists, not that naga can compile everything
        match result {
            Ok(spirv) => {
                assert!(!spirv.is_empty());
                assert_eq!(spirv[0], 0x07230203); // SPIR-V magic number
            }
            Err(e) => {
                // Naga GLSL support is limited, this is expected
                println!("GLSL compilation failed (expected with naga): {}", e);
            }
        }
    }
}