//! The fighter state machine: the [`FighterState`] trait every state
//! implements, the context they inspect, and the generic system that runs the
//! interrupt check. One state per submodule.

use bevy::ecs::query::QueryData;
use bevy::prelude::*;

use crate::fighter::animation::{AnimKind, FighterAnimationFrame};
use crate::fighter::attack::{
    FighterAttackScriptAssets, FighterKnockbackVelocity, FighterPendingHitReaction,
    FighterStaleMoveQueue,
};
use crate::fighter::baked_animation::BakedFighterAnimations;
use crate::fighter::ecb::FighterPreviousECB;
use crate::fighter::manifest::FighterManifest;
use crate::fighter::state::hitstun::HitstunState;
use crate::fighter::visual::FighterAnimations;
use crate::fighter::{
    Fighter, FighterECB, FighterFacingDirection, FighterHitboxes, FighterHitstop,
    FighterPreviousTranslation, FighterTranslation, FighterVelocity,
};
use crate::game_settings::GameSettings;
use crate::input::FighterInput;
use crate::scripting::FighterAttackScript;
use crate::stage::line::StageCollision;

pub mod air;
pub mod air_dodge;
pub mod dash;
pub mod fall;
pub mod ground;
pub mod ground_attack;
pub mod hitstun;
pub mod jump;
pub mod land;
pub mod run;
pub mod turn;
pub mod wait;
pub mod walk;

macro_rules! fighter_states {
    ($($variant:ident => $ty:ty),* $(,)?) => {
        #[derive(Component, Clone, Debug, Hash)]
        pub enum FighterState { $($variant($ty)),* }

        impl FighterState {
            pub fn name(&self) -> &'static str {
                match self { $(Self::$variant(_) => <$ty>::NAME),* }
            }
            pub fn check_interrupt(&self, ctx: &mut FighterStateContext) -> Option<FighterState> {
                match self { $(Self::$variant(s) => s.check_interrupt(ctx)),* }
            }

            pub fn update(&mut self, ctx: &mut FighterStateContext) {
                match self { $(Self::$variant(s) => s.update(ctx)),* }
            }
            pub fn on_enter(&mut self, ctx: &mut FighterStateContext) {
                match self { $(Self::$variant(s) => s.on_enter(ctx)),* }
            }

            pub fn on_exit(&mut self, ctx: &mut FighterStateContext) {
                match self { $(Self::$variant(s) => s.on_exit(ctx)),* }
            }

            pub fn check_collision_interrupt(&mut self, ctx: &mut FighterStateContext) -> Option<FighterState> {
                match self { $(Self::$variant(s) => s.check_collision_interrupt(ctx)),* }
            }
        }

        $(impl From<$ty> for FighterState {
            fn from(s: $ty) -> Self { Self::$variant(s) }
        })*
    };
}

pub trait FighterStateImpl: Sized {
    const NAME: &'static str;
    fn check_interrupt(&self, state_context: &mut FighterStateContext) -> Option<FighterState>;

    fn on_enter(&mut self, _state_context: &mut FighterStateContext) {}
    fn on_exit(&mut self, _state_context: &mut FighterStateContext) {}
    fn name(&self) -> String {
        Self::NAME.to_owned()
    }
    fn update(&mut self, _state_context: &mut FighterStateContext);

    fn check_collision_interrupt(
        &mut self,
        _state_context: &mut FighterStateContext,
    ) -> Option<FighterState> {
        None
    }
}

pub struct FighterStateContext<'a> {
    pub translation: &'a mut FighterTranslation,
    pub prev_translation: &'a mut FighterPreviousTranslation,
    pub input: &'a mut FighterInput,
    pub hitboxes: &'a mut FighterHitboxes,
    pub velocity: &'a mut FighterVelocity,
    pub knockback_velocity: &'a mut FighterKnockbackVelocity,
    pub stale_move_queue: &'a mut FighterStaleMoveQueue,
    pub ecb: &'a mut FighterECB,
    pub prev_ecb: &'a mut FighterPreviousECB,
    pub facing_direction: &'a mut FighterFacingDirection,
    pub game_settings: &'a GameSettings,
    pub stage_collision: &'a StageCollision,
    pub animations: &'a FighterAnimations,
    pub animation_frame: &'a mut FighterAnimationFrame,
    pub baked_animations: &'a Assets<BakedFighterAnimations>,
    pub attack_scripts: &'a Assets<FighterAttackScript>,
    pub fighter_manifest: &'a FighterManifest,
    pub attack_script_assets: &'a FighterAttackScriptAssets,
}

impl FighterStateContext<'_> {
    pub fn play_animation(&mut self, kind: AnimKind, repeat: bool) {
        self.animation_frame.reset(kind, repeat);
    }

    pub fn get_current_animation_frame(&self) -> u32 {
        self.animation_frame.frame
    }

    pub fn is_current_animation_finished(&self) -> bool {
        let anims_data = self
            .baked_animations
            .get(&self.animations.baked)
            .expect("Animation should exist");
        anims_data
            .frame_count(self.animation_frame.kind)
            .expect("Current animation should have baked frames")
            <= self.animation_frame.frame
    }
}

fighter_states! {
    Wait => wait::WaitState,
    Walk => walk::WalkState,
    Dash => dash::DashState,
    Run => run::RunState,
    Fall => fall::FallState,
    JumpSquat => jump::JumpSquatState,
    Jump => jump::JumpState,
    Land => land::LandingState,
    AirDodge => air_dodge::AirDodgeState,
    Turn => turn::TurnState,
    GroundAttack => ground_attack::GroundAttackState,
    Hitstun => hitstun::HitstunState,
}

#[derive(QueryData)]
#[query_data(mutable)]
pub struct FighterFrameQuery {
    pub translation: &'static mut FighterTranslation,
    pub prev_translation: &'static mut FighterPreviousTranslation,
    pub input: &'static mut FighterInput,
    pub hitboxes: &'static mut FighterHitboxes,
    pub velocity: &'static mut FighterVelocity,
    pub knockback_velocity: &'static mut FighterKnockbackVelocity,
    pub stale_move_queue: &'static mut FighterStaleMoveQueue,
    pub ecb: &'static mut FighterECB,
    pub prev_ecb: &'static mut FighterPreviousECB,
    pub facing_direction: &'static mut FighterFacingDirection,
    pub animation_frame: &'static mut FighterAnimationFrame,
    pub fighter: &'static Fighter,
}

impl<'w, 's> FighterFrameQueryItem<'w, 's> {
    fn context<'a>(
        &'a mut self,
        game_settings: &'a GameSettings,
        stage_collision: &'a StageCollision,
        animations: &'a FighterAnimations,
        baked_animations: &'a Assets<BakedFighterAnimations>,
        attack_scripts: &'a Assets<FighterAttackScript>,
        attack_script_assets: &'a FighterAttackScriptAssets,
        manifest: &'a FighterManifest,
    ) -> FighterStateContext<'a> {
        FighterStateContext {
            translation: &mut self.translation,
            prev_translation: &mut self.prev_translation,
            input: &mut self.input,
            hitboxes: &mut self.hitboxes,
            velocity: &mut self.velocity,
            knockback_velocity: &mut self.knockback_velocity,
            stale_move_queue: &mut self.stale_move_queue,
            ecb: &mut self.ecb,
            prev_ecb: &mut self.prev_ecb,
            facing_direction: &mut self.facing_direction,
            game_settings,
            stage_collision,
            animations,
            animation_frame: &mut self.animation_frame,
            baked_animations,
            attack_scripts,
            attack_script_assets,
            fighter_manifest: manifest,
        }
    }
}

#[derive(QueryData)]
#[query_data(mutable)]
pub struct StatefulFighterQuery {
    pub state: &'static mut FighterState,
    pub frame: FighterFrameQuery,
    pub animations: &'static FighterAnimations,
    pub fighter: &'static Fighter,
    pub attack_script_assets: &'static FighterAttackScriptAssets,
    pub pending_hit_reaction: &'static mut FighterPendingHitReaction,
    pub hitstop: &'static mut FighterHitstop,
}

pub fn state_update_system(
    mut query: Query<StatefulFighterQuery>,
    stage_collision: Res<StageCollision>,
    game_settings: Res<GameSettings>,
    baked_animations: Res<Assets<BakedFighterAnimations>>,
    attack_scripts: Res<Assets<FighterAttackScript>>,
    manifests: Res<Assets<FighterManifest>>,
) {
    for fighter in &mut query {
        let StatefulFighterQueryItem {
            mut state,
            mut frame,
            animations,
            hitstop,
            ..
        } = fighter;

        if hitstop.remaining_frames > 0 {
            continue;
        }

        let mut update_ctx = frame.context(
            &game_settings,
            &stage_collision,
            animations,
            &baked_animations,
            &attack_scripts,
            &fighter.attack_script_assets,
            &manifests
                .get(&(fighter.fighter.manifest))
                .expect("Fighter manifest should be valid"),
        );
        state.update(&mut update_ctx);
    }
}

fn transition_state(
    current: &mut FighterState,
    mut next: FighterState,
    ctx: &mut FighterStateContext,
) {
    let old_name = current.name();
    let new_name = next.name();

    current.on_exit(ctx);
    next.on_enter(ctx);
    *current = next;
    info!("{} -> {}", old_name, new_name);
}

pub fn state_interrupt_system(
    mut query: Query<StatefulFighterQuery>,
    stage_collision: Res<StageCollision>,
    game_settings: Res<GameSettings>,
    baked_animations: Res<Assets<BakedFighterAnimations>>,
    fighter_manifests: Res<Assets<FighterManifest>>,
    attack_scripts: Res<Assets<FighterAttackScript>>,
) {
    for fighter in &mut query {
        let StatefulFighterQueryItem {
            mut state,
            mut frame,
            animations,
            fighter,
            attack_script_assets,
            mut pending_hit_reaction,
            mut hitstop,
        } = fighter;
        let mut state_context = frame.context(
            &game_settings,
            &stage_collision,
            animations,
            &baked_animations,
            &attack_scripts,
            attack_script_assets,
            fighter_manifests
                .get(&fighter.manifest)
                .expect("Fighter manifest should be valid"),
        );

        let new_state = if let Some(reaction) = pending_hit_reaction.take() {
            hitstop.remaining_frames = *reaction.hitstop_frames;
            let hitstop_kind = match &*state {
                FighterState::Wait(value) => {
                    hitstun::HitstopStateKind::Grounded(value.grounded_common.clone())
                }
                FighterState::Walk(value) => {
                    hitstun::HitstopStateKind::Grounded(value.grounded_common.clone())
                }
                FighterState::Dash(value) => {
                    hitstun::HitstopStateKind::Grounded(value.ground_common.clone())
                }
                FighterState::Run(value) => {
                    hitstun::HitstopStateKind::Grounded(value.grounded_common.clone())
                }
                FighterState::Turn(value) => {
                    hitstun::HitstopStateKind::Grounded(value.grounded_common.clone())
                }
                FighterState::GroundAttack(value) => {
                    hitstun::HitstopStateKind::Grounded(value.grounded_common.clone())
                }
                _ => hitstun::HitstopStateKind::Airborne,
            };
            Some(FighterState::Hitstun(HitstunState {
                hit_reaction: reaction,
                hitstop_kind,
            }))
        } else if hitstop.remaining_frames > 0 {
            None
        } else {
            state.check_interrupt(&mut state_context)
        };

        if let Some(new_state) = new_state {
            transition_state(&mut state, new_state, &mut state_context);
        };
    }
}

pub fn state_collision_interrupt_system(
    mut query: Query<StatefulFighterQuery>,
    stage_collision: Res<StageCollision>,
    game_settings: Res<GameSettings>,
    baked_animations: Res<Assets<BakedFighterAnimations>>,
    fighter_manifests: Res<Assets<FighterManifest>>,
    attack_scripts: Res<Assets<FighterAttackScript>>,
) {
    for fighter in &mut query {
        let StatefulFighterQueryItem {
            mut state,
            mut frame,
            animations,
            fighter,
            attack_script_assets,
            hitstop,
            ..
        } = fighter;
        let mut state_context = frame.context(
            &game_settings,
            &stage_collision,
            animations,
            &baked_animations,
            &attack_scripts,
            &attack_script_assets,
            fighter_manifests
                .get(&fighter.manifest)
                .expect("Fighter manifest should be valid"),
        );

        if hitstop.remaining_frames > 0 {
            continue;
        }

        if let Some(new_state) = state.check_collision_interrupt(&mut state_context) {
            transition_state(&mut state, new_state, &mut state_context);
        }
    }
}