use bevy::prelude::*; // Removed unused AppExit
use std::time::Duration;
// --- Game States ---
#[derive(Clone, Eq, PartialEq, Debug, Hash, Default, States)]
enum AppState {
    #[default]
    Playing,
    GameOver,
}
// --- Constants ---
const WINDOW_WIDTH: f32 = 600.0;
const WINDOW_HEIGHT: f32 = 800.0;
const PLAYER_SPEED: f32 = 300.0;
const BULLET_SPEED: f32 = 500.0;
const ENEMY_SPEED: f32 = 100.0;
const PLAYER_SIZE: Vec2 = Vec2::new(30.0, 30.0);
const ENEMY_SIZE: Vec2 = Vec2::new(40.0, 40.0);
const BULLET_SIZE: Vec2 = Vec2::new(5.0, 15.0); // Player bullet
const ENEMY_BULLET_SIZE: Vec2 = Vec2::new(8.0, 8.0); // Enemy bullet
const ENEMY_BULLET_SPEED: f32 = 300.0;
const ENEMY_SHOOT_INTERVAL: f32 = 1.5; // Seconds between shots for an attacking enemy

// Formation constants
const FORMATION_ROWS: usize = 4;
const FORMATION_COLS: usize = 8;
const FORMATION_ENEMY_COUNT: usize = FORMATION_ROWS * FORMATION_COLS;
const FORMATION_X_SPACING: f32 = 60.0;
const FORMATION_Y_SPACING: f32 = 50.0;
const FORMATION_BASE_Y: f32 = WINDOW_HEIGHT / 2.0 - 150.0; // Top area for formation
const STARTING_LIVES: u32 = 3;
const PLAYER_RESPAWN_DELAY: f32 = 2.0;
const PLAYER_INVINCIBILITY_DURATION: f32 = 2.0;

// Collision thresholds
const BULLET_ENEMY_COLLISION_THRESHOLD: f32 = (BULLET_SIZE.x + ENEMY_SIZE.x) * 0.5; // 22.5
const PLAYER_ENEMY_COLLISION_THRESHOLD: f32 = (PLAYER_SIZE.x + ENEMY_SIZE.x) * 0.5; // 35.0
const ENEMY_BULLET_PLAYER_COLLISION_THRESHOLD: f32 = (ENEMY_BULLET_SIZE.x + PLAYER_SIZE.x) * 0.5; // ~19.0

// --- Components ---
#[derive(Component)]
struct Player {
    speed: f32,
    shoot_cooldown: Timer,
}

#[derive(Component)]
struct Bullet;

#[derive(Component)]
struct Enemy {
    shoot_cooldown: Timer,
}

#[derive(Component)]
struct Invincible {
    timer: Timer,
}

#[derive(Component)]
struct FormationTarget {
    position: Vec3,
}

#[derive(Component, Clone, PartialEq)]
enum EnemyState {
    Entering,    // Flying onto the screen towards formation target
    InFormation, // Holding position in the formation
    Attacking,   // Diving towards the player
}

#[derive(Component)]
struct EnemyBullet;
// --- Resources ---
#[derive(Resource)]
struct EnemySpawnTimer {
    timer: Timer,
}

#[derive(Resource)]
struct PlayerLives {
    count: u32,
}

#[derive(Resource)]
struct PlayerRespawnTimer {
    timer: Timer,
}
#[derive(Resource)]
struct Score {
    value: u32,
}

#[derive(Resource)]
struct CurrentStage {
    number: u32,
    waiting_for_clear: bool, // Flag to check if we should check for stage clear
}

#[derive(Resource)]
struct FormationState {
    next_slot_index: usize,
    total_spawned_this_stage: usize,
    formation_complete: bool, // Flag to indicate if all enemies are in position
}

#[derive(Resource)]
struct AttackDiveTimer {
    timer: Timer,
}
// --- Game Over UI ---

#[derive(Component)]
struct GameOverUI;

fn setup_game_over_ui(mut commands: Commands) {
    println!("Entering GameOver state. Setting up UI.");
    commands.spawn((
        TextBundle::from_section(
            "GAME OVER",
            TextStyle {
                font_size: 100.0,
                color: Color::WHITE,
                ..default()
            },
        )
        .with_style(Style {
            position_type: PositionType::Absolute,
            align_self: AlignSelf::Center,
            justify_self: JustifySelf::Center,
            top: Val::Percent(40.0), // Center vertically roughly
            ..default()
        }),
        GameOverUI, // Tag the UI element
    ));
    // TODO: Add "Press R to Restart" text later
}

fn cleanup_game_over_ui(mut commands: Commands, query: Query<Entity, With<GameOverUI>>) {
    println!("Exiting GameOver state. Cleaning up UI.");
    for entity in query.iter() {
        commands.entity(entity).despawn_recursive();
    }
}

// --- Cleanup ---

fn cleanup_game_entities(
    mut commands: Commands,
    bullet_query: Query<Entity, With<Bullet>>,
    enemy_query: Query<Entity, With<Enemy>>,
    // Optionally despawn player too, or handle separately if needed for restart
    // player_query: Query<Entity, With<Player>>,
) {
    println!("Exiting Playing state. Cleaning up game entities.");
    for entity in bullet_query.iter() {
        commands.entity(entity).despawn();
    }
    for entity in enemy_query.iter() {
        commands.entity(entity).despawn();
    }
    // for entity in player_query.iter() {
    //     commands.entity(entity).despawn();
    // }
}

// --- Stage Management ---

// Helper to access world directly in check_stage_clear
fn check_stage_clear(world: &mut World) {
    // Use manual resource access because we need mutable access to multiple resources
    // Separate checks to manage borrows correctly
    let mut should_clear = false;
    if let Some(stage) = world.get_resource::<CurrentStage>() {
        if stage.waiting_for_clear {
            // Create the query *after* checking the flag
            let mut enemy_query = world.query_filtered::<Entity, With<Enemy>>();
            if enemy_query.iter(world).next().is_none() {
                should_clear = true;
            }
        }
    }

    if should_clear {
        // Get mutable resources only when needed
        if let Some(mut stage) = world.get_resource_mut::<CurrentStage>() {
            stage.number += 1;
            stage.waiting_for_clear = false;
            println!("Stage cleared! Starting Stage {}...", stage.number);
        }
        if let Some(mut formation_state) = world.get_resource_mut::<FormationState>() {
            formation_state.next_slot_index = 0;
            formation_state.total_spawned_this_stage = 0;
            formation_state.formation_complete = false; // Reset flag for new stage
        }
        if let Some(mut spawn_timer) = world.get_resource_mut::<EnemySpawnTimer>() {
            spawn_timer.timer.reset();
        }
    }
}

// --- Enemy Attack Logic ---

fn trigger_attack_dives(
    mut timer: ResMut<AttackDiveTimer>,
    time: Res<Time>,
    mut query: Query<(Entity, &mut EnemyState), With<Enemy>>,
    formation_state: Res<FormationState>, // Read formation state
) {
    timer.timer.tick(time.delta());

    // Timer is ticked implicitly by the system scheduler when the run condition is met,
    // or we can tick it explicitly *after* checking formation_complete.
    // Let's rely on the run condition and internal check for now.
    // timer.timer.tick(time.delta()); // Removed redundant tick

    // Only proceed if the timer finished AND the formation is complete
    if timer.timer.just_finished() && formation_state.formation_complete {
        // Find all enemies currently in formation
        let mut available_enemies: Vec<Entity> = Vec::new();
        for (entity, state) in query.iter() {
            if *state == EnemyState::InFormation {
                available_enemies.push(entity);
            }
        }

        // If there are enemies available, pick one randomly
        if !available_enemies.is_empty() {
            let random_index = fastrand::usize(..available_enemies.len());
            let chosen_entity = available_enemies[random_index];

            // Get the chosen enemy's state mutably and change it
            if let Ok((_, mut state)) = query.get_mut(chosen_entity) {
                println!("Enemy {:?} starting attack dive!", chosen_entity);
                *state = EnemyState::Attacking;
                // Timer will automatically repeat due to TimerMode::Repeating,
                // no need to reset here unless we wanted a different interval after the first dive.
            }
        }
    }
}

fn enemy_shoot(
    mut commands: Commands,
    time: Res<Time>,
    // Query attacking enemies: need their transform and mutable Enemy component for the timer
    mut enemy_query: Query<(&Transform, &mut Enemy, &EnemyState), Without<FormationTarget>>,
) {
    for (transform, mut enemy, state) in enemy_query.iter_mut() {
        // Only shoot if attacking
        if *state == EnemyState::Attacking {
            enemy.shoot_cooldown.tick(time.delta());
            if enemy.shoot_cooldown.finished() {
                println!("Enemy {:?} firing!", transform.translation); // Placeholder for entity ID if needed
                let bullet_start_pos = transform.translation
                    - Vec3::Y * (ENEMY_SIZE.y / 2.0 + ENEMY_BULLET_SIZE.y / 2.0);

                commands.spawn((
                    SpriteBundle {
                        sprite: Sprite {
                            color: Color::rgb(1.0, 0.5, 0.5), // Different color for enemy bullets
                            custom_size: Some(ENEMY_BULLET_SIZE),
                            ..default()
                        },
                        transform: Transform::from_translation(bullet_start_pos),
                        ..default()
                    },
                    EnemyBullet,
                ));

                // Reset the timer for the next shot
                enemy.shoot_cooldown.reset();
            }
        }
    }
}

fn move_enemy_bullets(
    mut commands: Commands,
    time: Res<Time>,
    mut query: Query<(Entity, &mut Transform), With<EnemyBullet>>,
) {
    for (entity, mut transform) in query.iter_mut() {
        transform.translation.y -= ENEMY_BULLET_SPEED * time.delta_seconds();

        // Despawn if off screen (bottom)
        if transform.translation.y < -WINDOW_HEIGHT / 2.0 - ENEMY_BULLET_SIZE.y {
            commands.entity(entity).despawn();
        }
    }
}

// System to check if all spawned enemies have reached their formation position
fn check_formation_complete(
    mut formation_state: ResMut<FormationState>,
    enemy_query: Query<&EnemyState, With<Enemy>>, // Query states directly
    mut attack_dive_timer: ResMut<AttackDiveTimer>, // Add timer to reset it
) {
    // Only run the check if the formation isn't already marked as complete
    if !formation_state.formation_complete {
        // Check if all enemies for the stage have been spawned
        if formation_state.total_spawned_this_stage == FORMATION_ENEMY_COUNT {
            // Check if any enemies are still in the Entering state
            let mut any_entering = false;
            for state in enemy_query.iter() {
                if *state == EnemyState::Entering {
                    any_entering = true;
                    break;
                }
            }

            // If all spawned and none are entering, formation is complete
            if !any_entering {
                println!(
                    "Formation complete! Resetting attack timer. (Spawned={})",
                    formation_state.total_spawned_this_stage
                );
                formation_state.formation_complete = true;
                attack_dive_timer.timer.reset(); // Reset the dive timer
                attack_dive_timer.timer.unpause(); // Start the timer now
                println!("Formation complete! Attack timer unpaused and reset.");
            }
            // else { // Optional log
            //     println!("Checking formation complete: Still entering...");
            // }
        }
        // else { // Optional log
        //     println!("Checking formation complete: Not all spawned yet ({}/{})", formation_state.total_spawned_this_stage, FORMATION_ENEMY_COUNT);
        // }
    }
}

// Check collisions between enemy bullets and the player
fn check_enemy_bullet_player_collisions(
    mut commands: Commands,
    mut lives: ResMut<PlayerLives>,
    mut respawn_timer: ResMut<PlayerRespawnTimer>,
    mut next_state: ResMut<NextState<AppState>>,
    // Player query matching the run_if condition
    player_query: Query<(Entity, &Transform), (With<Player>, Without<Invincible>)>,
    enemy_bullet_query: Query<(Entity, &Transform), With<EnemyBullet>>,
) {
    if let Ok((player_entity, player_transform)) = player_query.get_single() {
        for (bullet_entity, bullet_transform) in enemy_bullet_query.iter() {
            let distance = player_transform
                .translation
                .distance(bullet_transform.translation);

            if distance < ENEMY_BULLET_PLAYER_COLLISION_THRESHOLD {
                println!("Player hit by enemy bullet!");
                commands.entity(bullet_entity).despawn(); // Despawn bullet

                lives.count = lives.count.saturating_sub(1);
                println!("Lives remaining: {}", lives.count);

                commands.entity(player_entity).despawn(); // Despawn player

                if lives.count > 0 {
                    respawn_timer.timer.reset();
                    respawn_timer.timer.unpause();
                    println!("Respawn timer started.");
                } else {
                    println!("GAME OVER!");
                    next_state.set(AppState::GameOver);
                }
                // Break because the player can only be hit once per frame
                break;
            }
        }
    }
}

fn main() {
    App::new() // Start App builder
        .init_state::<AppState>() // Initialize the AppState *after* App::new()
        .insert_resource(ClearColor(Color::rgb(0.0, 0.0, 0.0)))
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "Galaga :: Stage: 1 Lives: 3 Score: 0".into(), // Initial title with stage
                resolution: (WINDOW_WIDTH, WINDOW_HEIGHT).into(),
                ..default()
            }),
            ..default()
        }))
        // Add Resources
        .insert_resource(PlayerLives {
            count: STARTING_LIVES,
        })
        .insert_resource(PlayerRespawnTimer {
            // Start paused and finished
            timer: Timer::new(
                Duration::from_secs_f32(PLAYER_RESPAWN_DELAY),
                TimerMode::Once,
            ),
        })
        .insert_resource(EnemySpawnTimer {
            timer: Timer::new(Duration::from_secs_f32(0.5), TimerMode::Repeating),
        })
        .insert_resource(Score { value: 0 }) // Initialize score
        .insert_resource(CurrentStage {
            number: 1,
            waiting_for_clear: false,
        }) // Initialize stage and flag
        .insert_resource(FormationState {
            next_slot_index: 0,
            total_spawned_this_stage: 0,
            formation_complete: false,
        }) // Initialize formation state with flag
        .insert_resource(AttackDiveTimer {
            timer: {
                let mut timer = Timer::new(Duration::from_secs_f32(3.0), TimerMode::Repeating);
                timer.pause(); // Start paused until formation is complete
                timer
            },
        }) // Initialize attack timer (e.g., every 3 seconds), initially paused
        // Add Systems
        .add_systems(Startup, setup)
        // Systems running only when Playing
        .add_systems(
            Update,
            (
                move_player,
                player_shoot.run_if(player_exists),
                move_bullets,
                move_enemies,
                spawn_enemies,
                check_bullet_collisions,
                check_player_enemy_collisions.run_if(player_vulnerable),
                respawn_player.run_if(should_respawn_player),
                manage_invincibility,
                // check_formation_complete and trigger_attack_dives added separately below
                enemy_shoot, // Add enemy shooting system
                move_enemy_bullets,
                check_enemy_bullet_player_collisions.run_if(player_vulnerable), // Add enemy bullet collision check
                                                                                // check_stage_clear is now added separately below using .pipe(system)
                                                                                // Game Over check is now implicit in check_player_enemy_collisions
            )
                .run_if(in_state(AppState::Playing)),
        )
        // Add formation check, dive trigger (with run condition), and stage clear separately
        .add_systems(
            Update,
            check_formation_complete.run_if(in_state(AppState::Playing)),
        )
        .add_systems(
            Update,
            trigger_attack_dives
                .run_if(in_state(AppState::Playing).and_then(is_formation_complete)),
        )
        .add_systems(
            Update,
            check_stage_clear.run_if(in_state(AppState::Playing)),
        ) // check_stage_clear uses world access
        // Systems running regardless of state (or managing state transitions)
        .add_systems(
            Update,
            (
                update_window_title, // Keep title updated
                // Add system to check for restart input in GameOver state later
                bevy::window::close_on_esc, // Allow closing anytime
            ),
        )
        // Systems for entering/exiting states
        .add_systems(OnEnter(AppState::GameOver), setup_game_over_ui)
        .add_systems(OnExit(AppState::GameOver), cleanup_game_over_ui)
        .add_systems(OnExit(AppState::Playing), cleanup_game_entities) // Cleanup when leaving Playing
        .run();
}

// --- Run Conditions ---
fn player_exists(query: Query<&Player>) -> bool {
    !query.is_empty()
}

fn player_vulnerable(query: Query<&Player, Without<Invincible>>) -> bool {
    !query.is_empty()
}

fn should_respawn_player(lives: Res<PlayerLives>, player_query: Query<&Player>) -> bool {
    player_query.is_empty() && lives.count > 0
}

// New run condition: Check if the formation is complete
fn is_formation_complete(formation_state: Res<FormationState>) -> bool {
    formation_state.formation_complete
}

// --- Systems ---

fn setup(mut commands: Commands) {
    commands.spawn(Camera2dBundle::default());
    spawn_player_ship(&mut commands); // Use helper function
}

// Helper to spawn player (used in setup and respawn)
fn spawn_player_ship(commands: &mut Commands) {
    commands.spawn((
        SpriteBundle {
            sprite: Sprite {
                color: Color::rgb(0.0, 0.5, 1.0),
                custom_size: Some(PLAYER_SIZE),
                ..default()
            },
            transform: Transform::from_translation(Vec3::new(
                0.0,
                -WINDOW_HEIGHT / 2.0 + PLAYER_SIZE.y / 2.0 + 20.0,
                0.0,
            )),
            ..default()
        },
        Player {
            speed: PLAYER_SPEED,
            shoot_cooldown: Timer::new(Duration::from_secs_f32(0.3), TimerMode::Once),
        },
        // Player starts invincible for a short time
        Invincible {
            timer: Timer::new(
                Duration::from_secs_f32(PLAYER_INVINCIBILITY_DURATION),
                TimerMode::Once,
            ),
        },
    ));
    println!("Player spawned!");
}

fn move_player(
    keyboard_input: Res<ButtonInput<KeyCode>>,
    mut query: Query<(&mut Transform, &Player)>,
    time: Res<Time>,
) {
    // Using get_single_mut handles the case where player might not exist yet (or was just destroyed)
    if let Ok((mut transform, player)) = query.get_single_mut() {
        let mut direction = 0.0;

        if keyboard_input.pressed(KeyCode::KeyA) || keyboard_input.pressed(KeyCode::ArrowLeft) {
            direction -= 1.0;
        }

        if keyboard_input.pressed(KeyCode::KeyD) || keyboard_input.pressed(KeyCode::ArrowRight) {
            direction += 1.0;
        }

        transform.translation.x += direction * player.speed * time.delta_seconds();
        let half_player_width = PLAYER_SIZE.x / 2.0;
        transform.translation.x = transform.translation.x.clamp(
            -WINDOW_WIDTH / 2.0 + half_player_width,
            WINDOW_WIDTH / 2.0 - half_player_width,
        );
    }
}

fn player_shoot(
    keyboard_input: Res<ButtonInput<KeyCode>>,
    mut query: Query<(&Transform, &mut Player)>, // Should only run if player exists due to run_if
    mut commands: Commands,
    time: Res<Time>,
) {
    if let Ok((player_transform, mut player)) = query.get_single_mut() {
        player.shoot_cooldown.tick(time.delta());

        if (keyboard_input.just_pressed(KeyCode::Space)
            || keyboard_input.just_pressed(KeyCode::ArrowUp))
            && player.shoot_cooldown.finished()
        {
            player.shoot_cooldown.reset();

            let bullet_start_pos = player_transform.translation
                + Vec3::Y * (PLAYER_SIZE.y / 2.0 + BULLET_SIZE.y / 2.0);

            commands.spawn((
                SpriteBundle {
                    sprite: Sprite {
                        color: Color::rgb(1.0, 1.0, 1.0),
                        custom_size: Some(BULLET_SIZE),
                        ..default()
                    },
                    transform: Transform::from_translation(bullet_start_pos),
                    ..default()
                },
                Bullet,
            ));
        }
    }
}

fn move_bullets(
    mut query: Query<(Entity, &mut Transform), With<Bullet>>,
    time: Res<Time>,
    mut commands: Commands,
) {
    for (entity, mut transform) in query.iter_mut() {
        transform.translation.y += BULLET_SPEED * time.delta_seconds();

        if transform.translation.y > WINDOW_HEIGHT / 2.0 + BULLET_SIZE.y {
            commands.entity(entity).despawn();
        }
    }
}

fn spawn_enemies(
    mut commands: Commands,
    time: Res<Time>,
    mut timer: ResMut<EnemySpawnTimer>,
    mut stage: ResMut<CurrentStage>,
    mut formation_state: ResMut<FormationState>, // Add formation state resource
) {
    // Tick the timer every frame
    timer.timer.tick(time.delta());

    // Only spawn if we haven't spawned the full formation for this stage yet
    // AND the timer just finished this frame
    if formation_state.total_spawned_this_stage < FORMATION_ENEMY_COUNT
        && timer.timer.just_finished()
    {
        // Calculate grid position
        let slot_index = formation_state.next_slot_index;
        let row = slot_index / FORMATION_COLS;
        let col = slot_index % FORMATION_COLS;

        // Calculate target position in formation
        let target_x = (col as f32 - (FORMATION_COLS as f32 - 1.0) / 2.0) * FORMATION_X_SPACING;
        let target_y = FORMATION_BASE_Y - (row as f32 * FORMATION_Y_SPACING);
        let target_pos = Vec3::new(target_x, target_y, 0.0);

        // Spawn position (still random at the top for now)
        let spawn_x = (fastrand::f32() - 0.5) * (WINDOW_WIDTH - ENEMY_SIZE.x);
        let spawn_y = WINDOW_HEIGHT / 2.0 - ENEMY_SIZE.y / 2.0;

        commands.spawn((
            SpriteBundle {
                sprite: Sprite {
                    color: Color::rgb(1.0, 0.2, 0.2),
                    custom_size: Some(ENEMY_SIZE),
                    ..default()
                },
                transform: Transform::from_translation(Vec3::new(spawn_x, spawn_y, 0.0)),
                ..default()
            },
            Enemy {
                // Initialize cooldown, maybe slightly randomized later
                shoot_cooldown: Timer::new(
                    Duration::from_secs_f32(ENEMY_SHOOT_INTERVAL),
                    TimerMode::Once,
                ),
            },
            FormationTarget {
                position: target_pos,
            },
            EnemyState::Entering, // Start in Entering state
        ));

        formation_state.next_slot_index = (slot_index + 1) % FORMATION_ENEMY_COUNT; // Cycle through slots if needed, though we stop spawning
        formation_state.total_spawned_this_stage += 1;

        // Mark that we are now waiting for enemies to be cleared
        stage.waiting_for_clear = true;

        // Reset timer only if we are still spawning more enemies for the formation
        if formation_state.total_spawned_this_stage < FORMATION_ENEMY_COUNT {
            timer.timer.reset();
        } else {
            println!("Full formation spawned for Stage {}", stage.number);
        }
    }
}

fn move_enemies(
    // Query now includes FormationTarget
    // Query now includes EnemyState, but FormationTarget is optional
    // We need separate queries or logic paths for different states
    mut entering_query: Query<
        (Entity, &mut Transform, &FormationTarget, &mut EnemyState),
        (With<Enemy>, With<FormationTarget>),
    >,
    mut attacking_query: Query<
        (Entity, &mut Transform, &EnemyState), // Add &EnemyState here
        (With<Enemy>, Without<FormationTarget>), // Keep With<EnemyState> filter if desired, but we check explicitly below
    >, // Query potential attackers
    time: Res<Time>,
    mut commands: Commands,
    stage: Res<CurrentStage>, // Add stage resource for speed calculation
) {
    // Increase speed based on stage number
    let current_speed = ENEMY_SPEED + (stage.number - 1) as f32 * 10.0;
    let arrival_threshold = current_speed * time.delta_seconds() * 1.1; // Slightly more than one frame's movement

    // --- Handle Entering Enemies ---
    for (entity, mut transform, target, mut state) in entering_query.iter_mut() {
        // Ensure we only process Entering state here, though query filters it mostly
        if *state == EnemyState::Entering {
            let current_pos = transform.translation;
            let target_pos = target.position;
            let direction = target_pos - current_pos;
            let distance = direction.length();

            if distance < arrival_threshold {
                // Arrived at target
                transform.translation = target_pos;
                commands.entity(entity).remove::<FormationTarget>(); // Remove target component
                *state = EnemyState::InFormation; // Change state
                println!(
                    "Enemy {:?} reached formation target and is now InFormation.",
                    entity
                );
            } else {
                // Move towards target
                let move_delta = direction.normalize() * current_speed * time.delta_seconds();
                transform.translation += move_delta;
            }
        }
    }

    // --- Handle Attacking Enemies ---
    let attack_speed = current_speed * 1.5; // Make them dive faster
    for (entity, mut transform, state) in attacking_query.iter_mut() { // Get state from query
        // *** Explicitly check if the enemy is actually in the Attacking state ***
        if *state == EnemyState::Attacking {
            // Move straight down ONLY if attacking
            transform.translation.y -= attack_speed * time.delta_seconds();
        }
        // Enemies that are InFormation but Without<FormationTarget> will now be ignored by this movement logic.

        // Despawn if off screen
        if transform.translation.y < -WINDOW_HEIGHT / 2.0 - ENEMY_SIZE.y {
            println!(
                "Despawning attacking enemy {:?} that went off screen.",
                entity
            );
            commands.entity(entity).despawn();
            // TODO: Later, attacking enemies might return to formation or loop
        }
    }
}

fn check_bullet_collisions(
    mut commands: Commands,
    bullet_query: Query<(Entity, &Transform), With<Bullet>>,
    enemy_query: Query<(Entity, &Transform), With<Enemy>>,
    mut score: ResMut<Score>, // Add Score resource
) {
    for (bullet_entity, bullet_transform) in bullet_query.iter() {
        for (enemy_entity, enemy_transform) in enemy_query.iter() {
            let distance = bullet_transform
                .translation
                .distance(enemy_transform.translation);

            if distance < BULLET_ENEMY_COLLISION_THRESHOLD {
                commands.entity(bullet_entity).despawn();
                commands.entity(enemy_entity).despawn();
                score.value += 100; // Increment score
                println!("Enemy hit! Score: {}", score.value); // Log score update
                break;
            }
        }
    }
}

// Modified Collision Check for Player
fn check_player_enemy_collisions(
    mut commands: Commands,
    mut lives: ResMut<PlayerLives>,
    mut respawn_timer: ResMut<PlayerRespawnTimer>,
    mut next_state: ResMut<NextState<AppState>>, // Resource to change state
    // Query player without Invincible component - relies on run_if condition too
    player_query: Query<(Entity, &Transform), (With<Player>, Without<Invincible>)>,
    enemy_query: Query<(Entity, &Transform), With<Enemy>>,
) {
    // This system only runs if player exists and is not invincible, due to run_if
    if let Ok((player_entity, player_transform)) = player_query.get_single() {
        for (enemy_entity, enemy_transform) in enemy_query.iter() {
            let distance = player_transform
                .translation
                .distance(enemy_transform.translation);

            if distance < PLAYER_ENEMY_COLLISION_THRESHOLD {
                println!("Player hit!");
                commands.entity(enemy_entity).despawn(); // Despawn enemy

                lives.count = lives.count.saturating_sub(1); // Decrement lives safely
                println!("Lives remaining: {}", lives.count);

                commands.entity(player_entity).despawn(); // Despawn player

                if lives.count > 0 {
                    // Start the respawn timer
                    respawn_timer.timer.reset();
                    respawn_timer.timer.unpause();
                    println!("Respawn timer started.");
                } else {
                    println!("GAME OVER!");
                    next_state.set(AppState::GameOver); // Transition to GameOver state
                }
                // Important: Break after handling one collision per frame for the player
                break;
            }
        }
    }
}

// New System: Respawn Player
fn respawn_player(
    mut commands: Commands,
    time: Res<Time>,
    mut respawn_timer: ResMut<PlayerRespawnTimer>,
    // No player query needed here due to run_if condition
) {
    // Tick the timer only if it's actually running
    if respawn_timer.timer.tick(time.delta()).just_finished() {
        println!("Respawn timer finished. Spawning player.");
        spawn_player_ship(&mut commands);
        respawn_timer.timer.pause(); // Pause timer until next death
    }
}

// New System: Manage Invincibility
fn manage_invincibility(
    mut commands: Commands,
    time: Res<Time>,
    mut query: Query<(Entity, &mut Invincible, Option<&mut Visibility>), With<Player>>,
) {
    for (entity, mut invincible, mut visibility) in query.iter_mut() {
        invincible.timer.tick(time.delta());

        // Blinking effect (optional)
        if let Some(ref mut vis) = visibility {
            // Blink roughly 5 times per second
            let elapsed_secs = invincible.timer.elapsed_secs();
            **vis = if (elapsed_secs * 10.0).floor() % 2.0 == 0.0 {
                Visibility::Visible
            } else {
                Visibility::Hidden
            };
        }

        if invincible.timer.finished() {
            println!("Invincibility finished.");
            commands.entity(entity).remove::<Invincible>();
            // Ensure player is visible when invincibility ends
            if let Some(ref mut vis) = visibility {
                **vis = Visibility::Visible;
            }
        }
    }
}

// New System: Update Window Title with Lives
fn update_window_title(
    lives: Res<PlayerLives>,
    score: Res<Score>,
    stage: Res<CurrentStage>, // Add CurrentStage resource
    mut windows: Query<&mut Window>,
) {
    // Update if lives or score changed
    // Update if lives, score, or stage changed
    if lives.is_changed() || score.is_changed() || stage.is_changed() {
        let mut window = windows.single_mut();
        window.title = format!(
            "Galaga :: Stage: {} Lives: {} Score: {}",
            stage.number, lives.count, score.value
        );
    }
}