rdpe 0.1.0

//! Custom uniforms for passing runtime data to shaders.
#![allow(clippy::too_many_arguments)]
//!
//! Custom uniforms let you pass dynamic values to your particle rules,
//! enabling interactive and reactive simulations.
//!
//! # Example
//!
//! ```ignore
//! Simulation::<Particle>::new()
//!     .with_uniform("attractor", Vec3::ZERO)
//!     .with_uniform("strength", 1.0f32)
//!     .with_update(|ctx| {
//!         // Update uniforms based on input - mouse_world_pos() maps to world coordinates!
//!         if ctx.input.mouse_held(MouseButton::Left) {
//!             ctx.set("attractor", ctx.mouse_world_pos());
//!         }
//!         ctx.set("strength", (ctx.time() * 2.0).sin() * 0.5 + 1.0);
//!     })
//!     .with_rule(Rule::Custom(r#"
//!         let dir = uniforms.attractor - p.position;
//!         p.velocity += normalize(dir) * uniforms.strength * uniforms.delta_time;
//!     "#.into()))
//!     .run();
//! ```

use crate::input::{Input, KeyCode, MouseButton};
use glam::{Vec2, Vec3, Vec4};
use std::collections::HashMap;

/// Supported uniform value types.
#[derive(Clone, Copy, Debug)]
pub enum UniformValue {
    F32(f32),
    I32(i32),
    U32(u32),
    Vec2(Vec2),
    Vec3(Vec3),
    Vec4(Vec4),
}

impl UniformValue {
    /// Get the WGSL type name for this value.
    pub fn wgsl_type(&self) -> &'static str {
        match self {
            UniformValue::F32(_) => "f32",
            UniformValue::I32(_) => "i32",
            UniformValue::U32(_) => "u32",
            UniformValue::Vec2(_) => "vec2<f32>",
            UniformValue::Vec3(_) => "vec3<f32>",
            UniformValue::Vec4(_) => "vec4<f32>",
        }
    }

    /// Get the byte size of this value (without trailing padding).
    pub fn byte_size(&self) -> usize {
        match self {
            UniformValue::F32(_) => 4,
            UniformValue::I32(_) => 4,
            UniformValue::U32(_) => 4,
            UniformValue::Vec2(_) => 8,
            UniformValue::Vec3(_) => 12, // 12 bytes, aligned to 16
            UniformValue::Vec4(_) => 16,
        }
    }

    /// Write this value to a byte buffer.
    pub fn write_bytes(&self, buf: &mut Vec<u8>) {
        match self {
            UniformValue::F32(v) => buf.extend_from_slice(&v.to_le_bytes()),
            UniformValue::I32(v) => buf.extend_from_slice(&v.to_le_bytes()),
            UniformValue::U32(v) => buf.extend_from_slice(&v.to_le_bytes()),
            UniformValue::Vec2(v) => {
                buf.extend_from_slice(&v.x.to_le_bytes());
                buf.extend_from_slice(&v.y.to_le_bytes());
            }
            UniformValue::Vec3(v) => {
                buf.extend_from_slice(&v.x.to_le_bytes());
                buf.extend_from_slice(&v.y.to_le_bytes());
                buf.extend_from_slice(&v.z.to_le_bytes());
                // No padding here - next value handles its own alignment
                // In std140, scalars can fit in vec3's trailing bytes
            }
            UniformValue::Vec4(v) => {
                buf.extend_from_slice(&v.x.to_le_bytes());
                buf.extend_from_slice(&v.y.to_le_bytes());
                buf.extend_from_slice(&v.z.to_le_bytes());
                buf.extend_from_slice(&v.w.to_le_bytes());
            }
        }
    }
}

// Conversion traits for ergonomic API
impl From<f32> for UniformValue {
    fn from(v: f32) -> Self {
        UniformValue::F32(v)
    }
}

impl From<i32> for UniformValue {
    fn from(v: i32) -> Self {
        UniformValue::I32(v)
    }
}

impl From<u32> for UniformValue {
    fn from(v: u32) -> Self {
        UniformValue::U32(v)
    }
}

impl From<Vec2> for UniformValue {
    fn from(v: Vec2) -> Self {
        UniformValue::Vec2(v)
    }
}

impl From<Vec3> for UniformValue {
    fn from(v: Vec3) -> Self {
        UniformValue::Vec3(v)
    }
}

impl From<Vec4> for UniformValue {
    fn from(v: Vec4) -> Self {
        UniformValue::Vec4(v)
    }
}

/// Collection of custom uniform values.
#[derive(Clone, Debug, Default)]
pub struct CustomUniforms {
    /// Ordered list of (name, value) pairs.
    /// Order matters for WGSL struct layout.
    values: Vec<(String, UniformValue)>,
    /// Quick lookup by name.
    indices: HashMap<String, usize>,
}

impl CustomUniforms {
    /// Create empty custom uniforms.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add or update a uniform value.
    pub fn set<V: Into<UniformValue>>(&mut self, name: &str, value: V) {
        let value = value.into();
        if let Some(&idx) = self.indices.get(name) {
            self.values[idx].1 = value;
        } else {
            let idx = self.values.len();
            self.values.push((name.to_string(), value));
            self.indices.insert(name.to_string(), idx);
        }
    }

    /// Get a uniform value by name.
    pub fn get(&self, name: &str) -> Option<&UniformValue> {
        self.indices.get(name).map(|&idx| &self.values[idx].1)
    }

    /// Check if any custom uniforms are defined.
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    /// Get the number of custom uniforms.
    pub fn len(&self) -> usize {
        self.values.len()
    }

    /// Iterate over all uniforms.
    pub fn iter(&self) -> impl Iterator<Item = (&str, &UniformValue)> {
        self.values.iter().map(|(n, v)| (n.as_str(), v))
    }

    /// Generate WGSL struct fields for custom uniforms.
    pub(crate) fn to_wgsl_fields(&self) -> String {
        self.values
            .iter()
            .map(|(name, value)| format!("    {}: {},", name, value.wgsl_type()))
            .collect::<Vec<_>>()
            .join("\n")
    }

    /// Serialize all values to bytes for GPU upload.
    pub(crate) fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        for (_, value) in &self.values {
            // Add padding for alignment
            let align = match value {
                UniformValue::Vec4(_) | UniformValue::Vec3(_) => 16,
                UniformValue::Vec2(_) => 8,
                _ => 4,
            };
            while buf.len() % align != 0 {
                buf.push(0);
            }
            value.write_bytes(&mut buf);
        }
        buf
    }

    /// Calculate total byte size with alignment.
    pub(crate) fn byte_size(&self) -> usize {
        let bytes = self.to_bytes();
        // Round up to 16-byte alignment for uniform buffer
        (bytes.len() + 15) & !15
    }
}

/// Context passed to the update callback each frame.
///
/// Provides access to input state and allows updating custom uniforms.
///
/// # Input Access
///
/// The context provides full access to input state through the `input` field:
///
/// ```ignore
/// .with_update(|ctx| {
///     // Check if a key was just pressed this frame
///     if ctx.input.key_pressed(KeyCode::Space) {
///         ctx.set("burst", 1.0);
///     }
///
///     // Check if left mouse is held down
///     if ctx.input.mouse_held(MouseButton::Left) {
///         ctx.set("attractor", ctx.mouse_world_pos());
///     }
///
///     // Get mouse movement delta
///     let delta = ctx.input.mouse_delta();
/// })
/// ```
///
/// # CPU Readback
///
/// You can request GPU particle data to be read back to the CPU:
///
/// ```ignore
/// .with_update(|ctx| {
///     // Request readback every second
///     if ctx.time() as u32 != (ctx.time() - ctx.delta_time()) as u32 {
///         ctx.request_readback();
///     }
///
///     // Access previous frame's readback data (if available)
///     if let Some(bytes) = ctx.particles_raw() {
///         let particles: &[MyParticleGpu] = bytemuck::cast_slice(bytes);
///         println!("First particle position: {:?}", particles[0].position);
///     }
/// })
/// ```
pub struct UpdateContext<'a> {
    /// Custom uniforms that can be modified.
    pub(crate) uniforms: &'a mut CustomUniforms,
    /// Input state for the current frame.
    pub input: &'a Input,
    /// Current simulation time in seconds.
    pub(crate) time: f32,
    /// Time since last frame in seconds.
    pub(crate) delta_time: f32,
    /// Simulation bounds (half-size of bounding cube).
    pub(crate) bounds: f32,
    /// Window aspect ratio (width / height).
    pub(crate) aspect_ratio: f32,
    /// Grid opacity to set (None = no change).
    pub(crate) grid_opacity: &'a mut Option<f32>,
    /// Whether to perform readback after this frame.
    pub(crate) readback_requested: &'a mut bool,
    /// Previous frame's readback data (if any).
    pub(crate) readback_data: Option<&'a [u8]>,
}

impl<'a> UpdateContext<'a> {
    /// Create a new update context.
    pub(crate) fn new(
        uniforms: &'a mut CustomUniforms,
        input: &'a Input,
        time: f32,
        delta_time: f32,
        bounds: f32,
        aspect_ratio: f32,
        grid_opacity: &'a mut Option<f32>,
        readback_requested: &'a mut bool,
        readback_data: Option<&'a [u8]>,
    ) -> Self {
        Self {
            uniforms,
            input,
            time,
            delta_time,
            bounds,
            aspect_ratio,
            grid_opacity,
            readback_requested,
            readback_data,
        }
    }

    /// Get the current simulation time in seconds.
    pub fn time(&self) -> f32 {
        self.time
    }

    /// Get the time since last frame in seconds.
    pub fn delta_time(&self) -> f32 {
        self.delta_time
    }

    // ========== Convenience methods that delegate to Input ==========

    /// Get the mouse position in normalized device coordinates (-1 to 1).
    ///
    /// This is a convenience method. For full input access, use `ctx.input`.
    pub fn mouse_ndc(&self) -> Vec2 {
        self.input.mouse_ndc()
    }

    /// Get the mouse position in world coordinates.
    ///
    /// Maps the mouse position from screen space to world space using the
    /// simulation bounds. The z-coordinate is set to 0 (front-facing plane).
    ///
    /// This accounts for window aspect ratio, so the mapping is correct
    /// regardless of window shape.
    ///
    /// # Example
    ///
    /// ```ignore
    /// .with_update(|ctx| {
    ///     if ctx.input.mouse_held(MouseButton::Left) {
    ///         // Get mouse position in world coordinates
    ///         let pos = ctx.mouse_world_pos();
    ///         ctx.set("attractor", pos);
    ///     }
    /// })
    /// ```
    pub fn mouse_world_pos(&self) -> Vec3 {
        let ndc = self.input.mouse_ndc();
        // Scale by bounds and correct for aspect ratio
        // The camera view scales uniformly, so we apply aspect ratio correction
        Vec3::new(
            ndc.x * self.bounds * self.aspect_ratio,
            ndc.y * self.bounds,
            0.0,
        )
    }

    /// Get the simulation bounds (half-size of the bounding cube).
    pub fn bounds(&self) -> f32 {
        self.bounds
    }

    /// Get the window aspect ratio (width / height).
    pub fn aspect_ratio(&self) -> f32 {
        self.aspect_ratio
    }

    /// Check if the left mouse button is currently held down.
    ///
    /// This is a convenience method. For full input access, use `ctx.input`.
    pub fn mouse_pressed(&self) -> bool {
        self.input.mouse_held(MouseButton::Left)
    }

    /// Check if a key was just pressed this frame.
    ///
    /// This is a convenience method. For full input access, use `ctx.input`.
    pub fn key_pressed(&self, key: KeyCode) -> bool {
        self.input.key_pressed(key)
    }

    /// Check if a key is currently held down.
    ///
    /// This is a convenience method. For full input access, use `ctx.input`.
    pub fn key_held(&self, key: KeyCode) -> bool {
        self.input.key_held(key)
    }

    // ========== Uniform methods ==========

    /// Set a custom uniform value.
    pub fn set<V: Into<UniformValue>>(&mut self, name: &str, value: V) {
        self.uniforms.set(name, value);
    }

    /// Get a custom uniform value.
    pub fn get(&self, name: &str) -> Option<&UniformValue> {
        self.uniforms.get(name)
    }

    // ========== Visual methods ==========

    /// Set the spatial grid visualization opacity.
    ///
    /// Use 0.0 to hide the grid, 1.0 for full visibility.
    /// The change takes effect on the next frame.
    pub fn set_grid_opacity(&mut self, opacity: f32) {
        *self.grid_opacity = Some(opacity.clamp(0.0, 1.0));
    }

    // ========== CPU Readback methods ==========

    /// Request particle data to be read back from GPU after this frame.
    ///
    /// The data will be available via `particles_raw()` on the next frame.
    /// This is an expensive operation that stalls the GPU pipeline.
    /// Use sparingly (e.g., once per second, or on user request).
    ///
    /// # Example
    ///
    /// ```ignore
    /// .with_update(|ctx| {
    ///     // Request readback when space is pressed
    ///     if ctx.key_pressed(KeyCode::Space) {
    ///         ctx.request_readback();
    ///     }
    /// })
    /// ```
    pub fn request_readback(&mut self) {
        *self.readback_requested = true;
    }

    /// Get the raw bytes of particle data from the previous readback request.
    ///
    /// Returns `None` if no readback was requested on the previous frame.
    /// The bytes can be cast to your particle's GPU type using `bytemuck::cast_slice`.
    ///
    /// **Note:** This returns a reference, so you cannot call `ctx.set()` while
    /// holding the returned slice. Use `with_particles()` for the common pattern
    /// of reading data and then updating uniforms.
    ///
    /// # Example
    ///
    /// ```ignore
    /// .with_update(|ctx| {
    ///     if let Some(bytes) = ctx.particles_raw() {
    ///         let particles: &[MyParticleGpu] = bytemuck::cast_slice(bytes);
    ///         for p in particles.iter().take(10) {
    ///             println!("Position: {:?}", p.position);
    ///         }
    ///     }
    /// })
    /// ```
    pub fn particles_raw(&self) -> Option<&[u8]> {
        self.readback_data
    }

    /// Process particle data and return a result, handling borrow scope automatically.
    ///
    /// This is the recommended way to read particle data when you also need to
    /// update uniforms based on the results. The closure receives the raw bytes
    /// and returns any value you need; after the closure completes, you can
    /// freely call `ctx.set()` with the results.
    ///
    /// Returns `None` if no readback data is available.
    ///
    /// # Example
    ///
    /// ```ignore
    /// .with_update(|ctx| {
    ///     // Compute stats from particles
    ///     if let Some((center, count)) = ctx.with_particles(|bytes| {
    ///         let particles: &[MyParticleGpu] = bytemuck::cast_slice(bytes);
    ///         let sum: Vec3 = particles.iter()
    ///             .map(|p| Vec3::from_array(p.position))
    ///             .sum();
    ///         let count = particles.len();
    ///         (sum / count as f32, count)
    ///     }) {
    ///         // Now we can update uniforms with the results
    ///         ctx.set("center_x", center.x);
    ///         ctx.set("center_y", center.y);
    ///         ctx.set("particle_count", count as f32);
    ///     }
    /// })
    /// ```
    pub fn with_particles<T, F>(&self, f: F) -> Option<T>
    where
        F: FnOnce(&[u8]) -> T,
    {
        self.readback_data.map(f)
    }
}