📄 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, In, Query, Res, ResMut, Single},
    },
    gltf::GltfAssetLabel,
    log::warn,
    math::{Dir3, Vec3, primitives::Cuboid},
    mesh::{Mesh, Mesh3d, MeshBuilder, Meshable},
    scene::SceneRoot,
    time::Time,
    transform::components::Transform,
};

use rand::RngExt;

use crate::{
    Rng,
    bunnies::{
        heart::{add_heart_system, spawn_heart},
        locator::BunnyLocator,
    },
    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, look_for_partner, breed));
    add_heart_system(app)
}

#[derive(Component)]
pub struct Bunny;

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

#[derive(Component)]
enum Breeding {
    Initiator { done: Duration, mid_point: Vec3 },
    Partner { done: 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 BREED_DISTANCE: f32 = 3.0;
const BREED_PROBABILITY: f64 = 0.05;
const BREEDING_DURATION: f64 = 1.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 IdleBunnyFilter = (With<Bunny>, Without<JumpState>, Without<Breeding>);

fn calculate(mut commands: Commands, mut bunnies: Query<Entity, IdleBunnyFilter>) {
    bunnies.iter_mut().for_each(|bunny| {
        commands.run_system_cached_with(calculate_next_move, bunny);
    });
}

fn calculate_next_move(
    In(bunny): In<Entity>,
    mut commands: Commands,
    locator: Res<BunnyLocator>,
    dog: Single<&Transform, With<Dog>>,
    transforms: Query<&mut Transform, Without<Dog>>,
) {
    let Ok(bunny_transform) = transforms.get(bunny) else {
        warn!("Could not find transform for known bunny {bunny}");
        return;
    };

    if let Some(direction) = is_dog_nearby(&dog, bunny_transform) {
        commands.run_system_cached_with(jump_away_from_dog, (bunny, direction));
    } else if let Some(nearby_bunnies) = is_near_others(&locator, bunny, bunny_transform) {
        commands.run_system_cached_with(jump_with_others, (bunny, nearby_bunnies));
    } else {
        commands.run_system_cached_with(jump_alone, bunny);
    }
}

fn is_dog_nearby(dog: &Transform, bunny_transform: &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(
    In((bunny, direction)): In<(Entity, Vec3)>,
    mut commands: Commands,
    mut transforms: Query<&mut Transform>,
    mut obstacles: Obstacles,
) {
    let Ok(mut bunny_transform) = transforms.get_mut(bunny) else {
        warn!("Could not find transform for known bunny {bunny}");
        return;
    };

    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(&mut commands, bunny, &mut bunny_transform, to);
}

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

fn jump_with_others(
    In((bunny, nearby_bunnies)): In<(Entity, Vec<(Entity, (Vec3, Dir3))>)>,
    mut commands: Commands,
    mut transforms: Query<&mut Transform>,
    time: Res<Time>,
    mut obstacles: Obstacles,
    mut rng: ResMut<Rng>,
) {
    let Ok(mut bunny_transform) = transforms.get_mut(bunny) else {
        warn!("Could not find transform for known bunny {bunny}");
        return;
    };

    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(&mut commands, bunny, &mut 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(
    In(bunny): In<Entity>,
    mut commands: Commands,
    mut transforms: Query<&mut Transform>,
    time: Res<Time>,
    mut obstacles: Obstacles,
    mut rng: ResMut<Rng>,
) {
    let Ok(mut bunny_transform) = transforms.get_mut(bunny) else {
        warn!("Could not find transform for known bunny {bunny}");
        return;
    };

    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(&mut commands, bunny, &mut 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>();
    }
}

fn look_for_partner(
    mut commands: Commands,
    bunnies: Query<(Entity, &mut Transform), IdleBunnyFilter>,
    time: Res<Time>,
    locator: Res<BunnyLocator>,
    mut rng: ResMut<Rng>,
) {
    bunnies.iter().for_each(|(bunny, bunny_transform)| {
        if !rng.random_bool((time.delta_secs_f64() * BREED_PROBABILITY).clamp(0.0, 1.0)) {
            return;
        }

        let Some((partner, _)) = locator
            .get_nearby_except_self(bunny, bunny_transform.translation, BREED_DISTANCE)
            .drain(..)
            .find(|(_, (p, _))| bunny_transform.translation.x < p.x)
        else {
            return;
        };

        commands.run_system_cached_with(start_breeding, (bunny, partner));
    });
}

fn start_breeding(
    In((a, b)): In<(Entity, Entity)>,
    mut commands: Commands,
    time: Res<Time>,
    mut transforms: Query<&mut Transform>,
) {
    let Ok([mut a_transform, mut b_transform]) = transforms.get_many_mut([a, b]) else {
        return;
    };

    let mid_point = (a_transform.translation + b_transform.translation) / 2.0;
    let breeding_duration = Duration::from_secs_f64(BREEDING_DURATION);
    let done = time.elapsed() + breeding_duration;

    let dir = a_transform.translation - b_transform.translation;
    a_transform.look_to(dir, Vec3::Y);
    let dir = b_transform.translation - a_transform.translation;
    b_transform.look_to(dir, Vec3::Y);

    commands
        .entity(a)
        .insert(Breeding::Initiator { done, mid_point });
    commands.entity(b).insert(Breeding::Partner { done });
    commands.run_system_cached_with(spawn_heart, (mid_point, breeding_duration));
}

fn breed(
    mut commands: Commands,
    time: Res<Time>,
    bunnies: Query<(Entity, &Breeding, &SceneRoot, &Mesh3d), With<Bunny>>,
) {
    bunnies
        .iter()
        .for_each(|(bunny, breeding, model, mesh)| match breeding {
            Breeding::Initiator { done, mid_point } => {
                if time.elapsed() < *done {
                    return;
                }

                commands.entity(bunny).remove::<Breeding>();
                commands.spawn((
                    model.clone(),
                    Transform::from_translation(*mid_point),
                    mesh.clone(),
                    Bunny,
                ));
            }
            Breeding::Partner { done } => {
                if time.elapsed() < *done {
                    return;
                }

                commands.entity(bunny).remove::<Breeding>();
            }
        });
}