HEDZReloaded/src/character.rs

use crate::{
    GameState, heads_database::HeadsDatabase, loading_assets::GameAssets, utils::trail::Trail,
};
use bevy::{
    ecs::system::SystemParam, platform::collections::HashMap, prelude::*, scene::SceneInstanceReady,
};
use std::{f32::consts::PI, time::Duration};

#[derive(Component, Debug)]
pub struct ProjectileOrigin;

#[derive(Component, Debug)]
pub struct AnimatedCharacter {
    head: usize,
    rotate_180: bool,
}

impl AnimatedCharacter {
    pub fn new(head: usize) -> Self {
        Self {
            head,
            rotate_180: false,
        }
    }

    pub fn with_rotation(self) -> Self {
        let mut s = self;
        s.rotate_180 = true;
        s
    }
}

#[derive(Component, Debug)]
struct AnimatedCharacterAsset(pub Handle<Gltf>);

#[derive(SystemParam)]
pub struct CharacterHierarchy<'w, 's> {
    descendants: Query<'w, 's, &'static Children>,
    projectile_origin: Query<'w, 's, &'static GlobalTransform, With<ProjectileOrigin>>,
}

impl CharacterHierarchy<'_, '_> {
    pub fn projectile_origin(&self, entity: Entity) -> Option<&GlobalTransform> {
        self.descendants
            .iter_descendants(entity)
            .find_map(|child| self.projectile_origin.get(child).ok())
    }
}

#[derive(Component, Debug, Reflect)]
#[reflect(Component)]
pub struct CharacterAnimations {
    pub idle: AnimationNodeIndex,
    pub run: AnimationNodeIndex,
    pub jump: AnimationNodeIndex,
    pub shooting: Option<AnimationNodeIndex>,
    pub graph: Handle<AnimationGraph>,
}

pub fn plugin(app: &mut App) {
    app.add_systems(
        Update,
        (spawn, setup_once_loaded).run_if(in_state(GameState::Playing)),
    );
    #[cfg(feature = "dbg")]
    app.add_systems(
        Update,
        debug_show_projectile_origin_and_trial.run_if(in_state(GameState::Playing)),
    );
}

fn spawn(
    mut commands: Commands,
    query: Query<(Entity, &AnimatedCharacter), Added<AnimatedCharacter>>,
    gltf_assets: Res<Assets<Gltf>>,
    assets: Res<GameAssets>,
    heads_db: Res<HeadsDatabase>,
) {
    for (entity, character) in query.iter() {
        let key = heads_db.head_key(character.head);

        let handle = assets
            .characters
            .get(format!("{key}.glb").as_str())
            .unwrap_or_else(|| {
                //TODO: remove once we use the new format for all
                info!("Character not found, using default [{}]", key);
                &assets.characters["angry demonstrator.glb"]
            });
        let asset = gltf_assets.get(handle).unwrap();

        let mut t =
            Transform::from_translation(Vec3::new(0., -1.45, 0.)).with_scale(Vec3::splat(1.2));

        if character.rotate_180 {
            t.rotate_y(PI);
        }

        commands
            .entity(entity)
            .insert((
                t,
                SceneRoot(asset.scenes[0].clone()),
                AnimatedCharacterAsset(handle.clone()),
            ))
            .observe(find_marker_bones);
    }
}

fn find_marker_bones(
    trigger: Trigger<SceneInstanceReady>,
    mut commands: Commands,
    descendants: Query<&Children>,
    name: Query<&Name>,
    mut gizmo_assets: ResMut<Assets<GizmoAsset>>,
) {
    let entity = trigger.target();

    let mut origin_found = false;
    for child in descendants.iter_descendants(entity) {
        let Ok(name) = name.get(child) else {
            continue;
        };

        if name.as_str() == "ProjectileOrigin" {
            commands.entity(child).insert(ProjectileOrigin);
            origin_found = true;
        } else if name.as_str().starts_with("Trail") {
            commands.entity(child).insert((
                Trail::new(
                    20,
                    LinearRgba::new(1., 1.0, 1., 0.5),
                    LinearRgba::new(1., 1., 1., 0.5),
                ),
                Gizmo {
                    handle: gizmo_assets.add(GizmoAsset::default()),
                    line_config: GizmoLineConfig {
                        width: 24.,
                        ..default()
                    },
                    ..default()
                },
            ));
        }
    }

    if !origin_found {
        error!("ProjectileOrigin not found");
    }
}

fn setup_once_loaded(
    mut commands: Commands,
    mut query: Query<(Entity, &mut AnimationPlayer), Added<AnimationPlayer>>,
    parent: Query<&ChildOf>,
    animated_character: Query<(&AnimatedCharacter, &AnimatedCharacterAsset)>,

    gltf_assets: Res<Assets<Gltf>>,
    mut graphs: ResMut<Assets<AnimationGraph>>,
) {
    for (entity, mut player) in &mut query {
        let Some((_character, asset)) = parent
            .iter_ancestors(entity)
            .find_map(|ancestor| animated_character.get(ancestor).ok())
        else {
            continue;
        };

        let asset = gltf_assets.get(asset.0.id()).unwrap();

        let animations = asset
            .named_animations
            .iter()
            .map(|(name, animation)| (name.to_string(), animation.clone()))
            .collect::<HashMap<_, _>>();

        let mut clips = vec![
            animations["idle"].clone(),
            animations["run"].clone(),
            animations["jump"].clone(),
        ];

        if let Some(shooting_animation) = animations.get("shoot") {
            clips.push(shooting_animation.clone());
        }

        let (graph, node_indices) = AnimationGraph::from_clips(clips);

        // Insert a resource with the current scene information
        let graph_handle = graphs.add(graph);
        let animations = CharacterAnimations {
            idle: node_indices[0],
            run: node_indices[1],
            jump: node_indices[2],
            shooting: if node_indices.len() == 4 {
                Some(node_indices[3])
            } else {
                None
            },
            graph: graph_handle.clone(),
        };

        let mut transitions = AnimationTransitions::new();
        transitions
            .play(&mut player, animations.idle, Duration::ZERO)
            .repeat();
        commands
            .entity(entity)
            .insert(AnimationGraphHandle(animations.graph.clone()))
            .insert(transitions)
            .insert(animations);
    }
}

#[cfg(feature = "dbg")]
fn debug_show_projectile_origin_and_trial(
    mut gizmos: Gizmos,
    query: Query<&GlobalTransform, Or<(With<ProjectileOrigin>, With<Trail>)>>,
) {
    for projectile_origin in query.iter() {
        gizmos.sphere(
            Isometry3d::from_translation(projectile_origin.translation()),
            0.1,
            Color::linear_rgb(0., 1., 0.),
        );
    }
}