shinespark/src/match_loading.rs

//! Match selection, asset loading, preparation, and cleanup.

use std::collections::HashMap;

use bevy::{asset::UntypedHandle, gltf::Gltf, prelude::*};
use bevy_asset_loader::prelude::*;
use bevy_ggrs::prelude::Session;
use serde::{Deserialize, Serialize};

use crate::{
    args::Args,
    fighter::{
        animation::AnimKind,
        attack::{AttackKind, FighterAttackScriptAssets},
        baked_animation::{BakedFighterAnimations, BAKED_ANIMATION_FPS},
        manifest::{FighterManifest, FighterManifestRegistry},
        spawn_fighter,
        visual::FighterAnimations,
        FighterId, FighterRoster,
    },
    input::{LocalInputAssignments, LocalInputSource},
    math::{int::FGi32, vec::FGVec2},
    netcode::{
        session::{create_session, ConnectionStatus, SessionConfig},
        GGRSCfg,
    },
    scripting::FighterAttackScript,
    stage::{
        self,
        manifest::{StageId, StageManifest, StageManifestRegistry},
    },
    AppState,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct MatchPlayer {
    pub handle: usize,
    pub fighter: FighterId,
}

#[derive(Resource, Debug, Clone, PartialEq, Eq)]
pub struct PendingMatch {
    pub stage: StageId,
    pub players: Vec<MatchPlayer>,
}

pub struct SelectedFighterAssets {
    pub fighter: FighterId,
    pub model: Handle<Gltf>,
    pub baked_animations: Handle<BakedFighterAnimations>,
    pub attack_scripts: HashMap<AttackKind, Handle<FighterAttackScript>>,
}

#[derive(Resource)]
pub struct MatchAssets {
    pub stage: Handle<WorldAsset>,
    pub fighters: Vec<SelectedFighterAssets>,
}

impl MatchAssets {
    fn requested(world: &World) -> Self {
        let request = world.resource::<PendingMatch>();
        let asset_server = world.resource::<AssetServer>();
        let registry = world.resource::<FighterManifestRegistry>();
        let stage_registry = world.resource::<StageManifestRegistry>();
        let manifests = world.resource::<Assets<FighterManifest>>();
        let stage_manifests = world.resource::<Assets<StageManifest>>();

        let stage_handle = stage_registry
            .get(request.stage)
            .expect("Startup validation guarantees every stage has a manifest");
        let stage_manifest = stage_manifests
            .get(&stage_handle)
            .expect("startup-loaded stage manifest should remain available");

        Self {
            stage: asset_server.load(stage_manifest.model_path.clone()),
            fighters: request
                .players
                .iter()
                .map(|player| {
                    let handle = registry
                        .get(player.fighter)
                        .expect("startup validation guarantees every fighter has a manifest");
                    let manifest = manifests
                        .get(handle)
                        .expect("startup-loaded fighter manifest should remain available");

                    SelectedFighterAssets {
                        fighter: player.fighter,
                        model: asset_server.load(manifest.model_path.clone()),
                        baked_animations: asset_server.load(manifest.baked_animation_path.clone()),
                        attack_scripts: manifest
                            .attack_scripts
                            .clone()
                            .into_iter()
                            .map(|(kind, path)| (kind, asset_server.load(path)))
                            .collect(),
                    }
                })
                .collect(),
        }
    }

    fn into_loading_handles(self) -> Vec<UntypedHandle> {
        let mut handles = Vec::with_capacity(
            1 + self
                .fighters
                .iter()
                .map(|fighter| 2 + fighter.attack_scripts.len())
                .sum::<usize>(),
        );
        handles.push(self.stage.untyped());
        for character in self.fighters {
            handles.push(character.model.untyped());
            handles.push(character.baked_animations.untyped());
            handles.extend(
                character
                    .attack_scripts
                    .into_values()
                    .map(|handle| handle.untyped()),
            );
        }
        handles
    }
}

impl AssetCollection for MatchAssets {
    fn create(world: &mut World) -> Self {
        Self::requested(world)
    }

    fn load(world: &mut World) -> Vec<UntypedHandle> {
        Self::requested(world).into_loading_handles()
    }
}

#[derive(Component)]
pub struct MatchRoot;

pub fn initiate_match(
    commands: &mut Commands,
    next_state: &mut NextState<AppState>,
    request: PendingMatch,
) {
    commands.remove_resource::<crate::netcode::session::MatchDisconnected>();
    commands.insert_resource(request);
    next_state.set(AppState::LoadingMatch);
}

pub fn initiate_default_match(
    mut commands: Commands,
    args: Res<Args>,
    gamepads: Query<Entity, With<Gamepad>>,
    mut assigments: ResMut<LocalInputAssignments>,
    mut next_state: ResMut<NextState<AppState>>,
    completed: Option<Res<crate::replay::ReplayComplete>>,
) {
    if !args.debug_match {
        next_state.set(AppState::MainMenu);
        return;
    }
    if completed.is_some() {
        return;
    }
    if let Some(path) = &args.play_replay {
        match crate::replay::load(path) {
            Ok(replay) => {
                let request = PendingMatch {
                    stage: replay.match_config.stage,
                    players: replay.match_config.players.clone(),
                };
                commands.insert_resource(crate::replay::LoadedReplay(replay));
                initiate_match(&mut commands, &mut next_state, request);
            }
            Err(error) => {
                error!("Replay load failed: {error}");
                next_state.set(AppState::MatchLoadFailed);
            }
        }
        return;
    }
    let session_config = match SessionConfig::from_args(&args, args.players) {
        Ok(config) => config,
        Err(error) => {
            error!("Invalid network configuration: {error}");
            next_state.set(AppState::MatchLoadFailed);
            return;
        }
    };
    let players = (0..args.players)
        .map(|handle| MatchPlayer {
            handle,
            fighter: FighterId::TestFighter,
        })
        .collect();

    assigments.clear();
    let local_handles = session_config
        .players
        .iter()
        .enumerate()
        .filter_map(|(handle, player)| {
            matches!(player, crate::netcode::session::PlayerConnection::Local).then_some(handle)
        });
    let mut local_handles = local_handles;
    if let Some(handle) = local_handles.next() {
        assigments.push((handle, LocalInputSource::Keyboard));
    }

    for (entity, handle) in gamepads.iter().zip(local_handles) {
        assigments.push((handle, LocalInputSource::Gamepad(entity)));
    }

    initiate_match(
        &mut commands,
        &mut next_state,
        PendingMatch {
            stage: StageId::TestStage,
            players,
        },
    );
}

pub fn prepare_match(
    mut commands: Commands,
    request: Res<PendingMatch>,
    assets: Res<MatchAssets>,
    gltfs: Res<Assets<Gltf>>,
    baked_animations: Res<Assets<BakedFighterAnimations>>,
    attack_scripts: Res<Assets<FighterAttackScript>>,
    registry: Res<FighterManifestRegistry>,
    stage_registry: Res<StageManifestRegistry>,
    manifests: Res<Assets<FighterManifest>>,
    stage_manifests: Res<Assets<StageManifest>>,
    args: Res<Args>,
    mut next_state: ResMut<NextState<AppState>>,
) {
    if !valid_player_handles(&request.players) {
        fail_match(
            &mut next_state,
            "player handles must be unique and cover 0..player_count",
        );
        return;
    }

    if assets.fighters.len() != request.players.len() {
        fail_match(
            &mut next_state,
            "loaded character assets do not match the requested player slots",
        );
        return;
    }

    let mut prepared = Vec::with_capacity(assets.fighters.len());

    for character in &assets.fighters {
        let manifest_handle = registry
            .get(character.fighter)
            .expect("startup validation guarantees every fighter has a manifest");
        let Some(manifest) = manifests.get(manifest_handle) else {
            fail_match(&mut next_state, "a fighter manifest is unavailable");
            return;
        };
        let Some(gltf) = gltfs.get(&character.model) else {
            fail_match(&mut next_state, "a loaded character GLTF is unavailable");
            return;
        };
        let Some(baked) = baked_animations.get(&character.baked_animations) else {
            fail_match(
                &mut next_state,
                "a fighter baked animation asset is unavailable",
            );
            return;
        };
        if baked.frame_rate != BAKED_ANIMATION_FPS {
            fail_match(
                &mut next_state,
                &format!(
                    "{:?} baked animation rate is {}, expected {}",
                    character.fighter, baked.frame_rate, BAKED_ANIMATION_FPS
                ),
            );
            return;
        }
        if !baked.has_unique_bone_names() {
            fail_match(
                &mut next_state,
                &format!(
                    "{:?} baked animation contains duplicate bone names",
                    character.fighter
                ),
            );
            return;
        }
        if let Err(error) = manifest.attributes.validate() {
            fail_match(
                &mut next_state,
                &format!("{:?} has invalid attributes: {error}", character.fighter),
            );
            return;
        }
        for (index, hurtbox) in manifest.hurtboxes.iter().enumerate() {
            if let Err(error) = hurtbox.validate() {
                fail_match(
                    &mut next_state,
                    &format!(
                        "{:?} hurtbox {index} is invalid: {error}",
                        character.fighter
                    ),
                );
                return;
            }
            if !baked.contains_bone(&hurtbox.bone) {
                fail_match(
                    &mut next_state,
                    &format!(
                        "{:?} hurtbox {index} references missing bone {:?}",
                        character.fighter, hurtbox.bone
                    ),
                );
                return;
            }
        }
        for (kind, handle) in &character.attack_scripts {
            let Some(script) = attack_scripts.get(handle) else {
                fail_match(
                    &mut next_state,
                    &format!(
                        "{:?} {kind:?} attack script is unavailable",
                        character.fighter
                    ),
                );
                return;
            };
            if let Err(error) = script.validate() {
                fail_match(
                    &mut next_state,
                    &format!(
                        "{:?} {kind:?} attack script is invalid: {error}",
                        character.fighter
                    ),
                );
                return;
            }
            for hitbox in &script.hitboxes {
                if !baked.contains_bone(&hitbox.bone) {
                    fail_match(
                        &mut next_state,
                        &format!(
                            "{:?} {kind:?} hitbox {} references missing bone {:?}",
                            character.fighter, hitbox.id, hitbox.bone
                        ),
                    );
                    return;
                }
            }
        }
        let Some(scene) = gltf.scenes.first().cloned() else {
            fail_match(&mut next_state, "a character GLTF contains no scene");
            return;
        };

        for kind in manifest.animations.keys() {
            if baked.frame_count(*kind).is_none() {
                fail_match(
                    &mut next_state,
                    &format!(
                        "{:?} baked data has no {kind:?} animation",
                        character.fighter
                    ),
                );
                return;
            }
        }

        if baked.frame_count(AnimKind::Wait).is_none() {
            fail_match(&mut next_state, "a character has no Wait animation");
            return;
        }

        let animations = FighterAnimations {
            baked: character.baked_animations.clone(),
        };

        let attack_scripts = FighterAttackScriptAssets {
            scripts: character.attack_scripts.clone(),
        };

        prepared.push((
            scene,
            animations,
            attack_scripts,
            manifest_handle,
            manifest.model_scale,
        ));
    }

    let session_config = match SessionConfig::from_args(&args, request.players.len()) {
        Ok(config) => config,
        Err(error) => {
            fail_match(
                &mut next_state,
                &format!("invalid network configuration: {error}"),
            );
            return;
        }
    };
    let session = match create_session(&args, request.players.len()) {
        Ok(session) => session,
        Err(error) => {
            fail_match(
                &mut next_state,
                &format!("failed to create GGRS session: {error}"),
            );
            return;
        }
    };

    let root = commands.spawn((Name::new("Match"), MatchRoot)).id();

    let stage_handle = stage_registry
        .get(request.stage)
        .expect("Initialization check should ensure all stages exist");
    let stage_manifest = stage_manifests
        .get(&stage_handle)
        .expect("Initialization check should ensure stage manifests are kept alive");
    stage::spawn_stage(
        &mut commands,
        root,
        assets.stage.clone(),
        stage_manifest.to_collision(),
        stage_manifest.camera.clone(),
    );

    let mut fighter_roster = Vec::with_capacity(request.players.len());
    for (index, (player, (scene, animations, attack_scripts, manifest_handle, model_scale))) in
        request.players.iter().zip(prepared).enumerate()
    {
        let spawn_x = (index as i32 * 2 + 1 - request.players.len() as i32) * 5;
        fighter_roster.push(spawn_fighter(
            &mut commands,
            player.handle,
            FGVec2::new(FGi32::from_num(spawn_x), FGi32::lit("12.5")),
            manifest_handle.clone(),
            model_scale,
            animations,
            attack_scripts,
            WorldAssetRoot(scene),
        ));
    }

    commands.insert_resource(FighterRoster(fighter_roster));
    commands.insert_resource(session_config);
    commands.insert_resource(ConnectionStatus::default());
    commands.insert_resource::<Session<GGRSCfg>>(session);
    next_state.set(AppState::InMatch);
}

fn fail_match(next_state: &mut NextState<AppState>, message: &str) {
    error!("Match preparation failed: {message}");
    next_state.set(AppState::MatchLoadFailed);
}

fn valid_player_handles(players: &[MatchPlayer]) -> bool {
    let mut handles = players
        .iter()
        .map(|player| player.handle)
        .collect::<Vec<_>>();
    handles.sort_unstable();
    handles == (0..players.len()).collect::<Vec<_>>()
}

pub fn cleanup_match(
    mut commands: Commands,
    roots: Query<Entity, With<MatchRoot>>,
    fighters: Query<Entity, With<crate::player::Player>>,
) {
    for root in &roots {
        commands.entity(root).despawn();
    }
    for fighter in &fighters {
        commands.entity(fighter).despawn();
    }
    commands.remove_resource::<Session<GGRSCfg>>();
    commands.remove_resource::<SessionConfig>();
    commands.remove_resource::<ConnectionStatus>();
    commands.remove_resource::<stage::line::StageCollision>();
    commands.remove_resource::<MatchAssets>();
    commands.remove_resource::<PendingMatch>();
    commands.remove_resource::<FighterRoster>();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn match_request_keeps_duplicate_fighter_slots() {
        let request = PendingMatch {
            stage: StageId::TestStage,
            players: vec![
                MatchPlayer {
                    handle: 0,
                    fighter: FighterId::TestFighter,
                },
                MatchPlayer {
                    handle: 1,
                    fighter: FighterId::TestFighter,
                },
            ],
        };

        assert_eq!(request.players.len(), 2);
        assert_eq!(request.players[0].handle, 0);
        assert_eq!(request.players[1].handle, 1);
        assert!(valid_player_handles(&request.players));
    }

    #[test]
    fn player_handles_must_match_the_session_range() {
        let players = [
            MatchPlayer {
                handle: 0,
                fighter: FighterId::TestFighter,
            },
            MatchPlayer {
                handle: 0,
                fighter: FighterId::TestFighter,
            },
        ];

        assert!(!valid_player_handles(&players));
    }

    #[test]
    fn match_loading_waits_for_attack_scripts() {
        let mut attack_scripts = HashMap::new();
        attack_scripts.insert(AttackKind::Jab, Handle::default());
        let assets = MatchAssets {
            stage: Handle::default(),
            fighters: vec![SelectedFighterAssets {
                fighter: FighterId::TestFighter,
                model: Handle::default(),
                baked_animations: Handle::default(),
                attack_scripts,
            }],
        };

        assert_eq!(assets.into_loading_handles().len(), 4);
    }
}