shdrlib 0.2.0

use crate::front::storage::ShaderObject;
use crate::front::compiler::{compile_ir, Language};
use crate::back::{GPUDevice, WGPUShader, RenderPipelineBuilder};
use std::collections::HashMap;

/// A pair of shader and its corresponding render pipeline
pub struct Duo {
    pub shader_object: ShaderObject,
    pub pipeline: wgpu::RenderPipeline,
    pub vertex_entry: String,
    pub fragment_entry: Option<String>,
}

impl Duo {
    pub fn new(shader_object: ShaderObject, pipeline: wgpu::RenderPipeline) -> Self {
        let vertex_entry = shader_object.ir.metadata.entry_point.clone();
        Self {
            shader_object,
            pipeline,
            vertex_entry,
            fragment_entry: None,
        }
    }

    pub fn with_fragment_entry(mut self, entry: String) -> Self {
        self.fragment_entry = Some(entry);
        self
    }
}

/// AssetManager handles shader compilation and pipeline creation
/// Stores everything by name for easy retrieval
pub struct AssetManager {
    gpu: GPUDevice,
    shaders: HashMap<String, ShaderObject>,
    pipelines: HashMap<String, Duo>,
}

impl AssetManager {
    /// Create a new asset manager with GPU initialization
    pub fn new() -> Self {
        Self {
            gpu: GPUDevice::new(),
            shaders: HashMap::new(),
            pipelines: HashMap::new(),
        }
    }

    /// Get GPU info
    pub fn gpu_info(&self) -> wgpu::AdapterInfo {
        self.gpu.info()
    }

    /// Add a shader from source code
    pub fn add_shader(
        &mut self,
        name: impl Into<String>,
        source: &str,
        stage: naga::ShaderStage,
        language: Language,
    ) -> Result<(), String> {
        let ir = compile_ir(source, stage, language)?;
        let shader_obj = ShaderObject::new(ir, vec![]);
        self.shaders.insert(name.into(), shader_obj);
        Ok(())
    }

    /// Add a pre-compiled shader object
    pub fn add_shader_object(&mut self, name: impl Into<String>, shader_obj: ShaderObject) {
        self.shaders.insert(name.into(), shader_obj);
    }

    /// Get a reference to a shader by name
    pub fn get_shader(&self, name: &str) -> Option<&ShaderObject> {
        self.shaders.get(name)
    }

    /// Create a render pipeline from vertex and optional fragment shaders
    /// Stores the pipeline with shader objects in a Duo
    pub fn create_pipeline(
        &mut self,
        pipeline_name: impl Into<String>,
        vertex_shader_name: &str,
        fragment_shader_name: Option<&str>,
        vertex_buffers: Vec<wgpu::VertexBufferLayout<'static>>,
        color_format: wgpu::TextureFormat,
    ) -> Result<(), String> {
        // Get vertex shader
        let vertex_obj = self.shaders.get(vertex_shader_name)
            .ok_or_else(|| format!("Vertex shader '{}' not found", vertex_shader_name))?;

        // Create vertex wgpu shader
        let vertex_shader = WGPUShader::from_shader_object(&self.gpu.device, vertex_obj);

        // Build pipeline (with optional fragment shader)
        let mut builder = RenderPipelineBuilder::new()
            .vertex_shader(vertex_shader)
            .vertex_buffers(vertex_buffers)
            .color_target_format(color_format);

        // Add fragment shader if provided
        if let Some(frag_name) = fragment_shader_name {
            let frag_obj = self.shaders.get(frag_name)
                .ok_or_else(|| format!("Fragment shader '{}' not found", frag_name))?;
            let fragment_shader = WGPUShader::from_shader_object(&self.gpu.device, frag_obj);
            builder = builder.fragment_shader(fragment_shader);
        }

        // Create pipeline layout (empty for now, can be customized)
        let pipeline_layout = self.gpu.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Pipeline Layout"),
            bind_group_layouts: &[],
            push_constant_ranges: &[],
        });

        // Build the pipeline
        let pipeline = builder.build(&self.gpu.device, &pipeline_layout)?;

        // Create Duo with the vertex shader object
        // (In a real app, you might want to store both vertex and fragment)
        let duo = Duo::new(vertex_obj.clone(), pipeline);

        self.pipelines.insert(pipeline_name.into(), duo);
        Ok(())
    }

    /// Get a reference to a pipeline Duo by name
    pub fn get_pipeline(&self, name: &str) -> Option<&Duo> {
        self.pipelines.get(name)
    }

    /// Get a mutable reference to a pipeline Duo by name
    pub fn get_pipeline_mut(&mut self, name: &str) -> Option<&mut Duo> {
        self.pipelines.get_mut(name)
    }

    /// List all registered shader names
    pub fn list_shaders(&self) -> Vec<&String> {
        self.shaders.keys().collect()
    }

    /// List all registered pipeline names
    pub fn list_pipelines(&self) -> Vec<&String> {
        self.pipelines.keys().collect()
    }

    /// Get reference to the GPU device
    pub fn device(&self) -> &wgpu::Device {
        &self.gpu.device
    }

    /// Get reference to the GPU queue
    pub fn queue(&self) -> &wgpu::Queue {
        &self.gpu.queue
    }

    /// Get reference to the adapter
    pub fn adapter(&self) -> &wgpu::Adapter {
        self.gpu.adapter()
    }

    /// Get reference to the instance
    pub fn instance(&self) -> &wgpu::Instance {
        self.gpu.instance()
    }
}

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