📄 src/bunnies/bunny.rs
use std::{ops::DerefMut, time::Duration};

use bevy::{
    app::{App, Startup, Update},
    asset::{AssetServer, Assets},
    ecs::{
        component::Component,
        entity::Entity,
        query::{With, Without},
        system::{Commands, Query, Res, ResMut, Single},
    },
    gltf::GltfAssetLabel,
    math::{Dir3, Vec3, primitives::Cuboid},
    mesh::{Mesh, Mesh3d, MeshBuilder, Meshable},
    scene::SceneRoot,
    time::Time,
    transform::components::Transform,
};

use rand::RngExt;

use crate::{Rng, bunnies::locator::Locator, dog::Dog, obstacles::Obstacles};

pub(super) fn add_bunny_systems(app: &mut App) -> &mut App {
    app.add_systems(Startup, setup)
        .add_systems(Update, (calculate, jump))
}

#[derive(Component)]
pub(super) struct Bunny;

#[derive(Component)]
enum JumpState {
    Jumping { from: Vec3, to: Vec3 },
    Cooldown { ready: Duration },
}

const BUNNY_COUNT: usize = 5;
const DETECTION_DISTANCE: f32 = 5.0;
const SHORT_JUMP_DISTANCE: f32 = 1.0;
const LONG_JUMP_DISTANCE: f32 = 3.0;
const LONE_JUMPS_PER_SECOND: f64 = 1.0;
const COLLECTIVE_JUMPS_PER_SECOND: f64 = 1.0;
const JUMP_HEIGHT: f32 = 1.0;
const JUMP_COOLDOWN: f32 = 0.05;
const DESIRED_SPEED: f32 = 10.0;
const SPEED: f32 = LONG_JUMP_DISTANCE / (LONG_JUMP_DISTANCE / DESIRED_SPEED - JUMP_COOLDOWN);

fn setup(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut rng: ResMut<Rng>,
) {
    let model =
        asset_server.load(GltfAssetLabel::Scene(0).from_asset("models/cube-pets/animal-bunny.glb"));
    let mesh = meshes.add(Cuboid::from_length(2.0).mesh().build());

    for _ in 0..BUNNY_COUNT {
        commands.spawn((
            SceneRoot(model.clone()),
            Transform::from_xyz(
                rng.random_range(-15.0..=15.0),
                0.0,
                rng.random_range(-15.0..=15.0),
            ),
            Mesh3d(mesh.clone()),
            Bunny,
        ));
    }
}

type IdleBunniesQuery<'w, 's, 't> =
    Query<'w, 's, (Entity, &'t mut Transform), (With<Bunny>, Without<JumpState>)>;

fn calculate(
    mut commands: Commands,
    time: Res<Time>,
    locator: Res<Locator>,
    dog: Single<&Transform, (With<Dog>, Without<Bunny>)>,
    mut bunnies: IdleBunniesQuery,
    mut obstacles: Obstacles,
    mut rng: ResMut<Rng>,
) {
    bunnies.iter_mut().for_each(|(bunny, mut bunny_transform)| {
        calculate_next_move(
            &mut commands,
            &time,
            &locator,
            &dog,
            &mut bunny_transform,
            bunny,
            &mut obstacles,
            &mut rng,
        );
    });
}

fn calculate_next_move(
    commands: &mut Commands,
    time: &Time,
    locator: &Locator,
    dog: &Transform,
    bunny_transform: &mut Transform,
    bunny: Entity,
    obstacles: &mut Obstacles,
    rng: &mut Rng,
) {
    if let Some(direction) = is_dog_nearby(dog, bunny_transform) {
        jump_away_from_dog(commands, bunny_transform, bunny, obstacles, direction);
    } else if let Some(nearby_bunnies) = is_near_others(locator, bunny, bunny_transform) {
        jump_with_others(
            commands,
            time,
            obstacles,
            bunny,
            bunny_transform,
            nearby_bunnies,
            rng,
        );
    } else {
        jump_alone(commands, time, obstacles, bunny, bunny_transform, rng);
    }
}

fn is_dog_nearby(dog: &Transform, bunny_transform: &mut Transform) -> Option<Vec3> {
    let direction = bunny_transform.translation - dog.translation;
    if direction.length_squared() > DETECTION_DISTANCE * DETECTION_DISTANCE {
        return None;
    }
    Some(direction)
}

fn jump_away_from_dog(
    commands: &mut Commands,
    bunny_transform: &mut Transform,
    bunny: Entity,
    obstacles: &mut Obstacles,
    direction: Vec3,
) {
    let dog_direction = direction.with_y(0.0).normalize();
    let Some(to) = obstacles.avoid(
        bunny_transform.translation,
        bunny_transform.translation + dog_direction * LONG_JUMP_DISTANCE,
    ) else {
        return;
    };
    start_jump(commands, bunny, bunny_transform, to);
}

fn is_near_others(
    locator: &Locator,
    bunny: Entity,
    bunny_transform: &Transform,
) -> Option<Vec<(Entity, (Vec3, Dir3))>> {
    let nearby_bunnies = locator.get_nearby(bunny, bunny_transform.translation, DETECTION_DISTANCE);
    if nearby_bunnies.is_empty() {
        None
    } else {
        Some(nearby_bunnies)
    }
}

fn jump_with_others(
    commands: &mut Commands,
    time: &Time,
    obstacles: &mut Obstacles,
    bunny: Entity,
    bunny_transform: &mut Transform,
    nearby_bunnies: Vec<(Entity, (Vec3, Dir3))>,
    rng: &mut Rng,
) {
    if !rng.random_bool((COLLECTIVE_JUMPS_PER_SECOND * time.delta_secs_f64()).clamp(0.0, 1.0)) {
        return;
    }

    let average_direction = calculate_average_direction(nearby_bunnies);
    let Some(to) = obstacles.avoid(
        bunny_transform.translation,
        bunny_transform.translation + average_direction * SHORT_JUMP_DISTANCE,
    ) else {
        return;
    };

    start_jump(commands, bunny, bunny_transform, to);

    fn calculate_average_direction(nearby_bunnies: Vec<(Entity, (Vec3, Dir3))>) -> Vec3 {
        let (sin, cos) = nearby_bunnies
            .iter()
            .skip(1)
            .fold((0.0, 0.0), |a, (_, (_, d))| {
                let angle = d.angle_between(Vec3::Z);
                (a.0 + angle.sin(), a.1 + angle.cos())
            });
        let average_angle = sin.atan2(cos);
        let average_direction = Dir3::Z.rotate_y(average_angle);
        average_direction
    }
}

fn jump_alone(
    commands: &mut Commands,
    time: &Time,
    obstacles: &mut Obstacles,
    bunny: Entity,
    bunny_transform: &mut Transform,
    rng: &mut Rng,
) {
    if !rng.random_bool((LONE_JUMPS_PER_SECOND * time.delta_secs_f64()).clamp(0.0, 1.0)) {
        return;
    }

    let angle = rng.random_range(-std::f32::consts::PI..=std::f32::consts::PI);
    let direction = Dir3::Z.rotate_y(angle);
    let Some(to) = obstacles.avoid(
        bunny_transform.translation,
        bunny_transform.translation + direction * LONG_JUMP_DISTANCE,
    ) else {
        return;
    };

    start_jump(commands, bunny, bunny_transform, to);
}

fn start_jump(commands: &mut Commands, bunny: Entity, bunny_transform: &mut Transform, to: Vec3) {
    commands.entity(bunny).insert(JumpState::Jumping {
        from: bunny_transform.translation,
        to,
    });
    bunny_transform.look_to(-(to - bunny_transform.translation), Vec3::Y);
}

fn jump(
    mut commands: Commands,
    time: Res<Time>,
    mut bunnies: Query<(Entity, &mut Transform, &mut JumpState), With<Bunny>>,
) {
    bunnies.iter_mut().for_each(
        |(bunny, mut bunny_transform, ref mut jump_state)| match jump_state.deref_mut() {
            JumpState::Jumping { from, to } => {
                let from = *from;
                let to = *to;
                animate_jumping(
                    &time,
                    &mut bunny_transform,
                    jump_state.deref_mut(),
                    from,
                    to,
                )
            }
            JumpState::Cooldown { ready } => {
                let ready = *ready;
                cooldown(&mut commands, &time, bunny, ready);
            }
        },
    );
}

fn animate_jumping(
    time: &Time,
    bunny_transform: &mut Transform,
    jump_state: &mut JumpState,
    from: Vec3,
    to: Vec3,
) {
    let current = bunny_transform.translation.with_y(to.y);
    let direction = to - current;
    let delta = SPEED * time.delta_secs();

    let (direction, length) = direction.normalize_and_length();
    if length <= delta {
        bunny_transform.translation = to;
        *jump_state = JumpState::Cooldown {
            ready: time.elapsed() + Duration::from_secs_f32(JUMP_COOLDOWN),
        };
        return;
    }

    let next = current + direction.normalize() * delta;
    let total_distance = from.distance(to);
    bunny_transform.translation = next
        + Vec3::Y
            * JUMP_HEIGHT
            * (std::f32::consts::PI * from.distance(next) / total_distance).sin();
}

fn cooldown(commands: &mut Commands, time: &Time, bunny: Entity, ready: Duration) {
    if time.elapsed() >= ready {
        commands.entity(bunny).remove::<JumpState>();
    }
}