shinespark/src/fighter/baked_animation/visual_binding.rs

use bevy::prelude::*;
use bevy::world_serialization::{WorldInstanceReady, WorldInstanceSpawner};
use bevy_ggrs::GgrsFrameTiming;
use std::collections::{HashMap, HashSet};

use super::{BakedFighterAnimations, FighterBoneMatrices};
use crate::fighter::visual::FighterAnimations;

#[derive(Component)]
pub(super) struct BakedFighterBone {
    fighter: Entity,
    index: usize,
}

pub(super) fn bind_fighter_visual_bones(
    ready: On<WorldInstanceReady>,
    mut commands: Commands,
    spawner: Res<WorldInstanceSpawner>,
    fighters: Query<&FighterAnimations>,
    baked_assets: Res<Assets<BakedFighterAnimations>>,
    nodes: Query<(&Name, Option<&ChildOf>)>,
    animation_players: Query<(), With<AnimationPlayer>>,
) {
    let fighter = ready.entity;
    let Ok(animations) = fighters.get(fighter) else {
        return;
    };
    let baked = baked_assets
        .get(&animations.baked)
        .expect("fighter baked animation asset should be loaded before its visual scene");
    let entities = spawner
        .iter_instance_entities(ready.instance_id)
        .collect::<HashSet<_>>();
    for &entity in &entities {
        if animation_players.contains(entity) {
            commands.entity(entity).remove::<AnimationPlayer>();
        }
    }
    let mut indices = HashMap::with_capacity(baked.bone_names.len());
    for (index, name) in baked.bone_names.iter().enumerate() {
        let key = (baked.parents[index], name.as_str());
        assert!(
            indices.insert(key, index).is_none(),
            "fighter GLTF has duplicate node names under the same parent"
        );
    }
    let mut assigned = HashMap::new();
    let mut remaining = entities.iter().copied().collect::<Vec<_>>();
    while !remaining.is_empty() {
        let mut next = Vec::new();
        let mut progress = false;
        for entity in remaining {
            let Ok((name, parent)) = nodes.get(entity) else {
                continue;
            };
            let parent = parent.map(ChildOf::parent);
            let parent_index = match parent {
                Some(parent) if entities.contains(&parent) => {
                    if let Some(index) = assigned.get(&parent) {
                        Some(*index)
                    } else if nodes.get(parent).ok().is_some_and(|(name, _)| {
                        baked.bone_names.iter().any(|bone| bone == name.as_str())
                    }) {
                        next.push(entity);
                        continue;
                    } else {
                        None
                    }
                }
                _ => None,
            };
            if let Some(&index) = indices.get(&(parent_index, name.as_str())) {
                assigned.insert(entity, index);
                progress = true;
            }
        }
        if !progress {
            break;
        }
        remaining = next;
    }
    if assigned.len() != baked.bone_names.len() {
        warn!(fighter = ?fighter, bound = assigned.len(), expected = baked.bone_names.len(), "could not bind every baked fighter node to its visual scene");
    }
    for (entity, index) in assigned {
        commands
            .entity(entity)
            .insert(BakedFighterBone { fighter, index });
    }
}

pub(super) fn apply_baked_fighter_visual_pose(
    mut nodes: Query<(&mut Transform, &BakedFighterBone)>,
    fighters: Query<&FighterBoneMatrices>,
    timing: Res<GgrsFrameTiming>,
) {
    let overstep_fraction = timing.overstep_fraction();
    for (mut transform, bone) in &mut nodes {
        let Ok(matrices) = fighters.get(bone.fighter) else {
            continue;
        };
        let Some(local) = matrices.presentation_local(bone.index, overstep_fraction) else {
            continue;
        };
        *transform = local;
    }
}