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::{
    AppState, args::Args, fighter::{
        FighterId, FighterRoster,
        animation::AnimKind,
        attack::{AttackKind, FighterAttackScriptAssets},
        baked_animation::{BAKED_ANIMATION_FPS, BakedFighterAnimations},
        manifest::{FighterManifest, FighterManifestRegistry},
        spawn_fighter,
        visual::FighterAnimations,
    }, input::{LocalInputAssignments, LocalInputSource, MatchInputAssignments}, match_runtime::MatchState, netcode::{
        GGRSCfg,
        session::{ConnectionStatus, SessionConfig, create_session},
    }, scripting::FighterAttackScript, stage::{
        self,
        manifest::{StageId, StageManifest, StageManifestRegistry},
    },
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct MatchPlayer {
    /// Dense index used by GGRS and rollback-simulated player state.
    pub handle: usize,
    /// Physical controller port selected in the menu.
    pub controller_port: usize,
    pub fighter: FighterId,
}

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

impl MatchConfig {
    /// Maps a GGRS handle to the controller port used by gameplay.
    ///
    /// Multiplayer uses the stable, dense GGRS handle as the port. A
    /// single-player match preserves the physical controller port selected in
    /// the menu (for example, port 3 remains port 3).
    pub fn controller_port_for_handle(&self, handle: usize) -> Option<usize> {
        let player = self.players.iter().find(|player| player.handle == handle)?;
        Some(if self.players.len() > 1 {
            handle
        } else {
            player.controller_port
        })
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum MatchRules {
    Stock {
        starting_stocks: u8,
        time_limit_frames: Option<u32>,
    },
    Timed {
        time_limit_frames: u32,
    },
}

impl Default for MatchRules {
    fn default() -> Self {
        Self::Stock {
            starting_stocks: 3,
            time_limit_frames: None,
        }
    }
}

impl MatchRules {
    fn validate(&self) -> Result<(), &'static str> {
        match self {
            Self::Stock {
                starting_stocks: 0, ..
            } => Err("stock matches must start with at least one stock"),
            Self::Stock {
                time_limit_frames: Some(0),
                ..
            } => Err("stock match time limits must be greater than zero frames"),
            Self::Timed {
                time_limit_frames: 0,
            } => Err("timed match limits must be greater than zero frames"),
            _ => Ok(()),
        }
    }
}

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 config = world.resource::<MatchConfig>();
        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(config.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: config
                .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>,
    config: MatchConfig,
) {
    commands.remove_resource::<crate::netcode::session::MatchDisconnected>();
    commands.insert_resource(config);
    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 config = replay.match_config.clone();
                commands.insert_resource(crate::replay::LoadedReplay(replay));
                initiate_match(&mut commands, &mut next_state, config);
            }
            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,
            controller_port: 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,
        MatchConfig {
            stage: StageId::TestStage,
            players,
            rules: MatchRules::default(),
        },
    );
}

pub fn prepare_match(
    mut commands: Commands,
    config: Res<MatchConfig>,
    input_assignments: Res<LocalInputAssignments>,
    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(&config.players) {
        fail_match(
            &mut next_state,
            "player handles must be dense, unique, and non-empty",
        );
        return;
    }

    if let Err(error) = config.rules.validate() {
        fail_match(&mut next_state, error);
        return;
    }

    let stage_handle = stage_registry
        .get(config.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");
    let spawn_positions = match stage_manifest.neutral_spawn_positions(config.players.len()) {
        Ok(positions) => positions,
        Err(error) => {
            fail_match(&mut next_state, &error);
            return;
        }
    };

    if assets.fighters.len() != config.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 !baked.contains_bone(&manifest.root_motion_bone) {
            fail_match(
                &mut next_state,
                &format!(
                    "{:?} references missing root motion bone {:?}",
                    character.fighter, manifest.root_motion_bone
                ),
            );
            return;
        }

        let ecb_bone_indices = baked.bone_names.iter().enumerate().filter_map(|(idx, name)| {
            if manifest.ecb_bones.contains(name) {
                Some(idx)
            } else {
                None
            }
        }).collect::<Vec<_>>();

        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,
            ecb_bone_indices,
        ));
    }

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

    commands.insert_resource(create_match_input_assignments(
        &input_assignments,
        &config.players,
    ));

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

    stage::spawn_stage(
        &mut commands,
        root,
        assets.stage.clone(),
        stage_manifest.to_collision(),
        stage_manifest.camera.clone(),
    );

    let mut fighter_roster = Vec::with_capacity(config.players.len());

    for (spawn_index, (player, (scene, animations, attack_scripts, manifest_handle, model_scale, ecb_bone_indices))) in
        config.players.iter().zip(prepared).enumerate()
    {
        let spawn_point = spawn_positions[spawn_index];
        fighter_roster.push(spawn_fighter(
            &mut commands,
            player.handle,
            config
                .controller_port_for_handle(player.handle)
                .expect("validated match config contains every player handle"),
            spawn_point,
            manifest_handle.clone(),
            model_scale,
            animations,
            attack_scripts,
            WorldAssetRoot(scene),
            ecb_bone_indices
        ));
    }

    commands.insert_resource(FighterRoster(fighter_roster));
    commands.insert_resource(MatchState::from_config(&config));
    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 {
    if players.is_empty() {
        return false;
    }

    let mut handles = players.iter().map(|player| player.handle).collect::<Vec<_>>();
    handles.sort_unstable();
    handles.into_iter().eq(0..players.len())
}

fn player_count_for_handles(players: &[MatchPlayer]) -> usize {
    players.len()
}

fn create_match_input_assignments(
    assignments: &LocalInputAssignments,
    players: &[MatchPlayer],
) -> MatchInputAssignments {
    let mut match_assignments = vec![None; players.len()];
    for player in players {
        match_assignments[player.handle] = assignments
            .get_from_slot(player.controller_port)
            .copied();
    }
    MatchInputAssignments(match_assignments)
}

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::<MatchConfig>();
    commands.remove_resource::<MatchState>();
    commands.remove_resource::<MatchInputAssignments>();
    commands.remove_resource::<FighterRoster>();
}

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

    #[test]
    fn match_config_keeps_duplicate_fighter_slots() {
        let config = MatchConfig {
            stage: StageId::TestStage,
            players: vec![
                MatchPlayer {
                    handle: 0,
                    controller_port: 0,
                    fighter: FighterId::TestFighter,
                },
                MatchPlayer {
                    handle: 1,
                    controller_port: 1,
                    fighter: FighterId::TestFighter,
                },
            ],
            rules: MatchRules::default(),
        };

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

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

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

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

        assert!(!valid_player_handles(&players));
        assert_eq!(player_count_for_handles(&players), 2);
    }

    #[test]
    fn controller_ports_use_ggrs_handles_only_in_multiplayer() {
        let single_player = MatchConfig {
            stage: StageId::TestStage,
            players: vec![MatchPlayer {
                handle: 0,
                controller_port: 3,
                fighter: FighterId::TestFighter,
            }],
            rules: MatchRules::default(),
        };
        assert_eq!(single_player.controller_port_for_handle(0), Some(3));

        let multiplayer = MatchConfig {
            stage: StageId::TestStage,
            players: vec![
                MatchPlayer {
                    handle: 0,
                    controller_port: 0,
                    fighter: FighterId::TestFighter,
                },
                MatchPlayer {
                    handle: 1,
                    controller_port: 3,
                    fighter: FighterId::TestFighter,
                },
            ],
            rules: MatchRules::default(),
        };
        assert_eq!(multiplayer.controller_port_for_handle(0), Some(0));
        assert_eq!(multiplayer.controller_port_for_handle(1), Some(1));
    }

    #[test]
    fn match_input_assignments_map_ports_to_ggrs_handles_without_mutating_ports() {
        let players = [
            MatchPlayer {
                handle: 0,
                controller_port: 0,
                fighter: FighterId::TestFighter,
            },
            MatchPlayer {
                handle: 1,
                controller_port: 3,
                fighter: FighterId::TestFighter,
            },
        ];
        let assignments = LocalInputAssignments(vec![
            (0, LocalInputSource::Keyboard),
            (3, LocalInputSource::Gamepad(Entity::PLACEHOLDER)),
        ]);

        let match_assignments = create_match_input_assignments(&assignments, &players);

        assert_eq!(match_assignments.get_for_handle(0), Some(LocalInputSource::Keyboard));
        assert_eq!(
            match_assignments.get_for_handle(1),
            Some(LocalInputSource::Gamepad(Entity::PLACEHOLDER))
        );
        assert_eq!(assignments[0].0, 0);
        assert_eq!(assignments[1].0, 3);
    }

    #[test]
    fn match_rules_reject_zero_limits() {
        assert!(
            MatchRules::Stock {
                starting_stocks: 0,
                time_limit_frames: None,
            }
            .validate()
            .is_err()
        );
        assert!(
            MatchRules::Stock {
                starting_stocks: 3,
                time_limit_frames: Some(0),
            }
            .validate()
            .is_err()
        );
        assert!(
            MatchRules::Timed {
                time_limit_frames: 0,
            }
            .validate()
            .is_err()
        );
    }

    #[test]
    fn match_rules_accept_valid_modes() {
        assert!(MatchRules::default().validate().is_ok());
        assert!(
            MatchRules::Stock {
                starting_stocks: 3,
                time_limit_frames: Some(60),
            }
            .validate()
            .is_ok()
        );
        assert!(
            MatchRules::Timed {
                time_limit_frames: 60,
            }
            .validate()
            .is_ok()
        );
    }

    #[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);
    }
}