ballin 0.1.2

//! Physics world wrapper around rapier2d.
//!
//! Encapsulates all rapier2d components needed for the simulation,
//! providing a clean, high-level interface for the application layer.

use rand::Rng;
use rapier2d::prelude::*;
use ratatui::style::Color;

use super::config::PhysicsConfig;
use crate::error::AppResult;

/// Ball color for Color Mode rendering.
///
/// Balls start as White (default) and can be converted to one of
/// the 6 geyser colors when hit by a geyser burst in Color Mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum BallColor {
    /// Default white color (no geyser interaction).
    #[default]
    White,
    /// Red (geyser 1).
    Red,
    /// Green (geyser 2).
    Green,
    /// Yellow (geyser 3).
    Yellow,
    /// Blue (geyser 4).
    Blue,
    /// Magenta (geyser 5).
    Magenta,
    /// Cyan (geyser 6).
    Cyan,
}

impl BallColor {
    /// Converts a geyser digit (1-6) to the corresponding ball color.
    ///
    /// # Arguments
    ///
    /// * `digit` - Geyser digit (1-6)
    ///
    /// # Returns
    ///
    /// The ball color for that geyser, or White if digit is invalid.
    pub fn from_geyser(digit: u8) -> Self {
        match digit {
            1 => BallColor::Red,
            2 => BallColor::Green,
            3 => BallColor::Yellow,
            4 => BallColor::Blue,
            5 => BallColor::Magenta,
            6 => BallColor::Cyan,
            _ => BallColor::White,
        }
    }

    /// Converts the ball color to a ratatui Color for rendering.
    pub fn to_ratatui_color(self) -> Color {
        match self {
            BallColor::White => Color::White,
            BallColor::Red => Color::Red,
            BallColor::Green => Color::Green,
            BallColor::Yellow => Color::Yellow,
            BallColor::Blue => Color::Blue,
            BallColor::Magenta => Color::Magenta,
            BallColor::Cyan => Color::Cyan,
        }
    }

    /// Returns a numeric index for color comparison/counting.
    /// Used to efficiently track the most common color in a cell.
    pub fn index(self) -> usize {
        match self {
            BallColor::White => 0,
            BallColor::Red => 1,
            BallColor::Green => 2,
            BallColor::Yellow => 3,
            BallColor::Blue => 4,
            BallColor::Magenta => 5,
            BallColor::Cyan => 6,
        }
    }

    /// Converts a color index back to a BallColor.
    pub fn from_index(index: usize) -> Self {
        match index {
            1 => BallColor::Red,
            2 => BallColor::Green,
            3 => BallColor::Yellow,
            4 => BallColor::Blue,
            5 => BallColor::Magenta,
            6 => BallColor::Cyan,
            _ => BallColor::White,
        }
    }

    /// Returns a random non-white ball color.
    ///
    /// Randomly selects one of the 6 geyser colors (Red, Green, Yellow, Blue, Magenta, Cyan).
    pub fn random_color() -> Self {
        let mut rng = rand::thread_rng();
        match rng.gen_range(1..=6) {
            1 => BallColor::Red,
            2 => BallColor::Green,
            3 => BallColor::Yellow,
            4 => BallColor::Blue,
            5 => BallColor::Magenta,
            _ => BallColor::Cyan,
        }
    }
}

/// Wrapper around rapier2d physics components.
///
/// Manages the complete physics simulation state including:
/// - Rigid bodies (balls)
/// - Colliders (ball shapes and walls)
/// - Physics pipeline (solver, broad/narrow phase)
///
/// This struct owns all rapier2d state and provides methods for
/// spawning balls, applying forces, and querying positions.
pub struct PhysicsWorld {
    // Core rapier2d component sets
    rigid_body_set: RigidBodySet,
    collider_set: ColliderSet,

    // Simulation infrastructure
    physics_pipeline: PhysicsPipeline,
    island_manager: IslandManager,
    broad_phase: DefaultBroadPhase,
    narrow_phase: NarrowPhase,
    impulse_joint_set: ImpulseJointSet,
    multibody_joint_set: MultibodyJointSet,
    ccd_solver: CCDSolver,

    // Simulation parameters
    integration_parameters: IntegrationParameters,
    gravity: Vector<Real>,

    // Ball tracking: maps rigid body handles to their collider handles and colors
    // Pre-allocated with expected capacity to avoid reallocation
    // Tuple: (body_handle, collider_handle, ball_color)
    ball_handles: Vec<(RigidBodyHandle, ColliderHandle, BallColor)>,

    // Boundary collider handles (walls: bottom, top, left, right)
    boundary_handles: Vec<ColliderHandle>,

    // Configuration reference for creating new balls
    config: PhysicsConfig,

    // World bounds in physics coordinates
    world_width: Real,
    world_height: Real,
}

impl PhysicsWorld {
    /// Creates a new physics world with the given dimensions.
    ///
    /// Initializes all rapier2d components and creates boundary walls.
    /// The world uses Y-up coordinate convention (0,0 is bottom-left).
    ///
    /// # Arguments
    ///
    /// * `width` - World width in physics units
    /// * `height` - World height in physics units
    /// * `config` - Physics configuration (gravity, friction, etc.)
    ///
    /// # Returns
    ///
    /// A new `PhysicsWorld` ready for simulation.
    pub fn new(width: Real, height: Real, config: PhysicsConfig) -> Self {
        let rigid_body_set = RigidBodySet::new();
        let collider_set = ColliderSet::new();

        let physics_pipeline = PhysicsPipeline::new();
        let island_manager = IslandManager::new();
        let broad_phase = DefaultBroadPhase::new();
        let narrow_phase = NarrowPhase::new();
        let impulse_joint_set = ImpulseJointSet::new();
        let multibody_joint_set = MultibodyJointSet::new();
        let ccd_solver = CCDSolver::new();

        // Configure integration for stable simulation at 60 FPS
        // Tuned for stability with high ball counts (5000+):
        // - num_solver_iterations: Reduced from 4 to 2 for better throughput
        // - num_internal_pgs_iterations: Limited to 1 to prevent solver explosion
        let integration_parameters = IntegrationParameters {
            num_solver_iterations: 2,
            num_internal_pgs_iterations: 1,
            ..Default::default()
        };

        let gravity = config.gravity;

        let mut world = Self {
            rigid_body_set,
            collider_set,
            physics_pipeline,
            island_manager,
            broad_phase,
            narrow_phase,
            impulse_joint_set,
            multibody_joint_set,
            ccd_solver,
            integration_parameters,
            gravity,
            // Pre-allocate for 10,000 balls to avoid reallocation
            ball_handles: Vec::with_capacity(10_000),
            // 4 walls + 4 corner colliders
            boundary_handles: Vec::with_capacity(8),
            config,
            world_width: width,
            world_height: height,
        };

        // Create boundary walls
        world.create_boundaries();
        world
    }

    /// Creates boundary walls around the physics world.
    ///
    /// Four walls are created as static colliders:
    /// - Bottom: y = 0
    /// - Top: y = world_height
    /// - Left: x = 0
    /// - Right: x = world_width
    fn create_boundaries(&mut self) {
        // Clear existing boundaries if any
        for handle in self.boundary_handles.drain(..) {
            self.collider_set.remove(
                handle,
                &mut self.island_manager,
                &mut self.rigid_body_set,
                true,
            );
        }

        let half_width = self.world_width / 2.0;
        let half_height = self.world_height / 2.0;
        let wall_thickness = 1.0;

        // Bottom wall (slightly below y=0 so balls rest at y=ball_radius)
        let bottom = ColliderBuilder::cuboid(half_width + wall_thickness, wall_thickness)
            .translation(vector![half_width, -wall_thickness])
            .friction(self.config.friction)
            .restitution(self.config.restitution)
            .build();

        // Top wall
        let top = ColliderBuilder::cuboid(half_width + wall_thickness, wall_thickness)
            .translation(vector![half_width, self.world_height + wall_thickness])
            .friction(self.config.friction)
            .restitution(self.config.restitution)
            .build();

        // Left wall
        let left = ColliderBuilder::cuboid(wall_thickness, half_height + wall_thickness)
            .translation(vector![-wall_thickness, half_height])
            .friction(self.config.friction)
            .restitution(self.config.restitution)
            .build();

        // Right wall
        let right = ColliderBuilder::cuboid(wall_thickness, half_height + wall_thickness)
            .translation(vector![self.world_width + wall_thickness, half_height])
            .friction(self.config.friction)
            .restitution(self.config.restitution)
            .build();

        self.boundary_handles.push(self.collider_set.insert(bottom));
        self.boundary_handles.push(self.collider_set.insert(top));
        self.boundary_handles.push(self.collider_set.insert(left));
        self.boundary_handles.push(self.collider_set.insert(right));

        // Corner colliders to prevent balls getting stuck in corners
        // Using larger ball-shaped colliders (5x ball size) to effectively
        // round off the corners and deflect balls before they can get trapped
        let corner_radius = self.config.ball_radius * 5.0;
        let corners = [
            (0.0, 0.0),                            // bottom-left
            (self.world_width, 0.0),               // bottom-right
            (0.0, self.world_height),              // top-left
            (self.world_width, self.world_height), // top-right
        ];
        for (cx, cy) in corners {
            let corner = ColliderBuilder::ball(corner_radius)
                .translation(vector![cx, cy])
                .friction(self.config.friction)
                .restitution(self.config.restitution)
                .build();
            self.boundary_handles.push(self.collider_set.insert(corner));
        }
    }

    /// Steps the physics simulation forward by one timestep.
    ///
    /// Uses rapier2d's parallel collision detection when the `parallel`
    /// feature is enabled. The timestep is determined by integration_parameters.
    /// After stepping, ball velocities are clamped to prevent infinite speeds.
    pub fn step(&mut self) {
        self.physics_pipeline.step(
            &self.gravity,
            &self.integration_parameters,
            &mut self.island_manager,
            &mut self.broad_phase,
            &mut self.narrow_phase,
            &mut self.rigid_body_set,
            &mut self.collider_set,
            &mut self.impulse_joint_set,
            &mut self.multibody_joint_set,
            &mut self.ccd_solver,
            &(),
            &(),
        );

        self.clamp_velocities();
    }

    /// Clamps all ball velocities to the configured maximum.
    ///
    /// This prevents balls from achieving unreasonably high speeds
    /// due to physics edge cases or accumulated impulses.
    fn clamp_velocities(&mut self) {
        let max_vel = self.config.max_velocity;
        let max_vel_sq = max_vel * max_vel;

        for (body_handle, _, _) in &self.ball_handles {
            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let vel = body.linvel();
                let vel_sq = vel.x * vel.x + vel.y * vel.y;

                if vel_sq > max_vel_sq {
                    // Scale velocity down to max magnitude
                    let scale = max_vel / vel_sq.sqrt();
                    body.set_linvel(vector![vel.x * scale, vel.y * scale], true);
                }
            }
        }
    }

    /// Spawns a new ball at the given physics coordinates.
    ///
    /// Creates a dynamic rigid body with a ball collider attached.
    /// The ball inherits physics properties from the current config.
    ///
    /// # Arguments
    ///
    /// * `x` - X position in physics units
    /// * `y` - Y position in physics units
    ///
    /// # Returns
    ///
    /// The rigid body handle for the new ball.
    ///
    /// # Errors
    ///
    /// Returns `AppError::Physics` if ball creation fails.
    pub fn spawn_ball(&mut self, x: Real, y: Real) -> AppResult<RigidBodyHandle> {
        self.spawn_ball_with_velocity(x, y, 0.0, 0.0)
    }

    /// Spawns a new ball with initial velocity.
    ///
    /// # Arguments
    ///
    /// * `x` - X position in physics units
    /// * `y` - Y position in physics units
    /// * `vx` - Initial X velocity
    /// * `vy` - Initial Y velocity
    ///
    /// # Returns
    ///
    /// The rigid body handle for the new ball.
    pub fn spawn_ball_with_velocity(
        &mut self,
        x: Real,
        y: Real,
        vx: Real,
        vy: Real,
    ) -> AppResult<RigidBodyHandle> {
        // Clamp position to world bounds
        let x = x.clamp(
            self.config.ball_radius,
            self.world_width - self.config.ball_radius,
        );
        let y = y.clamp(
            self.config.ball_radius,
            self.world_height - self.config.ball_radius,
        );

        // Create dynamic rigid body
        let rigid_body = RigidBodyBuilder::dynamic()
            .translation(vector![x, y])
            .linvel(vector![vx, vy])
            .linear_damping(self.config.linear_damping)
            .build();

        let body_handle = self.rigid_body_set.insert(rigid_body);

        let collider = ColliderBuilder::ball(self.config.ball_radius)
            .restitution(self.config.restitution)
            .friction(self.config.friction)
            .density(self.config.density)
            .build();

        let collider_handle =
            self.collider_set
                .insert_with_parent(collider, body_handle, &mut self.rigid_body_set);

        self.ball_handles
            .push((body_handle, collider_handle, BallColor::White));

        Ok(body_handle)
    }

    /// Spawns a new ball with initial velocity and a specific color.
    ///
    /// # Arguments
    ///
    /// * `x` - X position in physics units
    /// * `y` - Y position in physics units
    /// * `vx` - Initial X velocity
    /// * `vy` - Initial Y velocity
    /// * `color` - The ball's color
    ///
    /// # Returns
    ///
    /// The rigid body handle for the new ball.
    pub fn spawn_ball_with_velocity_and_color(
        &mut self,
        x: Real,
        y: Real,
        vx: Real,
        vy: Real,
        color: BallColor,
    ) -> AppResult<RigidBodyHandle> {
        // Clamp position to world bounds
        let x = x.clamp(
            self.config.ball_radius,
            self.world_width - self.config.ball_radius,
        );
        let y = y.clamp(
            self.config.ball_radius,
            self.world_height - self.config.ball_radius,
        );

        // Create dynamic rigid body
        let rigid_body = RigidBodyBuilder::dynamic()
            .translation(vector![x, y])
            .linvel(vector![vx, vy])
            .linear_damping(self.config.linear_damping)
            .build();

        let body_handle = self.rigid_body_set.insert(rigid_body);

        let collider = ColliderBuilder::ball(self.config.ball_radius)
            .restitution(self.config.restitution)
            .friction(self.config.friction)
            .density(self.config.density)
            .build();

        let collider_handle =
            self.collider_set
                .insert_with_parent(collider, body_handle, &mut self.rigid_body_set);

        self.ball_handles
            .push((body_handle, collider_handle, color));

        Ok(body_handle)
    }

    /// Applies a radial impulse (burst) at the given position.
    ///
    /// All balls within the burst radius receive an outward impulse.
    /// Force magnitude decreases with distance (inverse square falloff).
    /// For stability with high ball counts, limits the number of affected balls
    /// and caps impulse magnitude.
    ///
    /// # Arguments
    ///
    /// * `center_x` - X coordinate of burst center
    /// * `center_y` - Y coordinate of burst center
    /// * `strength` - Base impulse magnitude
    /// * `radius` - Maximum effect radius
    pub fn apply_burst(&mut self, center_x: Real, center_y: Real, strength: Real, radius: Real) {
        let center: Vector<Real> = vector![center_x, center_y];
        let radius_sq = radius * radius;

        // Stability limits for high ball counts
        const MAX_AFFECTED_BALLS: usize = 500;
        const MAX_IMPULSE_MAGNITUDE: Real = 10.0;

        let mut affected_count = 0;

        for (body_handle, _, _) in &self.ball_handles {
            // Limit affected balls to prevent physics solver overflow
            if affected_count >= MAX_AFFECTED_BALLS {
                break;
            }

            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let pos = body.translation();
                let diff = pos - center;
                let dist_sq = diff.norm_squared();

                // Only affect balls within radius and not at center
                if dist_sq < radius_sq && dist_sq > 0.01 {
                    let dist = dist_sq.sqrt();
                    let direction = diff / dist;
                    // Inverse square falloff, clamped to prevent extreme forces
                    let force_magnitude = (strength / dist_sq.max(1.0)).min(MAX_IMPULSE_MAGNITUDE);
                    let impulse = direction * force_magnitude;

                    body.apply_impulse(impulse, true);
                    affected_count += 1;
                }
            }
        }
    }

    /// Applies a directional burst (biased in a specific direction).
    ///
    /// Similar to `apply_burst()` but biases the impulse toward a given direction.
    /// Used for number key bursts that push balls upward.
    /// For stability with high ball counts, limits the number of affected balls
    /// and caps impulse magnitude.
    ///
    /// # Arguments
    ///
    /// * `center_x` - X coordinate of burst center
    /// * `center_y` - Y coordinate of burst center
    /// * `dir_x` - X component of bias direction (normalized)
    /// * `dir_y` - Y component of bias direction (normalized)
    /// * `strength` - Base impulse magnitude
    /// * `radius` - Maximum effect radius
    pub fn apply_directional_burst(
        &mut self,
        center_x: Real,
        center_y: Real,
        dir_x: Real,
        dir_y: Real,
        strength: Real,
        radius: Real,
    ) {
        let center: Vector<Real> = vector![center_x, center_y];
        let bias: Vector<Real> = vector![dir_x, dir_y];
        let radius_sq = radius * radius;

        // Stability limits for high ball counts
        const MAX_AFFECTED_BALLS: usize = 500;
        const MAX_IMPULSE_MAGNITUDE: Real = 10.0;

        let mut affected_count = 0;

        for (body_handle, _, _) in &self.ball_handles {
            // Limit affected balls to prevent physics solver overflow
            if affected_count >= MAX_AFFECTED_BALLS {
                break;
            }

            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let pos = body.translation();
                let diff = pos - center;
                let dist_sq = diff.norm_squared();

                // Only affect balls within radius
                if dist_sq < radius_sq && dist_sq > 0.01 {
                    let dist = dist_sq.sqrt();
                    let outward = diff / dist;
                    // Blend outward direction with bias direction (50/50)
                    let direction = (outward + bias).normalize();
                    // Inverse square falloff, clamped
                    let force_magnitude = (strength / dist_sq.max(1.0)).min(MAX_IMPULSE_MAGNITUDE);
                    let impulse = direction * force_magnitude;

                    body.apply_impulse(impulse, true);
                    affected_count += 1;
                }
            }
        }
    }

    /// Applies a directional burst that also colors affected balls.
    ///
    /// Similar to `apply_directional_burst()` but additionally changes the color
    /// of all balls directly impacted by the burst. Used for Color Mode.
    ///
    /// # Arguments
    ///
    /// * `center_x` - X coordinate of burst center
    /// * `center_y` - Y coordinate of burst center
    /// * `dir_x` - X component of bias direction (normalized)
    /// * `dir_y` - Y component of bias direction (normalized)
    /// * `strength` - Base impulse magnitude
    /// * `radius` - Maximum effect radius
    /// * `color` - The color to apply to affected balls
    #[allow(clippy::too_many_arguments)]
    pub fn apply_directional_burst_with_color(
        &mut self,
        center_x: Real,
        center_y: Real,
        dir_x: Real,
        dir_y: Real,
        strength: Real,
        radius: Real,
        color: BallColor,
    ) {
        let center: Vector<Real> = vector![center_x, center_y];
        let bias: Vector<Real> = vector![dir_x, dir_y];
        let radius_sq = radius * radius;

        // Stability limits for high ball counts
        const MAX_AFFECTED_BALLS: usize = 500;
        const MAX_IMPULSE_MAGNITUDE: Real = 10.0;

        let mut affected_count = 0;

        for (body_handle, _, ball_color) in &mut self.ball_handles {
            // Limit affected balls to prevent physics solver overflow
            if affected_count >= MAX_AFFECTED_BALLS {
                break;
            }

            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let pos = body.translation();
                let diff = pos - center;
                let dist_sq = diff.norm_squared();

                // Only affect balls within radius
                if dist_sq < radius_sq && dist_sq > 0.01 {
                    let dist = dist_sq.sqrt();
                    let outward = diff / dist;
                    // Blend outward direction with bias direction (50/50)
                    let direction = (outward + bias).normalize();
                    // Inverse square falloff, clamped
                    let force_magnitude = (strength / dist_sq.max(1.0)).min(MAX_IMPULSE_MAGNITUDE);
                    let impulse = direction * force_magnitude;

                    body.apply_impulse(impulse, true);
                    *ball_color = color;
                    affected_count += 1;
                }
            }
        }
    }

    /// Applies forces to balls affected by boundary shrinkage.
    ///
    /// When the window shrinks, only balls near the moving boundary are pushed.
    /// Force magnitude is proportional to resize velocity and proximity to the edge.
    ///
    /// # Arguments
    ///
    /// * `old_width` - Previous world width
    /// * `old_height` - Previous world height
    /// * `new_width` - New world width
    /// * `new_height` - New world height
    /// * `velocity_x` - Horizontal resize velocity (negative = shrinking)
    /// * `velocity_y` - Vertical resize velocity (negative = shrinking from top in physics coords)
    pub fn apply_boundary_shrink_force(
        &mut self,
        old_width: Real,
        old_height: Real,
        new_width: Real,
        new_height: Real,
        velocity_x: Real,
        velocity_y: Real,
    ) {
        // Only apply forces when shrinking
        let shrinking_width = new_width < old_width;
        let shrinking_height = new_height < old_height;

        if !shrinking_width && !shrinking_height {
            return;
        }

        // Zone depth: how far from the boundary balls are affected
        // Balls within this distance from the shrinking edge receive force
        let ball_radius = self.config.ball_radius;
        let zone_depth = ball_radius * 8.0;

        // Force scaling factor
        const FORCE_SCALE: Real = 0.5;
        // Maximum impulse to prevent instability
        const MAX_IMPULSE: Real = 8.0;

        for (body_handle, _, _) in &self.ball_handles {
            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let pos = *body.translation();
                let mut impulse_x: Real = 0.0;
                let mut impulse_y: Real = 0.0;

                // Width shrinking: left wall moves right, push nearby balls right
                // When terminal shrinks from left side, it's like the left wall advancing
                if shrinking_width {
                    let dist_from_left = pos.x;
                    if dist_from_left < zone_depth && dist_from_left >= 0.0 {
                        // Force is stronger for balls closer to the edge
                        // Uses inverse relationship: closer = stronger push
                        let proximity = 1.0 - (dist_from_left / zone_depth);
                        // Force pushes right (positive X), proportional to shrink velocity
                        impulse_x = velocity_x.abs() * proximity * FORCE_SCALE;
                    }
                }

                // Height shrinking: floor rises, push balls near bottom upward
                // In physics coords (Y-up), floor is at y=0
                // When terminal shrinks from bottom, it's like the floor rising
                if shrinking_height {
                    let dist_from_bottom = pos.y;
                    if dist_from_bottom < zone_depth && dist_from_bottom >= 0.0 {
                        let proximity = 1.0 - (dist_from_bottom / zone_depth);
                        // Force pushes up (positive Y in physics coords)
                        impulse_y = velocity_y.abs() * proximity * FORCE_SCALE;
                    }
                }

                // Apply impulse if any force was calculated
                if impulse_x != 0.0 || impulse_y != 0.0 {
                    let impulse = vector![impulse_x, impulse_y];
                    // Clamp impulse magnitude for stability
                    let magnitude = impulse.norm();
                    let clamped = if magnitude > MAX_IMPULSE {
                        impulse * (MAX_IMPULSE / magnitude)
                    } else {
                        impulse
                    };
                    body.apply_impulse(clamped, true);
                }
            }
        }
    }

    /// Applies a nudge impulse to all balls (pinball-style).
    ///
    /// Used for arrow key inputs to push all balls in a direction.
    /// The impulse is applied uniformly to all balls.
    ///
    /// # Arguments
    ///
    /// * `dx` - Horizontal impulse component
    /// * `dy` - Vertical impulse component
    pub fn nudge_all(&mut self, dx: Real, dy: Real) {
        let impulse = vector![dx, dy];

        for (body_handle, _, _) in &self.ball_handles {
            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                body.apply_impulse(impulse, true);
            }
        }
    }

    /// Updates world boundaries to match new dimensions.
    ///
    /// Called when the terminal is resized. Recreates boundary
    /// colliders to match the new world size.
    ///
    /// # Arguments
    ///
    /// * `width` - New world width in physics units
    /// * `height` - New world height in physics units
    pub fn update_boundaries(&mut self, width: Real, height: Real) {
        self.world_width = width;
        self.world_height = height;
        self.create_boundaries();
    }

    /// Returns an iterator over all ball positions.
    ///
    /// Positions are in physics coordinates (Y-up convention).
    /// The iterator yields `(x, y)` tuples for each active ball.
    ///
    /// # Returns
    ///
    /// Iterator over ball positions as `(f32, f32)` tuples.
    pub fn ball_positions(&self) -> impl Iterator<Item = (Real, Real)> + '_ {
        self.ball_handles.iter().filter_map(|(body_handle, _, _)| {
            self.rigid_body_set.get(*body_handle).map(|body| {
                let pos = body.translation();
                (pos.x, pos.y)
            })
        })
    }

    /// Returns an iterator over all ball positions with their colors.
    ///
    /// Positions are in physics coordinates (Y-up convention).
    /// The iterator yields `(x, y, color)` tuples for each active ball.
    ///
    /// # Returns
    ///
    /// Iterator over ball positions and colors as `(f32, f32, BallColor)` tuples.
    pub fn ball_positions_with_colors(&self) -> impl Iterator<Item = (Real, Real, BallColor)> + '_ {
        self.ball_handles
            .iter()
            .filter_map(|(body_handle, _, color)| {
                self.rigid_body_set.get(*body_handle).map(|body| {
                    let pos = body.translation();
                    (pos.x, pos.y, *color)
                })
            })
    }

    /// Returns the current number of balls in the simulation.
    pub fn ball_count(&self) -> usize {
        self.ball_handles.len()
    }

    /// Updates physics configuration at runtime.
    ///
    /// Changes to gravity take effect immediately. Other properties
    /// (friction, restitution) only affect newly created balls.
    ///
    /// # Arguments
    ///
    /// * `config` - New physics configuration
    pub fn update_config(&mut self, config: PhysicsConfig) {
        self.gravity = config.gravity;
        self.config = config;
    }

    /// Removes all balls from the simulation.
    ///
    /// Used when resetting the simulation. Boundary walls are preserved.
    pub fn clear_balls(&mut self) {
        for (body_handle, collider_handle, _) in self.ball_handles.drain(..) {
            self.collider_set.remove(
                collider_handle,
                &mut self.island_manager,
                &mut self.rigid_body_set,
                true,
            );
            self.rigid_body_set.remove(
                body_handle,
                &mut self.island_manager,
                &mut self.collider_set,
                &mut self.impulse_joint_set,
                &mut self.multibody_joint_set,
                true,
            );
        }
    }

    /// Returns the world dimensions.
    ///
    /// # Returns
    ///
    /// `(width, height)` in physics units.
    pub fn dimensions(&self) -> (Real, Real) {
        (self.world_width, self.world_height)
    }

    /// Returns a reference to the current physics configuration.
    pub fn config(&self) -> &PhysicsConfig {
        &self.config
    }

    /// Returns a mutable reference to the rigid body set.
    ///
    /// Used for shape collider creation and manipulation.
    pub fn rigid_body_set_mut(&mut self) -> &mut RigidBodySet {
        &mut self.rigid_body_set
    }

    /// Returns a mutable reference to the collider set.
    ///
    /// Used for shape collider creation and manipulation.
    pub fn collider_set_mut(&mut self) -> &mut ColliderSet {
        &mut self.collider_set
    }

    /// Returns a reference to the collider set.
    ///
    /// Used for shape hit testing.
    pub fn collider_set(&self) -> &ColliderSet {
        &self.collider_set
    }

    /// Returns a mutable reference to the island manager.
    pub fn island_manager_mut(&mut self) -> &mut IslandManager {
        &mut self.island_manager
    }

    /// Returns a mutable reference to the impulse joint set.
    pub fn impulse_joint_set_mut(&mut self) -> &mut ImpulseJointSet {
        &mut self.impulse_joint_set
    }

    /// Returns a mutable reference to the multibody joint set.
    pub fn multibody_joint_set_mut(&mut self) -> &mut MultibodyJointSet {
        &mut self.multibody_joint_set
    }

    /// Returns mutable references to all physics components needed for shape operations.
    ///
    /// This is necessary because Rust's borrow checker prevents multiple mutable
    /// borrows of different struct fields through separate method calls.
    ///
    /// # Returns
    ///
    /// A tuple containing mutable references to:
    /// - RigidBodySet
    /// - ColliderSet
    /// - IslandManager
    /// - ImpulseJointSet
    /// - MultibodyJointSet
    #[allow(clippy::type_complexity)]
    pub fn shape_components_mut(
        &mut self,
    ) -> (
        &mut RigidBodySet,
        &mut ColliderSet,
        &mut IslandManager,
        &mut ImpulseJointSet,
        &mut MultibodyJointSet,
    ) {
        (
            &mut self.rigid_body_set,
            &mut self.collider_set,
            &mut self.island_manager,
            &mut self.impulse_joint_set,
            &mut self.multibody_joint_set,
        )
    }

    /// Pushes balls away from a shape's position when it is placed.
    ///
    /// This prevents balls from becoming trapped inside shapes or in sealed
    /// areas formed by multiple shapes. Balls within the displacement radius
    /// receive an outward impulse proportional to their proximity to the center.
    ///
    /// # Arguments
    ///
    /// * `center_x` - X coordinate of the shape center
    /// * `center_y` - Y coordinate of the shape center
    /// * `radius` - Radius around the center where balls should be displaced
    pub fn displace_balls_from_shape(&mut self, center_x: Real, center_y: Real, radius: Real) {
        let center: Vector<Real> = vector![center_x, center_y];
        let radius_sq = radius * radius;

        // Strong impulse to ensure balls escape the shape area
        const DISPLACEMENT_STRENGTH: Real = 15.0;

        for (body_handle, _, _) in &self.ball_handles {
            if let Some(body) = self.rigid_body_set.get_mut(*body_handle) {
                let pos = body.translation();
                let diff = pos - center;
                let dist_sq = diff.norm_squared();

                // Affect balls within the displacement radius
                if dist_sq < radius_sq {
                    let dist = dist_sq.sqrt().max(0.1); // Prevent division by zero
                    let direction = diff / dist;

                    // Stronger impulse for balls closer to center
                    // This ensures balls at the shape's center get pushed out
                    let proximity = 1.0 - (dist / radius);
                    let impulse_magnitude = DISPLACEMENT_STRENGTH * (proximity + 0.5);
                    let impulse = direction * impulse_magnitude;

                    body.apply_impulse(impulse, true);
                }
            }
        }
    }
}