bglga/src/enemy.rs

use bevy::prelude::*;
use std::time::Duration;

use crate::components::{
    AttackPattern, Captured, Enemy, EnemyBullet, EnemyState, EnemyType, FormationTarget,
    OriginalFormationPosition, Player, TractorBeam, TractorBeamSprite,
};
use crate::constants::{
    BEAM_CORE_COLOR, BEAM_GLOW_COLOR, BEAM_GLOW_WIDTH, BEAM_PULSE_AMPLITUDE, BEAM_PULSE_FREQ,
    CAPTURE_DURATION, ENEMY_BULLET_SIZE, ENEMY_SIZE, ENEMY_SPEED, TRACTOR_BEAM_DURATION,
    TRACTOR_BEAM_WIDTH, WINDOW_HEIGHT, WINDOW_WIDTH,
};
use crate::resources::{
    AttackDiveTimer, CurrentStage, EnemySpawnTimer, FormationState, StageConfigurations,
    is_special_stage,
};

const BOSS_BASE_CHANCE: f32 = 0.30;
const BOSS_CHANCE_PER_STAGE: f32 = 0.05;
const BOSS_CHANCE_BONUS_CAP: f32 = 0.50;
const BOSS_CAPTURE_BEAM_PROBABILITY: f32 = 0.80;
const ATTACKER_SPEED_MULTIPLIER: f32 = 1.5;
const RETURN_SPEED_MULTIPLIER: f32 = 0.7;
const SWOOP_HORIZONTAL_FACTOR: f32 = 0.5;
const KAMIKAZE_NEAR_DISTANCE: f32 = 50.0;
const RETURN_ARRIVAL_DISTANCE: f32 = 5.0;
const ARRIVAL_THRESHOLD_FACTOR: f32 = 1.1;
const BOSS_CAPTURE_BEAM_HOVER_Y_FACTOR: f32 = 0.25; // -WINDOW_HEIGHT * this
const BOSS_DESCENT_SPEED_FACTOR: f32 = 0.8;

pub fn spawn_enemies(
    mut commands: Commands,
    time: Res<Time>,
    mut timer: ResMut<EnemySpawnTimer>,
    mut stage: ResMut<CurrentStage>,
    mut formation_state: ResMut<FormationState>,
    stage_configs: Res<StageConfigurations>,
) {
    let config = stage_configs.for_stage(stage.number);
    timer.timer.tick(time.delta());

    if formation_state.total_spawned_this_stage >= config.enemy_count
        || !timer.timer.just_finished()
    {
        return;
    }

    let slot = formation_state.next_slot_index;
    if slot >= config.formation_layout.positions.len() {
        warn!(
            "slot_index {} out of bounds for formation '{}' (size {})",
            slot,
            config.formation_layout.name,
            config.formation_layout.positions.len()
        );
        timer.timer.reset();
        return;
    }

    let target_pos = config.formation_layout.positions[slot];
    let spawn_pos = Vec3::new(
        (fastrand::f32() - 0.5) * (WINDOW_WIDTH - ENEMY_SIZE.x),
        WINDOW_HEIGHT / 2.0 + ENEMY_SIZE.y / 2.0,
        0.0,
    );

    let boss_chance = BOSS_BASE_CHANCE
        + (stage.number as f32 * BOSS_CHANCE_PER_STAGE).min(BOSS_CHANCE_BONUS_CAP);
    let enemy_type = if fastrand::f32() < boss_chance {
        EnemyType::Boss
    } else {
        EnemyType::Grunt
    };

    let color = match enemy_type {
        EnemyType::Grunt => Color::srgb(1.0, 0.2, 0.2),
        EnemyType::Boss => Color::srgb(0.8, 0.2, 1.0),
    };

    commands.spawn((
        Sprite {
            color,
            custom_size: Some(ENEMY_SIZE),
            ..default()
        },
        Transform::from_translation(spawn_pos),
        Enemy {
            enemy_type,
            shoot_cooldown: Timer::new(
                Duration::from_secs_f32(config.enemy_shoot_interval),
                TimerMode::Once,
            ),
        },
        FormationTarget {
            position: target_pos,
        },
        EnemyState::Entering,
    ));

    formation_state.next_slot_index = (slot + 1) % config.enemy_count;
    formation_state.total_spawned_this_stage += 1;
    stage.waiting_for_clear = true;

    if formation_state.total_spawned_this_stage < config.enemy_count {
        timer.timer.reset();
    } else {
        info!(
            "Stage {} formation '{}' fully spawned",
            stage.number, config.formation_layout.name
        );
    }
}

pub fn move_enemies(
    mut entering_query: Query<
        (Entity, &mut Transform, &FormationTarget, &mut EnemyState),
        (With<Enemy>, With<FormationTarget>),
    >,
    mut attacking_query: Query<
        (
            Entity,
            &mut Transform,
            &mut EnemyState,
            &Enemy,
            Option<&OriginalFormationPosition>,
        ),
        (With<Enemy>, Without<FormationTarget>),
    >,
    time: Res<Time>,
    mut commands: Commands,
    stage: Res<CurrentStage>,
    stage_configs: Res<StageConfigurations>,
    has_beam_query: Query<&TractorBeam>,
) {
    let dt = time.delta_secs();
    let speed_multiplier = stage_configs.for_stage(stage.number).enemy_speed_multiplier;
    let base_speed = ENEMY_SPEED * speed_multiplier;
    let attack_speed = base_speed * ATTACKER_SPEED_MULTIPLIER;
    let arrival_threshold = base_speed * dt * ARRIVAL_THRESHOLD_FACTOR;

    for (entity, mut transform, target, mut state) in entering_query.iter_mut() {
        if *state != EnemyState::Entering {
            continue;
        }
        let direction = target.position - transform.translation;
        if direction.length() < arrival_threshold {
            transform.translation = target.position;
            commands.entity(entity).remove::<FormationTarget>();
            commands.entity(entity).insert(OriginalFormationPosition {
                position: target.position,
            });
            *state = EnemyState::InFormation;
        } else {
            transform.translation += direction.normalize() * base_speed * dt;
        }
    }

    for (entity, mut transform, mut state, enemy, original_pos) in attacking_query.iter_mut() {
        match state.as_ref() {
            EnemyState::Attacking(pattern) => {
                step_attacker(
                    &mut commands,
                    entity,
                    &mut transform,
                    enemy.enemy_type,
                    *pattern,
                    attack_speed,
                    dt,
                    &has_beam_query,
                );
            }
            EnemyState::ReturningWithCaptive => {
                if let Some(home) = original_pos {
                    let direction = home.position - transform.translation;
                    let return_speed = base_speed * RETURN_SPEED_MULTIPLIER;
                    if direction.length() < RETURN_ARRIVAL_DISTANCE {
                        transform.translation = home.position;
                        *state = EnemyState::InFormation;
                    } else {
                        transform.translation += direction.normalize() * return_speed * dt;
                    }
                }
            }
            _ => {}
        }

        if transform.translation.y < -WINDOW_HEIGHT / 2.0 - ENEMY_SIZE.y {
            commands.entity(entity).despawn();
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn step_attacker(
    commands: &mut Commands,
    entity: Entity,
    transform: &mut Transform,
    enemy_type: EnemyType,
    pattern: AttackPattern,
    attack_speed: f32,
    dt: f32,
    has_beam_query: &Query<&TractorBeam>,
) {
    match (enemy_type, pattern) {
        (EnemyType::Boss, AttackPattern::CaptureBeam) => {
            let hover_y = -WINDOW_HEIGHT * BOSS_CAPTURE_BEAM_HOVER_Y_FACTOR;
            if transform.translation.y > hover_y {
                transform.translation.y -= attack_speed * BOSS_DESCENT_SPEED_FACTOR * dt;
            } else if has_beam_query.get(entity).is_err() {
                commands.entity(entity).insert(TractorBeam {
                    target: None,
                    timer: Timer::new(
                        Duration::from_secs_f32(TRACTOR_BEAM_DURATION),
                        TimerMode::Once,
                    ),
                    width: TRACTOR_BEAM_WIDTH,
                    active: false,
                });
            }
        }
        (EnemyType::Boss, AttackPattern::SwoopDive) => {
            let target = Vec3::new(0.0, -WINDOW_HEIGHT / 2.0 - ENEMY_SIZE.y, 0.0);
            let direction = target - transform.translation;
            if direction.length() > 0.0 {
                transform.translation += direction.normalize() * attack_speed * dt;
            }
        }
        (_, AttackPattern::DirectDive) | (EnemyType::Grunt, AttackPattern::CaptureBeam) => {
            transform.translation.y -= attack_speed * dt;
        }
        (EnemyType::Grunt, AttackPattern::SwoopDive) => {
            transform.translation.y -= attack_speed * dt;
            let prev_x = transform.translation.x;
            let dx = -prev_x.signum() * attack_speed * SWOOP_HORIZONTAL_FACTOR * dt;
            transform.translation.x += dx;
            // Snap to center if we crossed it (prevents wobble)
            if prev_x != 0.0 && prev_x.signum() != transform.translation.x.signum() {
                transform.translation.x = 0.0;
            }
        }
        (_, AttackPattern::Kamikaze(target)) => {
            let direction = target - transform.translation;
            let dist = direction.length();
            if enemy_type == EnemyType::Boss && dist < KAMIKAZE_NEAR_DISTANCE {
                transform.translation.y -= attack_speed * dt;
            } else if dist > attack_speed * dt * ARRIVAL_THRESHOLD_FACTOR {
                transform.translation += direction.normalize() * attack_speed * dt;
            }
        }
    }
}

pub fn check_formation_complete(
    mut formation_state: ResMut<FormationState>,
    enemy_query: Query<&EnemyState, With<Enemy>>,
    mut attack_dive_timer: ResMut<AttackDiveTimer>,
    stage: Res<CurrentStage>,
    stage_configs: Res<StageConfigurations>,
) {
    if formation_state.formation_complete {
        return;
    }
    let config = stage_configs.for_stage(stage.number);
    if formation_state.total_spawned_this_stage != config.enemy_count {
        return;
    }
    let any_entering = enemy_query
        .iter()
        .any(|s| matches!(s, EnemyState::Entering));
    if any_entering {
        return;
    }

    info!("Stage {} formation complete", stage.number);
    formation_state.formation_complete = true;
    let dive_interval = Duration::from_secs_f32(config.attack_dive_interval);
    attack_dive_timer.timer.set_duration(dive_interval);
    attack_dive_timer.timer.reset();
    attack_dive_timer.timer.unpause();
}

pub fn trigger_attack_dives(
    mut timer: ResMut<AttackDiveTimer>,
    time: Res<Time>,
    mut enemy_query: Query<(Entity, &mut EnemyState, &Enemy)>,
    formation_state: Res<FormationState>,
    stage: Res<CurrentStage>,
    stage_configs: Res<StageConfigurations>,
    player_query: Query<&Transform, With<Player>>,
) {
    timer.timer.tick(time.delta());
    if !timer.timer.just_finished() || !formation_state.formation_complete {
        return;
    }

    let config = stage_configs.for_stage(stage.number);
    if config.attack_patterns.is_empty() {
        return;
    }

    let candidates: Vec<(Entity, EnemyType)> = enemy_query
        .iter()
        .filter(|(_, state, _)| matches!(state, EnemyState::InFormation))
        .map(|(e, _, enemy)| (e, enemy.enemy_type))
        .collect();
    if candidates.is_empty() {
        return;
    }

    let (chosen, enemy_type) = candidates[fastrand::usize(..candidates.len())];
    let pattern = pick_pattern(enemy_type, config.attack_patterns.as_slice(), &player_query);

    if let Ok((_, mut state, _)) = enemy_query.get_mut(chosen) {
        *state = EnemyState::Attacking(pattern);
    }
}

fn pick_pattern(
    enemy_type: EnemyType,
    available: &[AttackPattern],
    player_query: &Query<&Transform, With<Player>>,
) -> AttackPattern {
    let pattern = match enemy_type {
        EnemyType::Boss if fastrand::f32() < BOSS_CAPTURE_BEAM_PROBABILITY => {
            AttackPattern::CaptureBeam
        }
        _ => available[fastrand::usize(..available.len())],
    };

    // Resolve Kamikaze target (or fall back if player is gone).
    if matches!(pattern, AttackPattern::Kamikaze(_)) {
        match player_query.single() {
            Ok(t) => AttackPattern::Kamikaze(t.translation),
            Err(_) => AttackPattern::DirectDive,
        }
    } else {
        pattern
    }
}

pub fn enemy_shoot(
    mut commands: Commands,
    time: Res<Time>,
    mut enemy_query: Query<(&Transform, &mut Enemy, &EnemyState), Without<FormationTarget>>,
    stage: Res<CurrentStage>,
) {
    if is_special_stage(stage.number) {
        return;
    }
    for (transform, mut enemy, state) in enemy_query.iter_mut() {
        if !matches!(state, EnemyState::Attacking(_)) {
            continue;
        }
        enemy.shoot_cooldown.tick(time.delta());
        if !enemy.shoot_cooldown.is_finished() {
            continue;
        }

        let bullet_pos =
            transform.translation - Vec3::Y * (ENEMY_SIZE.y / 2.0 + ENEMY_BULLET_SIZE.y / 2.0);
        commands.spawn((
            Sprite {
                color: Color::srgb(1.0, 0.5, 0.5),
                custom_size: Some(ENEMY_BULLET_SIZE),
                ..default()
            },
            Transform::from_translation(bullet_pos),
            EnemyBullet,
        ));
        enemy.shoot_cooldown.reset();
    }
}

pub fn is_formation_complete(formation_state: Res<FormationState>) -> bool {
    formation_state.formation_complete
}

/// Beam spans from the boss down to the bottom of the screen.
pub(crate) fn calculate_beam_height(boss_y: f32) -> f32 {
    (boss_y - (-WINDOW_HEIGHT / 2.0)).max(0.0)
}

/// Sinusoidal alpha pulse, clamped to [0, 1].
pub(crate) fn beam_pulse_alpha(base_alpha: f32, time_secs: f32, freq: f32, amplitude: f32) -> f32 {
    (base_alpha + (time_secs * freq).sin() * amplitude).clamp(0.0, 1.0)
}

pub fn boss_capture_attack(
    mut commands: Commands,
    time: Res<Time>,
    mut boss_query: Query<(Entity, &Transform, &mut TractorBeam)>,
    player_query: Query<(Entity, &Transform), (With<Player>, Without<Captured>)>,
    enemy_query: Query<&EnemyState, With<Enemy>>,
    children_query: Query<&Children>,
) {
    for (boss_entity, boss_transform, mut beam) in boss_query.iter_mut() {
        beam.timer.tick(time.delta());

        if !beam.active && beam.target.is_none() {
            if let Ok((player_entity, _)) = player_query.single() {
                beam.target = Some(player_entity);
                beam.active = true;
                spawn_beam_visual(&mut commands, boss_entity, boss_transform.translation.y);
            }
        }

        if beam.active {
            if let Ok((player_entity, player_transform)) = player_query.single() {
                let dx = (player_transform.translation.x - boss_transform.translation.x).abs();
                if dx < beam.width / 2.0 {
                    info!("Player captured by boss {boss_entity:?}");
                    commands.entity(player_entity).insert(Captured {
                        boss_entity,
                        timer: Timer::new(
                            Duration::from_secs_f32(CAPTURE_DURATION),
                            TimerMode::Once,
                        ),
                    });
                    end_beam(&mut commands, boss_entity, &enemy_query, &children_query);
                    continue;
                }
            }
        }

        if beam.timer.is_finished() {
            end_beam(&mut commands, boss_entity, &enemy_query, &children_query);
        }
    }
}

fn spawn_beam_visual(commands: &mut Commands, boss_entity: Entity, boss_y: f32) {
    let beam_height = calculate_beam_height(boss_y);
    let core_height = beam_height * 0.6;
    commands.entity(boss_entity).with_children(|parent| {
        parent.spawn((
            Sprite {
                color: BEAM_GLOW_COLOR,
                custom_size: Some(Vec2::new(BEAM_GLOW_WIDTH, beam_height)),
                ..default()
            },
            Transform::from_xyz(0.0, -beam_height / 2.0, -0.5),
            TractorBeamSprite,
        ));
        parent.spawn((
            Sprite {
                color: BEAM_CORE_COLOR,
                custom_size: Some(Vec2::new(TRACTOR_BEAM_WIDTH, core_height)),
                ..default()
            },
            Transform::from_xyz(0.0, -core_height / 2.0 - beam_height * 0.2, -0.3),
            TractorBeamSprite,
        ));
    });
}

/// Tear down the beam component and visuals; if the boss was attacking, send it home.
fn end_beam(
    commands: &mut Commands,
    boss_entity: Entity,
    enemy_query: &Query<&EnemyState, With<Enemy>>,
    children_query: &Query<&Children>,
) {
    commands.entity(boss_entity).remove::<TractorBeam>();
    if let Ok(children) = children_query.get(boss_entity) {
        for child in children.iter() {
            commands.entity(child).despawn();
        }
    }
    if let Ok(EnemyState::Attacking(_)) = enemy_query.get(boss_entity) {
        commands
            .entity(boss_entity)
            .insert(EnemyState::ReturningWithCaptive);
    }
}

pub fn update_tractor_beam_visual(
    time: Res<Time>,
    boss_query: Query<(Entity, &Transform), (With<TractorBeam>, Without<TractorBeamSprite>)>,
    mut sprite_query: Query<(&mut Sprite, &mut Transform), With<TractorBeamSprite>>,
    children_query: Query<&Children>,
) {
    let time_secs = time.elapsed_secs();
    for (boss_entity, boss_transform) in boss_query.iter() {
        let beam_height = calculate_beam_height(boss_transform.translation.y);
        let Ok(children) = children_query.get(boss_entity) else {
            continue;
        };

        for &child in children {
            let Ok((mut sprite, mut transform)) = sprite_query.get_mut(child) else {
                continue;
            };

            let alpha = beam_pulse_alpha(
                sprite.color.alpha(),
                time_secs,
                BEAM_PULSE_FREQ,
                BEAM_PULSE_AMPLITUDE,
            );
            sprite.color.set_alpha(alpha);

            // Distinguish glow vs. core by their fixed widths.
            let width = sprite.custom_size.map(|s| s.x).unwrap_or(0.0);
            if (width - BEAM_GLOW_WIDTH).abs() < f32::EPSILON {
                sprite.custom_size = Some(Vec2::new(width, beam_height));
                transform.translation.y = -beam_height / 2.0;
            } else if (width - TRACTOR_BEAM_WIDTH).abs() < f32::EPSILON {
                let core_height = beam_height * 0.6;
                sprite.custom_size = Some(Vec2::new(width, core_height));
                transform.translation.y = -core_height / 2.0 - beam_height * 0.2;
            }
        }
    }
}

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

    #[test]
    fn beam_height_at_screen_center() {
        assert_eq!(calculate_beam_height(0.0), WINDOW_HEIGHT / 2.0);
    }

    #[test]
    fn beam_height_at_boss_position() {
        assert_eq!(calculate_beam_height(-200.0), 200.0);
    }

    #[test]
    fn beam_height_clamped_at_bottom() {
        assert_eq!(calculate_beam_height(-WINDOW_HEIGHT / 2.0), 0.0);
    }

    #[test]
    fn pulse_at_zero_returns_base() {
        assert_eq!(beam_pulse_alpha(0.6, 0.0, 3.0, 0.15), 0.6);
    }

    #[test]
    fn pulse_at_peak() {
        let t = std::f32::consts::FRAC_PI_2 / 3.0;
        let result = beam_pulse_alpha(0.5, t, 3.0, 0.15);
        assert!((result - 0.65).abs() < 0.001);
    }

    #[test]
    fn pulse_at_trough() {
        let t = std::f32::consts::FRAC_PI_2;
        let result = beam_pulse_alpha(0.5, t, 3.0, 0.15);
        assert!((result - 0.35).abs() < 0.001);
    }

    #[test]
    fn pulse_clamped_to_zero() {
        let t = 3.0 * std::f32::consts::FRAC_PI_2 / 3.0;
        let result = beam_pulse_alpha(0.05, t, 3.0, 1.0);
        assert_eq!(result, 0.0);
    }

    #[test]
    fn pulse_clamped_to_one() {
        let t = std::f32::consts::FRAC_PI_2 / 3.0;
        let result = beam_pulse_alpha(0.95, t, 3.0, 1.0);
        assert_eq!(result, 1.0);
    }
}