wsg/lib/src/pipeline/pipeline_cache.rs

//! # PipelineCache Module — Translation of resources/ Data Toward GPU Pipelines
//!
//! Defines `PipelineCache`, the library's shader compilation cache. It translates WGSL shader source and resources/ data types
//! into compiled RenderPipelines, storing them in a HashMap keyed by shader_id + texture format to avoid duplicate GPU work.
//! Materials request pipelines through this cache; if a pipeline for the given key exists, it is returned directly
//! via Arc cloning. Otherwise the cache compiles one on-the-fly, caches it, then returns it.
//!
//! ## Interaction with Other Modules
//! - **Material** calls `get_or_create()` during its own construction to obtain a shared RenderPipeline.
//! - **conf::STANDARD_SHADER** provides fallback WGSL source when an external file is not found.
//! - **vertex::Vertex** defines the CPU-side layout that `build_pipeline` uses as the vertex buffer contract.
//!
//! ## Technical Points
//! - Pipelines are stored behind `Arc` so multiple Materials share the same compiled object without copying.
//! - wgpu 30 requires `compilation_options` in VertexState/FragmentState and `depth_slice` in color attachments.
//! - **Batching**: Multiple Materials with the same shader_id share one pipeline, enabling material-level batching in Renderer.

use crate::resources::{Texture, Vertex};
use crate::utils::STANDARD_SHADER;

use std::collections::HashMap;
use std::sync::Arc;

/// Creates the two bind group layouts shared by **every** pipeline (Étape 3 — décision actée
/// « un seul layout pour tous »). Both buffers are `Uniform`, 16-byte aligned, no dynamic offset.
/// Matching CPU types: `FrameUniforms` (192 B) and `ObjectUniform` (64 B) in `resources::uniform`.
/// Returns `[frame_layout, object_layout]` in renderer binding order.
///
/// - `index 0` : per-frame uniforms (view/proj/light/options), visible in both shader stages.
/// - `index 1` : per-object uniforms (model matrix), visible in the vertex stage only.
pub fn create_uniform_bind_group_layouts(device: &wgpu::Device) -> [wgpu::BindGroupLayout; 2] {
    [
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("frame_uniform_layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        }),
        device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("object_uniform_layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        }),
    ]
}

/// Creates the texture bind group layout (group 2) shared by every pipeline (Étape 10, DRAFT D1).
/// Binds the diffuse texture + its sampler in the **fragment** stage only. Added to every pipeline
/// layout alongside the frame (@0) + object (@1) uniform groups, so « un seul layout pour tous »
/// (Étape 3) is preserved: a texture-less `Material` binds the white 1×1 placeholder instead.
///
/// - `binding 0` : sampler (filtering, linear/repeat — D3).
/// - `binding 1` : `texture_2d<f32>` diffuse.
pub fn create_texture_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("texture_bind_group_layout"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
                    view_dimension: wgpu::TextureViewDimension::D2,
                    multisampled: false,
                },
                count: None,
            },
        ],
    })
}

/// Creates the **shadow map** bind group layout (group 3) shared by every main pipeline (Étape 14,
/// DRAFT D1/D5). Binds a **comparison** sampler + a depth texture so the fragment can run a PCF
/// `textureSampleCompare` against the shadow map. Added to every pipeline layout alongside groups
/// 0–2, keeping « un seul layout pour tous » — shadows are simply a no-op when disabled.
///
/// - `binding 0` : `sampler_comparison` (compare fn drives the shadow test, D5).
/// - `binding 1` : `texture_depth_2d` (the shadow map).
pub fn create_shadow_map_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("shadow_map_bind_group_layout"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Comparison),
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    sample_type: wgpu::TextureSampleType::Depth,
                    view_dimension: wgpu::TextureViewDimension::D2,
                    multisampled: false,
                },
                count: None,
            },
        ],
    })
}

/// Creates the **shadow uniform** bind group layout (group 0 of the depth-only shadow pipeline,
/// Étape 14, D4): a single uniform buffer holding the light's `view_proj` matrix. Read in the
/// **vertex** stage only (the shadow shader transforms vertices into light-clip space).
pub fn create_shadow_uniform_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("shadow_uniform_layout"),
        entries: &[wgpu::BindGroupLayoutEntry {
            binding: 0,
            visibility: wgpu::ShaderStages::VERTEX,
            ty: wgpu::BindingType::Buffer {
                ty: wgpu::BufferBindingType::Uniform,
                has_dynamic_offset: false,
                min_binding_size: None,
            },
            count: None,
        }],
    })
}

/// The shared GPU `Vertex`-buffer layout used by **every** pipeline that renders mesh geometry
/// (both the main `build_pipeline` and the depth-only shadow pipeline). The array stride equals
/// `size_of::<Vertex>()` so it matches the mesh vertex buffers exactly; the four attributes are
/// declared position (loc 0), normal (1), uv (2), color (3).
pub fn vertex_buffer_layout() -> wgpu::VertexBufferLayout<'static> {
    wgpu::VertexBufferLayout {
        array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
        step_mode: wgpu::VertexStepMode::Vertex,
        attributes: &[
            wgpu::VertexAttribute {
                offset: 0,
                shader_location: 0,
                format: wgpu::VertexFormat::Float32x3,
            }, // position
            wgpu::VertexAttribute {
                offset: 12,
                shader_location: 1,
                format: wgpu::VertexFormat::Float32x3,
            }, // normal
            wgpu::VertexAttribute {
                offset: 24,
                shader_location: 2,
                format: wgpu::VertexFormat::Float32x2,
            }, // uv
            wgpu::VertexAttribute {
                offset: 32,
                shader_location: 3,
                format: wgpu::VertexFormat::Float32x4,
            }, // color
        ],
    }
}

/// Depth texture format shared by the whole library (Étape 9, décision D1 du 2026-09-18).
///
/// Single z-buffer format used for **both** the depth attachment textures (`Renderer`) and the
/// `DepthStencilState` of every pipeline (`build_pipeline`). Keeping them on the same constant
/// guarantees by construction that the pipeline depth format always matches the texture format
/// (wgpu validation error otherwise). `Depth32Float` = portée maximale (comparaison précise),
/// avec clear `1.0` (profondeur maximale au loin), `depth_compare: Less`, write enabled.
pub const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;

/// Shader pipeline cache: maps (shader_id, format) keys to compiled RenderPipelines.
/// Ensures each unique shader+format combination is compiled at most once; subsequent requests return cached instances.
pub struct PipelineCache {
    device: Arc<wgpu::Device>,
    /// Cached pipelines keyed by their shader identifier string. Multiple Materials sharing the same ID share one Arc-wrapped pipeline.
    pipelines: HashMap<String, Arc<wgpu::RenderPipeline>>,
    /// Maps shader IDs to file paths on disk for WGSL loading in `load_shader()`.
    shader_paths: HashMap<String, String>,
    /// Shared bind group layout for the texture group (`@group(2)`), used by every pipeline and by
    /// every Material's texture bind group (Étape 10, DRAFT D1 : « un seul layout pour tous »).
    texture_bind_group_layout: wgpu::BindGroupLayout,
    /// White 1×1 placeholder texture bound by materials that have no diffuse texture (DRAFT D1/D2).
    /// A white texel is the multiplicative identity, so sampling it reproduces the pre-Étape-10 look.
    placeholder: Arc<Texture>,
}

impl PipelineCache {
    /// Creates an empty pipeline cache with no pre-loaded shaders or pipelines, plus the shared
    /// texture bind group layout (group 2) and the white placeholder texture (Étape 10).
    /// Inputs: device (owned Arc reference to wgpu Device), queue (used once to upload the white
    /// placeholder). Returns a new PipelineCache ready for shader registration via register_shader().
    /// Called at application startup before any Material creation. Shader paths must be registered via register_shader() first.
    pub fn new(device: Arc<wgpu::Device>, queue: wgpu::Queue) -> Self {
        let placeholder = Texture::white_placeholder(&device, &queue).arc();
        let texture_bind_group_layout = create_texture_bind_group_layout(&device);
        Self {
            device,
            pipelines: HashMap::new(),
            // Maps shader IDs to file paths on disk for WGSL loading in load_shader().
            // When a path exists, it reads from it; otherwise falls back to STANDARD_SHADER constant.
            shader_paths: HashMap::new(),
            texture_bind_group_layout,
            placeholder,
        }
    }

    /// Returns the shared white placeholder texture, bound by `Material`s without a diffuse texture.
    /// Called by `Material` construction (through [`PipelineCache::texture_bind_group`]) and by
    /// `Scene::get_texture` fallbacks. Étape 10 (DRAFT D1/D2).
    pub fn placeholder(&self) -> &Arc<Texture> {
        &self.placeholder
    }

    /// Returns a reference to the shared group-2 bind group layout (sampler + texture), used by
    /// every Material to build its texture bind group. Étape 10 (DRAFT D1).
    pub fn texture_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
        &self.texture_bind_group_layout
    }

    /// Builds a group-2 bind group for a Material from its diffuse texture (or the white placeholder
    /// when `texture` is `None`). Centralizes the sampler+texture binding so `Material` never touches
    /// wgpu directly (Étape 10, DRAFT D4). Inputs: texture — the material's diffuse texture, `None`
    /// for a texture-less material (binds the placeholder). Returns the group-2 bind group.
    pub fn texture_bind_group(&self, texture: Option<Arc<Texture>>) -> wgpu::BindGroup {
        let tex = texture.unwrap_or_else(|| self.placeholder.clone());
        self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("texture bind group"),
            layout: &self.texture_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::Sampler(&tex.sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(&tex.view),
                },
            ],
        })
    }
    /// Registers an external WGSL shader file path associated with a given ID.
    /// Inputs: id (unique key for this shader), path (filesystem path to .wgsl file).
    /// Returns Ok(id) on success or Err(String) if the ID is already registered. Called during scene setup to register custom shaders.
    pub fn register_shader(&mut self, id: &str, path: &str) -> Result<String, String> {
        if self.shader_paths.contains_key(id) {
            return Err(format!("ID '{}' already exists.", id));
        }
        self.shader_paths.insert(id.to_string(), path.to_string());
        Ok(id.to_string())
    }

    /// Unregisters a shader by its ID, removing both the path reference and any cached compiled pipeline.
    /// Inputs: id (the shader identifier to remove).
    /// Returns Ok(id) on success or Err(String) if the ID does not exist. Called when a shader should be freed from GPU memory.
    pub fn unregister_shader(&mut self, id: &str) -> Result<String, String> {
        if self.shader_paths.remove(id).is_none() {
            return Err(format!("ID '{}' does not exist.", id));
        }
        // Remove cached pipeline so GPU memory is freed (wgpu drops it automatically)
        self.pipelines.remove(id);

        Ok(id.to_string())
    }

    /// Retrieves a cached RenderPipeline by shader_id, or creates one on-demand if not present.
    /// Inputs: format (surface texture format for fragment output), shader_id (unique key into the cache).
    /// Returns an Arc-wrapped RenderPipeline ready for rendering. Called by Material::new().
    /// Internal steps: 1) check pipelines HashMap for existing entry →
    ///   2a) if found: clone Arc and return →
    ///   2b) if not found: load_shader() + build_pipeline() → cache behind Arc → insert and return.
    pub fn get_or_create(
        &mut self,
        format: wgpu::TextureFormat,
        shader_id: &str,
    ) -> Arc<wgpu::RenderPipeline> {
        // Step 1: Return cached pipeline if it already exists for this shader_id
        if let Some(pipeline) = self.pipelines.get(shader_id) {
            return pipeline.clone();
        }

        // Step 2: Compile a new pipeline — loads shader and builds the GPU render pipeline
        let path = self
            .shader_paths
            .get(shader_id)
            .map(|s| s.as_str())
            .unwrap_or(shader_id);
        let shader = self.load_shader(&self.device, path);
        let pipeline = Self::build_pipeline(&self.device, format, &shader);

        // Step 3: Cache the new pipeline behind Arc and return it
        let pipeline_arc = Arc::new(pipeline);
        self.pipelines
            .insert(shader_id.to_string(), pipeline_arc.clone());
        pipeline_arc
    }

    /// Loads a WGSL shader module: reads from disk first, falls back to the embedded STANDARD_SHADER constant.
    /// Inputs: device (GPU command source for shader compilation), path (file path or shader_id string).
    /// Returns a compiled wgpu::ShaderModule. Called internally by `get_or_create()` when compiling a new pipeline.
    fn load_shader(&self, device: &wgpu::Device, path: &str) -> wgpu::ShaderModule {
        let source = std::fs::read_to_string(path).unwrap_or_else(|_| {
            println!("Shader not found: {}, falling back to default", path);
            STANDARD_SHADER.to_string()
        });

        device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(path),
            source: wgpu::ShaderSource::Wgsl(source.into()),
        })
    }

    /// Builds a RenderPipeline from a shader module, device, and surface texture format.
    /// Inputs: device (GPU command source), format (output texture format), shader (compiled WGSL module).
    /// Returns a fully configured RenderPipeline ready for draw calls. Called internally by `get_or_create()`.
    /// Internal steps: 1) define VertexBufferLayout from Vertex struct offsets →
    ///   2) create PipelineLayout with bind_group_layouts + immediate_size →
    ///   3) create RenderPipeline with vertex/fragment states, primitive config, multisample state.
    fn build_pipeline(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        shader: &wgpu::ShaderModule,
    ) -> wgpu::RenderPipeline {
        // Define vertex attribute layout — the contract between CPU vertex data and GPU shader inputs.
        // Must match Vertex struct field offsets exactly.
        let vertex_buffer_layout = vertex_buffer_layout();

        // Pipeline layout — the two uniform bind groups (frame @0 + object @1), the texture
        // bind group (@2, Étape 10 DRAFT D1) AND the shadow-map bind group (@3, Étape 14 D5) are
        // attached to EVERY pipeline (Étape 3, décision actée « un seul layout pour tous »), even
        // if a given shader does not read them.
        // `immediate_size` stays 0 (no var<immediate> used).
        let uniform_layouts = create_uniform_bind_group_layouts(device);
        let texture_layout = create_texture_bind_group_layout(device);
        let shadow_layout = create_shadow_map_bind_group_layout(device);
        let layout_refs: Vec<Option<&wgpu::BindGroupLayout>> = vec![
            Some(&uniform_layouts[0]), // frame @0
            Some(&uniform_layouts[1]), // object @1
            Some(&texture_layout),     // texture @2
            Some(&shadow_layout),      // shadow map @3
        ];
        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("render_pipeline_layout"),
                bind_group_layouts: &layout_refs,
                immediate_size: 0, // no var<immediate> used
            });

        // Create the full RenderPipeline — vertex state + fragment state + primitive configuration.
        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: shader,
                // entry_point is now Option<&str> — Some to specify explicitly, None for auto-detection.
                entry_point: Some("vs_main"),
                compilation_options: Default::default(), // required field in wgpu 30
                buffers: &[Some(vertex_buffer_layout)],
            },
            fragment: Some(wgpu::FragmentState {
                module: shader,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(), // required field in wgpu 30
                // targets is now &[Option<ColorTargetState>] — each wrapped in Some.
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState::default(),
            // Étape 9 (DRAFT 9.3) : depth test activé sur TOUTE pipeline. Le format doit matcher
            // l'attachment depth (DEPTH_FORMAT) — c'est garanti par la constante partagée D1.
            // depth_write_enabled + depth_compare sont des Option en wgpu 30 : Some(true) → on
            // écrit la profondeur ; Some(Less) → le fragment est gardé si son z est plus proche.
            depth_stencil: Some(wgpu::DepthStencilState {
                format: DEPTH_FORMAT,
                depth_write_enabled: Some(true),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState::default(),
            // multiview → replaced by multiview_mask (NonZeroU32) and cache fields in wgpu 30.
            multiview_mask: None,
            cache: None,
        })
    }

    /// Retrieves a cached RenderPipeline by shader_id without creating one.
    /// Inputs: shader_id (unique key into the cache).
    /// Returns Some(`Arc<RenderPipeline>`) if found, None otherwise. Called by renderer code for pipeline inspection.
    pub fn get(&self, shader_id: &str) -> Option<&Arc<wgpu::RenderPipeline>> {
        self.pipelines.get(shader_id)
    }
}

/// Builds the **depth-only shadow pipeline** (Étape 14, D4): a vertex-only pipeline (no fragment
/// stage) that transforms each mesh vertex into the shadow-casting light's clip space, writing only
/// depth. Its layout is [`shadow_uniform_layout`] (group 0 : light `view_proj`) + [`object_layout`]
/// (group 1 : per-entity model matrix — the SAME layout/bind groups the main renderer already caches
/// per entity, so the shadow pass reuses them directly).
///
/// `depth_stencil` writes depth with a slope-scaled bias (D5) to suppress acne on surfaces nearly
/// parallel to the light. The vertex buffer layout is the shared [`vertex_buffer_layout`], so the
/// same mesh vertex/index buffers are reused.
///
/// Inputs: device (GPU), object_layout (the shared per-object bind group layout, group 1).
/// Returns the compiled shadow pipeline, ready to render into a depth attachment.
pub fn build_shadow_pipeline(
    device: &wgpu::Device,
    object_layout: &wgpu::BindGroupLayout,
) -> wgpu::RenderPipeline {
    // Vertex-only shader : this pipeline sets `fragment: None`, so only the depth is produced.
    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("shadow_shader"),
        source: wgpu::ShaderSource::Wgsl(crate::utils::SHADOW_SHADER.into()),
    });

    let shadow_uniform_layout = create_shadow_uniform_layout(device);
    let layout_refs: Vec<Option<&wgpu::BindGroupLayout>> =
        vec![Some(&shadow_uniform_layout), Some(object_layout)];
    let shadow_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("shadow_pipeline_layout"),
        bind_group_layouts: &layout_refs,
        immediate_size: 0,
    });

    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("Shadow Pipeline"),
        layout: Some(&shadow_pipeline_layout),
        // wgpu 30 : vertex state requires `compilation_options`.
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: Some("vs_main"),
            compilation_options: Default::default(),
            buffers: &[Some(vertex_buffer_layout())],
        },
        // Depth-only : no fragment state (no color output, no color target).
        fragment: None,
        primitive: wgpu::PrimitiveState::default(),
        depth_stencil: Some(wgpu::DepthStencilState {
            format: DEPTH_FORMAT,
            depth_write_enabled: Some(true),
            depth_compare: Some(wgpu::CompareFunction::Less),
            stencil: wgpu::StencilState::default(),
            // Étape 14 (D5) : slope-scaled depth bias against acne — surfaces nearly parallel to
            // the light are pushed back slightly in the shadow map so they do not self-shadow.
            bias: wgpu::DepthBiasState {
                constant: 2,
                slope_scale: 2.0,
                clamp: 0.0,
            },
        }),
        multisample: wgpu::MultisampleState::default(),
        multiview_mask: None,
        cache: None,
    })
}