//! 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, HitStop,
};
use crate::fighter::attributes::FighterAttributes;
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,
    FighterHitstopShake, FighterHitstopShakeKind, FighterHitstopShakeRole,
    FighterPreviousTranslation, FighterSurfaceState, FighterTranslation, FighterVelocity,
};
use crate::fighter::surface::EcbLockPolicy;
use crate::game_settings::GameSettings;
use crate::input::FighterInput;
use crate::math::int::FGi32;
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;
pub mod crouch;

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub enum SurfacePolicy {
    Grounded,
    Airborne(EcbLockPolicy),
}

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 InterruptContext) -> Option<FighterState> {
                match self { $(Self::$variant(s) => s.check_interrupt(ctx)),* }
            }

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

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

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

            pub fn is_visually_grounded(&self) -> bool {
                match self { $(Self::$variant(s) => s.is_visually_grounded()),* }
            }

            pub fn surface_policy(&self) -> SurfacePolicy {
                match self { $(Self::$variant(s) => s.surface_policy()),* }
            }
        }

        $(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 InterruptContext) -> Option<FighterState>;

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

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

    fn is_visually_grounded<'a>(&'a self) -> bool;

    fn surface_policy(&self) -> SurfacePolicy {
        SurfacePolicy::Grounded
    }
}

pub struct AnimationStatus {
    pub frame: u32,
    pub finished: bool,
}

pub struct InterruptContext<'a> {
    pub input: &'a FighterInput,
    pub facing_direction: &'a mut FighterFacingDirection,
    pub game_settings: &'a GameSettings,
    pub fighter_attributes: &'a FighterAttributes,
    pub animation: AnimationStatus,
}

impl InterruptContext<'_> {
    pub fn get_current_animation_frame(&self) -> u32 {
        self.animation.frame
    }

    pub fn is_current_animation_finished(&self) -> bool {
        self.animation.finished
    }
}

pub struct EnterContext<'a> {
    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 facing_direction: &'a mut FighterFacingDirection,
    pub animation_frame: &'a mut FighterAnimationFrame,
    pub fighter_attributes: &'a FighterAttributes,
    pub attack_script_assets: &'a FighterAttackScriptAssets,
}

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

pub struct ExitContext<'a> {
    pub hitboxes: &'a mut FighterHitboxes,
}

pub struct UpdateContext<'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 surface: &'a mut FighterSurfaceState,
    pub knockback_velocity: &'a mut FighterKnockbackVelocity,
    pub game_settings: &'a GameSettings,
    pub animation_frame: u32,
    pub fighter_attributes: &'a FighterAttributes,
    pub attack_scripts: &'a Assets<FighterAttackScript>,
    pub attack_script_assets: &'a FighterAttackScriptAssets,
    pub facing_direction: &'a FighterFacingDirection,
}

impl UpdateContext<'_> {
    pub fn get_current_animation_frame(&self) -> u32 {
        self.animation_frame
    }
}

pub struct CollisionContext<'a> {
    pub translation: &'a mut FighterTranslation,
    pub prev_translation: &'a FighterPreviousTranslation,
    pub velocity: &'a mut FighterVelocity,
    pub surface: &'a mut FighterSurfaceState,
    pub ecb: &'a mut FighterECB,
    pub prev_ecb: &'a FighterPreviousECB,
    pub stage_collision: &'a StageCollision,
    pub leave_ground_ecb_lock_duration: u32,
}

impl CollisionContext<'_> {
    pub fn air_collide_with_stage(
        &self,
    ) -> Option<crate::fighter::collision::AirCollisionWithStageResult> {
        let ignored_polys = self.surface.ignored_polygon_ids();
        crate::fighter::collision::air_collide_with_stage(
            self.prev_translation,
            self.translation,
            self.prev_ecb,
            self.ecb,
            self.stage_collision,
            &ignored_polys,
        )
    }
}

fighter_states! {
    Wait => wait::WaitState,
    Walk => walk::WalkState,
    Crouch => crouch::CrouchState,
    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 surface: &'static mut FighterSurfaceState,
    pub knockback_velocity: &'static mut FighterKnockbackVelocity,
    pub stale_move_queue: &'static mut FighterStaleMoveQueue,
    pub ecb: &'static mut FighterECB,
    pub prev_ecb: &'static FighterPreviousECB,
    pub facing_direction: &'static mut FighterFacingDirection,
    pub animation_frame: &'static mut FighterAnimationFrame,
    pub fighter: &'static Fighter,
}

impl<'w, 's> FighterFrameQueryItem<'w, 's> {
    fn animation_status(
        &self,
        animations: &FighterAnimations,
        baked_animations: &Assets<BakedFighterAnimations>,
    ) -> AnimationStatus {
        let baked = baked_animations
            .get(&animations.baked)
            .expect("Animation should exist");
        let frame = self.animation_frame.frame;
        let finished = baked
            .frame_count(self.animation_frame.kind)
            .expect("Current animation should have baked frames")
            <= frame;
        AnimationStatus { frame, finished }
    }

    fn interrupt_context<'a>(
        &'a mut self,
        game_settings: &'a GameSettings,
        animations: &'a FighterAnimations,
        baked_animations: &'a Assets<BakedFighterAnimations>,
        attributes: &'a FighterAttributes,
    ) -> InterruptContext<'a> {
        let animation = self.animation_status(animations, baked_animations);
        InterruptContext {
            input: &self.input,
            facing_direction: &mut self.facing_direction,
            game_settings,
            fighter_attributes: attributes,
            animation,
        }
    }

    fn enter_context<'a>(
        &'a mut self,
        attributes: &'a FighterAttributes,
        attack_script_assets: &'a FighterAttackScriptAssets,
    ) -> EnterContext<'a> {
        EnterContext {
            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,
            facing_direction: &mut self.facing_direction,
            animation_frame: &mut self.animation_frame,
            fighter_attributes: attributes,
            attack_script_assets,
        }
    }

    fn exit_context(&mut self) -> ExitContext<'_> {
        ExitContext {
            hitboxes: &mut self.hitboxes,
        }
    }

    fn update_context<'a>(
        &'a mut self,
        game_settings: &'a GameSettings,
        attributes: &'a FighterAttributes,
        attack_scripts: &'a Assets<FighterAttackScript>,
        attack_script_assets: &'a FighterAttackScriptAssets,
    ) -> UpdateContext<'a> {
        UpdateContext {
            translation: &mut self.translation,
            prev_translation: &mut self.prev_translation,
            input: &mut self.input,
            hitboxes: &mut self.hitboxes,
            velocity: &mut self.velocity,
            surface: &mut self.surface,
            knockback_velocity: &mut self.knockback_velocity,
            game_settings,
            animation_frame: self.animation_frame.frame,
            fighter_attributes: attributes,
            attack_scripts,
            attack_script_assets,
            facing_direction: &self.facing_direction,
        }
    }

    fn collision_context<'a>(
        &'a mut self,
        stage_collision: &'a StageCollision,
        game_settings: &'a GameSettings,
    ) -> CollisionContext<'a> {
        CollisionContext {
            translation: &mut self.translation,
            prev_translation: &self.prev_translation,
            velocity: &mut self.velocity,
            surface: &mut self.surface,
            ecb: &mut self.ecb,
            prev_ecb: &self.prev_ecb,
            stage_collision,
            leave_ground_ecb_lock_duration: game_settings.fighter_common.leave_ground_ecb_frame_lock_duration
        }
    }
}

#[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 hitstop_shake: &'static mut FighterHitstopShake,
}

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

        if hitstop.remaining_frames > 0 {
            continue;
        }

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

fn transition_state(
    current: &mut FighterState,
    mut next: FighterState,
    frame: &mut FighterFrameQueryItem,
    attributes: &FighterAttributes,
    attack_script_assets: &FighterAttackScriptAssets,
    game_settings: &GameSettings,
) {
    let old_name = current.name();
    let new_name = next.name();

    current.on_exit(&mut frame.exit_context());
    if let SurfacePolicy::Airborne(lock_policy) = next.surface_policy() {
        frame.surface.apply_airborne_policy(
            lock_policy,
            game_settings.fighter_common.leave_ground_ecb_frame_lock_duration,
        );
    }
    next.on_enter(&mut frame.enter_context(attributes, attack_script_assets));
    *current = next;
    info!("{} -> {}", old_name, new_name);
}

pub fn state_interrupt_system(
    mut query: Query<StatefulFighterQuery>,
    game_settings: Res<GameSettings>,
    baked_animations: Res<Assets<BakedFighterAnimations>>,
    fighter_manifests: Res<Assets<FighterManifest>>,
) {
    for fighter in &mut query {
        let StatefulFighterQueryItem {
            mut state,
            mut frame,
            animations,
            fighter,
            attack_script_assets,
            mut pending_hit_reaction,
            mut hitstop,
            mut hitstop_shake,
        } = fighter;
        let manifest = 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 shake_frames = reaction
                .hitstop_frames
                .calculate_shake_frames(&game_settings);
            hitstop_shake.role = FighterHitstopShakeRole::Receiver;
            hitstop_shake.total_frames = shake_frames;
            hitstop_shake.remaining_frames = shake_frames;
            hitstop_shake.kind = if state.is_visually_grounded() {
                FighterHitstopShakeKind::Horizontal
            } else {
                FighterHitstopShakeKind::Vertical
            };

            let mut hitstop_kind = if frame.surface.is_grounded() {
                hitstun::HitstopStateKind::Grounded
            } else {
                hitstun::HitstopStateKind::Airborne
            };

            if let hitstun::HitstopStateKind::Grounded = hitstop_kind {
                let straight_up_angle = FGi32::lit("90.0");
                if reaction.hitstop_frames
                    >= HitStop::new(game_settings.fighter_common.tumble_bounce_hitstop_threshold)
                {
                    hitstop_kind = hitstun::HitstopStateKind::Airborne;
                // TODO: Use floor normal here
                } else if reaction.launch_angle.abs_diff(straight_up_angle)
                    < game_settings
                        .fighter_common
                        .attack_ground_to_airborne_transition_vertical_angle_threshold
                {
                    hitstop_kind = hitstun::HitstopStateKind::Airborne;
                } else {
                    // straight down is 270
                    let straight_down_angle = FGi32::lit("270.0");
                    let meteor_threshold = game_settings.fighter_common.meteor_angle_threshold;
                    if (reaction.launch_angle - straight_down_angle).abs() <= meteor_threshold {
                        // TODO: Floor bounce
                        hitstop_kind = hitstun::HitstopStateKind::Airborne;
                    }
                }
            }

            info!("Entering hitstop kind: {:?}", hitstop_kind);
            Some(FighterState::Hitstun(HitstunState::new(
                reaction,
                hitstop_kind,
            )))
        } else if hitstop.remaining_frames > 0 {
            None
        } else {
            state.check_interrupt(&mut frame.interrupt_context(
                &game_settings,
                animations,
                &baked_animations,
                &manifest.attributes,
            ))
        };

        if let Some(new_state) = new_state {
            transition_state(
                &mut state,
                new_state,
                &mut frame,
                &manifest.attributes,
                attack_script_assets,
                &game_settings,
            );
        };
    }
}

pub fn state_collision_interrupt_system(
    mut query: Query<StatefulFighterQuery>,
    stage_collision: Res<StageCollision>,
    fighter_manifests: Res<Assets<FighterManifest>>,
    game_settings: Res<GameSettings>,
) {
    for fighter in &mut query {
        let StatefulFighterQueryItem {
            mut state,
            mut frame,
            fighter,
            attack_script_assets,
            hitstop,
            ..
        } = fighter;
        if hitstop.remaining_frames > 0 {
            continue;
        }

        let manifest = fighter_manifests
            .get(&fighter.manifest)
            .expect("Fighter manifest should be valid");
        let new_state =
            state.check_collision_interrupt(&mut frame.collision_context(&stage_collision, &game_settings));
        if let Some(new_state) = new_state {
            transition_state(
                &mut state,
                new_state,
                &mut frame,
                &manifest.attributes,
                &attack_script_assets,
                &game_settings,
            );
        }
    }
}