//! Uses glTF extension processing to play an animation on a skinned glTF model of a fox.

use std::f32::consts::PI;

use bevy::{
    animation::{AnimationEvent, AnimationTargetId},
    asset::LoadContext,
    ecs::entity::EntityHashSet,
    gltf::extensions::{ErasedGltfExtensionHandler, GltfExtensionHandler, GltfExtensionHandlers},
    light::CascadeShadowConfigBuilder,
    platform::collections::{HashMap, HashSet},
    prelude::*,
    world_serialization::WorldInstanceReady,
};
use chacha20::ChaCha8Rng;
use rand::{RngExt, SeedableRng};

/// An example asset that contains a mesh and animation.
const GLTF_PATH: &str = "models/animated/Fox.glb";

fn main() {
    App::new()
        .insert_resource(GlobalAmbientLight {
            color: Color::WHITE,
            brightness: 2000.,
            ..default()
        })
        .add_plugins((DefaultPlugins, GltfExtensionHandlerAnimationPlugin))
        .init_resource::<ParticleAssets>()
        .add_systems(
            Startup,
            (setup_mesh_and_animation, setup_camera_and_environment),
        )
        .add_systems(Update, simulate_particles)
        .add_observer(observe_on_step)
        .run();
}

#[derive(Resource)]
struct SeededRng(ChaCha8Rng);

/// A component that stores a reference to an animation we want to play. This is
/// created when we start loading the mesh (see `setup_mesh_and_animation`) and
/// read when the mesh has spawned (see `play_animation_once_loaded`).
#[derive(Component, Reflect)]
#[reflect(Component)]
struct AnimationToPlay {
    graph_handle: Handle<AnimationGraph>,
    index: AnimationNodeIndex,
}

fn setup_mesh_and_animation(mut commands: Commands, asset_server: Res<AssetServer>) {
    // Spawn an entity with our components, and connect it to an observer that
    // will trigger when the scene is loaded and spawned.
    commands
        .spawn(WorldAssetRoot(
            asset_server.load(GltfAssetLabel::Scene(0).from_asset(GLTF_PATH)),
        ))
        .observe(play_animation_when_ready);
}

fn play_animation_when_ready(
    scene_ready: On<WorldInstanceReady>,
    mut commands: Commands,
    children: Query<&Children>,
    mut players: Query<(&mut AnimationPlayer, &AnimationToPlay)>,
) {
    for child in children.iter_descendants(scene_ready.entity) {
        let Ok((mut player, animation_to_play)) = players.get_mut(child) else {
            continue;
        };

        // Tell the animation player to start the animation and keep
        // repeating it.
        //
        // If you want to try stopping and switching animations, see the
        // `animated_mesh_control.rs` example.
        player.play(animation_to_play.index).repeat();

        // Add the animation graph. This only needs to be done once to
        // connect the animation player to the mesh.
        commands
            .entity(child)
            .insert(AnimationGraphHandle(animation_to_play.graph_handle.clone()));
    }
}

/// Spawn a camera and a simple environment with a ground plane and light.
fn setup_camera_and_environment(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    // Camera
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(100.0, 100.0, 150.0).looking_at(Vec3::new(0.0, 20.0, 0.0), Vec3::Y),
    ));

    // Plane
    commands.spawn((
        Mesh3d(meshes.add(Plane3d::default().mesh().size(500000.0, 500000.0))),
        MeshMaterial3d(materials.add(Color::srgb(0.3, 0.5, 0.3))),
    ));

    // Light
    commands.spawn((
        Transform::from_rotation(Quat::from_euler(EulerRot::ZYX, 0.0, 1.0, -PI / 4.)),
        DirectionalLight {
            shadow_maps_enabled: true,
            ..default()
        },
        CascadeShadowConfigBuilder {
            first_cascade_far_bound: 200.0,
            maximum_distance: 400.0,
            ..default()
        }
        .build(),
    ));

    // We're seeding the PRNG here to make this example deterministic for testing purposes.
    // This isn't strictly required in practical use unless you need your app to be deterministic.
    let seeded_rng = ChaCha8Rng::seed_from_u64(19878367467712);
    commands.insert_resource(SeededRng(seeded_rng));
}

struct GltfExtensionHandlerAnimationPlugin;

impl Plugin for GltfExtensionHandlerAnimationPlugin {
    fn build(&self, app: &mut App) {
        #[cfg(target_family = "wasm")]
        bevy::tasks::block_on(async {
            app.world_mut()
                .resource_mut::<GltfExtensionHandlers>()
                .0
                .write()
                .await
                .push(Box::new(GltfExtensionHandlerAnimation::default()))
        });
        #[cfg(not(target_family = "wasm"))]
        app.world_mut()
            .resource_mut::<GltfExtensionHandlers>()
            .0
            .write_blocking()
            .push(Box::new(GltfExtensionHandlerAnimation::default()));
    }
}

#[derive(Default, Clone)]
struct GltfExtensionHandlerAnimation {
    animation_root_indices: HashSet<usize>,
    animation_root_entities: EntityHashSet,
    clip: Option<Handle<AnimationClip>>,
}

impl GltfExtensionHandler for GltfExtensionHandlerAnimation {
    fn dyn_clone(&self) -> Box<dyn ErasedGltfExtensionHandler> {
        Box::new((*self).clone())
    }

    #[cfg(feature = "bevy_animation")]
    fn on_animation(
        &mut self,
        _load_context: &mut LoadContext<'_>,
        gltf_animation: &gltf::Animation,
        animation_clip: &mut AnimationClip,
    ) {
        if gltf_animation.name().is_some_and(|v| v == "Run") {
            let hip_node = ["root", "_rootJoint", "b_Root_00", "b_Hip_01"];
            let front_left_foot = hip_node.iter().chain(
                [
                    "b_Spine01_02",
                    "b_Spine02_03",
                    "b_LeftUpperArm_09",
                    "b_LeftForeArm_010",
                    "b_LeftHand_011",
                ]
                .iter(),
            );
            let front_right_foot = hip_node.iter().chain(
                [
                    "b_Spine01_02",
                    "b_Spine02_03",
                    "b_RightUpperArm_06",
                    "b_RightForeArm_07",
                    "b_RightHand_08",
                ]
                .iter(),
            );
            let back_left_foot = hip_node.iter().chain(
                [
                    "b_LeftLeg01_015",
                    "b_LeftLeg02_016",
                    "b_LeftFoot01_017",
                    "b_LeftFoot02_018",
                ]
                .iter(),
            );
            let back_right_foot = hip_node.iter().chain(
                [
                    "b_RightLeg01_019",
                    "b_RightLeg02_020",
                    "b_RightFoot01_021",
                    "b_RightFoot02_022",
                ]
                .iter(),
            );
            animation_clip.add_event_to_target(
                AnimationTargetId::from_iter(front_left_foot),
                0.625,
                Step,
            );
            animation_clip.add_event_to_target(
                AnimationTargetId::from_iter(front_right_foot),
                0.5,
                Step,
            );
            animation_clip.add_event_to_target(
                AnimationTargetId::from_iter(back_left_foot),
                0.0,
                Step,
            );
            animation_clip.add_event_to_target(
                AnimationTargetId::from_iter(back_right_foot),
                0.125,
                Step,
            );
        }
    }
    #[cfg(feature = "bevy_animation")]
    fn on_animations_collected(
        &mut self,
        _load_context: &mut LoadContext<'_>,
        _animations: &[Handle<AnimationClip>],
        named_animations: &HashMap<Box<str>, Handle<AnimationClip>>,
        animation_roots: &HashSet<usize>,
    ) {
        self.animation_root_indices = animation_roots.clone();

        if let Some(handle) = named_animations.get("Run") {
            self.clip = Some(handle.clone());
        }
    }

    fn on_gltf_node(
        &mut self,
        _load_context: &mut LoadContext<'_>,
        gltf_node: &gltf::Node,
        entity: &mut EntityWorldMut,
    ) {
        if self.animation_root_indices.contains(&gltf_node.index()) {
            self.animation_root_entities.insert(entity.id());
        }
    }

    /// Called when an individual Scene is done processing
    fn on_scene_completed(
        &mut self,
        load_context: &mut LoadContext<'_>,
        _scene: &gltf::Scene,
        _world_root_id: Entity,
        world: &mut World,
    ) {
        // Create an AnimationGraph from the desired clip
        let (graph, index) = AnimationGraph::from_clip(self.clip.clone().unwrap());
        // Store the animation graph as an asset with an arbitrary label
        // We only have one graph, so this label will be unique
        let graph_handle =
            load_context.add_labeled_asset("MyAnimationGraphLabel".to_string(), graph);

        // Create a component that stores a reference to our animation
        let animation_to_play = AnimationToPlay {
            graph_handle,
            index,
        };

        // Insert the `AnimationToPlay` component on the first animation root
        let mut entity = world.entity_mut(*self.animation_root_entities.iter().next().unwrap());
        entity.insert(animation_to_play);
    }
}

fn simulate_particles(
    mut commands: Commands,
    mut query: Query<(Entity, &mut Transform, &mut Particle)>,
    time: Res<Time>,
) {
    for (entity, mut transform, mut particle) in &mut query {
        if particle.lifetime_timer.tick(time.delta()).just_finished() {
            commands.entity(entity).despawn();
            return;
        }

        transform.translation += particle.velocity * time.delta_secs();
        transform.scale = Vec3::splat(particle.size.lerp(0.0, particle.lifetime_timer.fraction()));
        particle
            .velocity
            .smooth_nudge(&Vec3::ZERO, 4.0, time.delta_secs());
    }
}

#[derive(Component)]
struct Particle {
    lifetime_timer: Timer,
    size: f32,
    velocity: Vec3,
}

#[derive(Resource)]
struct ParticleAssets {
    mesh: Handle<Mesh>,
    material: Handle<StandardMaterial>,
}

impl FromWorld for ParticleAssets {
    fn from_world(world: &mut World) -> Self {
        Self {
            mesh: world.add_asset::<Mesh>(Sphere::new(10.0)),
            material: world.add_asset::<StandardMaterial>(StandardMaterial {
                base_color: Color::WHITE,
                ..Default::default()
            }),
        }
    }
}

#[derive(AnimationEvent, Reflect, Clone)]
struct Step;

fn observe_on_step(
    step: On<Step>,
    particle: Res<ParticleAssets>,
    mut commands: Commands,
    transforms: Query<&GlobalTransform>,
    mut seeded_rng: ResMut<SeededRng>,
) -> Result {
    let translation = transforms.get(step.trigger().target)?.translation();
    // Spawn a bunch of particles.
    for _ in 0..14 {
        let horizontal = seeded_rng.0.random::<Dir2>() * seeded_rng.0.random_range(8.0..12.0);
        let vertical = seeded_rng.0.random_range(0.0..4.0);
        let size = seeded_rng.0.random_range(0.2..1.0);

        commands.spawn((
            Particle {
                lifetime_timer: Timer::from_seconds(
                    seeded_rng.0.random_range(0.2..0.6),
                    TimerMode::Once,
                ),
                size,
                velocity: Vec3::new(horizontal.x, vertical, horizontal.y) * 10.0,
            },
            Mesh3d(particle.mesh.clone()),
            MeshMaterial3d(particle.material.clone()),
            Transform {
                translation,
                scale: Vec3::splat(size),
                ..Default::default()
            },
        ));
    }
    Ok(())
}