//! Rollback-simulated state for an active match.

use bevy::input::InputSystems;
use bevy::prelude::*;
use bevy_ggrs::{ConfirmedFrameCount, RollbackFrameCount, RunGgrsSystems, prelude::*};

use crate::{
    AppState,
    match_loading::{MatchConfig, MatchRules},
    netcode::session::SessionConfig,
};

#[derive(Resource, Clone, Debug, Hash, PartialEq, Eq)]
pub struct MatchState {
    pub elapsed_frames: u32,
    /// Indexed by the player's session handle.
    pub players: Vec<PlayerMatchState>,
    /// The rollback frame on which the match first finished.
    pub finished_at_frame: Option<i32>,
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct PlayerMatchState {
    pub stocks_remaining: Option<u8>,
    pub score: i32,
    pub deaths: u32,
}

#[derive(Resource, Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct PauseState {
    pub paused: bool,
}

impl MatchState {
    pub fn from_config(config: &MatchConfig) -> Self {
        let stocks_remaining = match &config.rules {
            MatchRules::Stock {
                starting_stocks, ..
            } => Some(*starting_stocks),
            MatchRules::Timed { .. } => None,
        };

        Self {
            elapsed_frames: 0,
            players: (0..config.players.len())
                .map(|_| PlayerMatchState {
                    stocks_remaining,
                    score: 0,
                    deaths: 0,
                })
                .collect(),
            finished_at_frame: None,
        }
    }

    pub fn is_match_finished(&self, config: &MatchConfig) -> bool {
        let time_limit = match config.rules {
            MatchRules::Stock {
                starting_stocks: _,
                time_limit_frames,
            } => time_limit_frames,
            MatchRules::Timed { time_limit_frames } => Some(time_limit_frames),
        };

        if let Some(time_limit) = time_limit {
            if self.elapsed_frames >= time_limit {
                return true;
            }
        }

        if config.players.len() > 1
            && let MatchRules::Stock {
                starting_stocks: _,
                time_limit_frames: _,
            } = config.rules
        {
            let alive_fighter_count = self
                .players
                .iter()
                .filter(|fighter| {
                    fighter
                        .stocks_remaining
                        .map(|stocks| stocks)
                        .unwrap_or_default()
                        > 0
                })
                .count();
            if alive_fighter_count <= 1 {
                return true;
            }
        }

        false
    }

    fn advance_frame(&mut self, config: &MatchConfig, frame: i32) {
        if self.finished_at_frame.is_some() {
            return;
        }

        self.elapsed_frames = self.elapsed_frames.saturating_add(1);

        if self.is_match_finished(config) {
            self.finished_at_frame = Some(frame);
        }
    }

    fn finish_is_confirmed(&self, confirmed_frame: i32) -> bool {
        self.finished_at_frame
            .is_some_and(|finished_at| confirmed_frame >= finished_at)
    }
}

pub struct MatchRuntimePlugin;

impl Plugin for MatchRuntimePlugin {
    fn build(&self, app: &mut App) {
        app.rollback_resource_with_clone::<MatchState>()
            .checksum_resource_with_hash::<MatchState>()
            .init_resource::<PauseState>()
            .configure_sets(PreUpdate, RunGgrsSystems.run_if(ggrs_should_run))
            .add_systems(OnEnter(AppState::InMatch), reset_pause)
            .add_systems(OnExit(AppState::InMatch), reset_pause)
            .add_systems(
                GgrsSchedule,
                update_match_rules.in_set(crate::schedule::GameplaySet::MatchRules),
            )
            .add_systems(
                PreUpdate,
                toggle_pause
                    .after(InputSystems)
                    .before(RunGgrsSystems)
                    .run_if(in_state(AppState::InMatch)),
            )
            .add_systems(
                Update,
                transition_after_confirmed_finish.run_if(in_state(AppState::InMatch)),
            );
    }
}

pub fn ggrs_should_run(pause: Option<Res<PauseState>>) -> bool {
    !pause.is_some_and(|pause| pause.paused)
}

pub(crate) fn match_is_ongoing(match_state: Option<Res<MatchState>>) -> bool {
    match_state.is_none_or(|match_state| match_state.finished_at_frame.is_none())
}

fn toggle_pause(
    key: Res<ButtonInput<KeyCode>>,
    gamepads: Query<&Gamepad>,
    session: Res<SessionConfig>,
    mut pause: ResMut<PauseState>,
    mut virtual_time: ResMut<Time<Virtual>>,
) {
    if session.is_multiplayer() {
        return;
    }

    let pressed = key.just_pressed(KeyCode::Escape)
        || gamepads
            .iter()
            .any(|gamepad| gamepad.just_pressed(GamepadButton::Start));
    if !pressed {
        return;
    }

    pause.paused = !pause.paused;
    virtual_time.toggle();
}

fn reset_pause(mut pause: ResMut<PauseState>, mut virtual_time: ResMut<Time<Virtual>>) {
    pause.paused = false;
    virtual_time.unpause();
}

fn update_match_rules(
    mut match_state: ResMut<MatchState>,
    match_config: Res<MatchConfig>,
    frame: Res<RollbackFrameCount>,
) {
    match_state.advance_frame(&match_config, frame.0);
}

fn transition_after_confirmed_finish(
    match_state: Res<MatchState>,
    confirmed_frame: Res<ConfirmedFrameCount>,
    mut next_state: ResMut<NextState<AppState>>,
) {
    if match_state.finish_is_confirmed(confirmed_frame.0) {
        next_state.set(AppState::MatchResults);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{fighter::FighterId, match_loading::MatchPlayer, stage::manifest::StageId};

    fn config(rules: MatchRules) -> MatchConfig {
        MatchConfig {
            stage: StageId::TestStage,
            players: vec![
                MatchPlayer {
                    handle: 1,
                    controller_port: 3,
                    fighter: FighterId::TestFighter,
                },
                MatchPlayer {
                    handle: 0,
                    controller_port: 0,
                    fighter: FighterId::TestFighter,
                },
            ],
            rules,
        }
    }

    #[test]
    fn stock_state_starts_with_configured_stocks() {
        let state = MatchState::from_config(&config(MatchRules::Stock {
            starting_stocks: 4,
            time_limit_frames: Some(60),
        }));

        assert_eq!(state.elapsed_frames, 0);
        assert_eq!(state.finished_at_frame, None);
        assert_eq!(state.players.len(), 2);
        assert!(state.players.iter().all(|player| {
            player.stocks_remaining == Some(4) && player.score == 0 && player.deaths == 0
        }));
    }

    #[test]
    fn timed_state_has_no_stock_counter() {
        let state = MatchState::from_config(&config(MatchRules::Timed {
            time_limit_frames: 60,
        }));

        assert!(
            state
                .players
                .iter()
                .all(|player| player.stocks_remaining.is_none())
        );
    }

    #[test]
    fn timed_match_finishes_on_its_rollback_frame() {
        let config = config(MatchRules::Timed {
            time_limit_frames: 2,
        });
        let mut state = MatchState::from_config(&config);

        state.advance_frame(&config, 10);
        assert_eq!(state.elapsed_frames, 1);
        assert_eq!(state.finished_at_frame, None);

        state.advance_frame(&config, 11);
        assert_eq!(state.elapsed_frames, 2);
        assert_eq!(state.finished_at_frame, Some(11));
        assert!(!state.finish_is_confirmed(10));
        assert!(state.finish_is_confirmed(11));

        state.advance_frame(&config, 12);
        assert_eq!(state.finished_at_frame, Some(11));
        assert_eq!(state.elapsed_frames, 2);
    }

    #[test]
    fn corrected_rollback_state_does_not_keep_speculative_finish() {
        let config = config(MatchRules::Stock {
            starting_stocks: 1,
            time_limit_frames: None,
        });
        let state_before_prediction = MatchState::from_config(&config);
        let mut predicted_state = state_before_prediction.clone();
        predicted_state.players[1].stocks_remaining = Some(0);
        predicted_state.advance_frame(&config, 20);
        assert_eq!(predicted_state.finished_at_frame, Some(20));

        let mut corrected_state = state_before_prediction;
        corrected_state.advance_frame(&config, 20);
        assert_eq!(corrected_state.finished_at_frame, None);
        assert!(!corrected_state.finish_is_confirmed(20));
    }

    #[test]
    fn single_player_stock_match_does_not_finish_by_elimination() {
        let mut config = config(MatchRules::Stock {
            starting_stocks: 1,
            time_limit_frames: None,
        });
        config.players.truncate(1);
        let mut state = MatchState::from_config(&config);

        state.advance_frame(&config, 1);
        assert_eq!(state.finished_at_frame, None);

        state.players[0].stocks_remaining = Some(0);
        state.advance_frame(&config, 2);
        assert_eq!(state.finished_at_frame, None);
    }

    #[test]
    fn single_player_stock_match_can_finish_by_time_limit() {
        let mut config = config(MatchRules::Stock {
            starting_stocks: 1,
            time_limit_frames: Some(1),
        });
        config.players.truncate(1);
        let mut state = MatchState::from_config(&config);

        state.advance_frame(&config, 7);
        assert_eq!(state.finished_at_frame, Some(7));
    }
}