snake/src/snake.rs

use bevy::prelude::*;
use rand::prelude::random;

use crate::components as c;

fn spawn_segment(
    mut commands: Commands,
    material: &Handle<ColorMaterial>,
    position: c::Position,
) -> Entity {
    commands
        .spawn_bundle(SpriteBundle {
            material: material.clone(),
            ..Default::default()
        })
        .insert(c::SnakeSegment)
        .insert(position)
        .insert(c::GridSize::square(0.65))
        .id()
}

pub fn spawn(
    mut commands: Commands,
    materials: Res<c::Materials>,
    mut segments: ResMut<c::SnakeSegments>,
) {
    segments.segments = vec![
        commands
            .spawn_bundle(SpriteBundle {
                material: materials.head_material.clone(),
                sprite: Sprite::new(Vec2::new(10.0, 10.0)),
                ..Default::default()
            })
            .insert(c::SnakeHead {
                direction: c::Direction::Up,
                intention: c::Direction::Up,
            })
            .insert(c::SnakeSegment)
            .insert(c::Position { x: 3, y: 3 })
            .insert(c::GridSize::square(0.8))
            .id(),
        spawn_segment(
            commands,
            &materials.segment_material,
            c::Position { x: 3, y: 2 },
        ),
    ]
}

pub fn input(keyboard_input: Res<Input<KeyCode>>, mut heads: Query<&mut c::SnakeHead>) {
    if let Some(mut head) = heads.iter_mut().next() {
        let dir = if keyboard_input.pressed(KeyCode::Left) {
            c::Direction::Left
        } else if keyboard_input.pressed(KeyCode::Down) {
            c::Direction::Down
        } else if keyboard_input.pressed(KeyCode::Up) {
            c::Direction::Up
        } else if keyboard_input.pressed(KeyCode::Right) {
            c::Direction::Right
        } else {
            head.intention
        };
        if dir != head.direction.opposite() {
            head.intention = dir;
        }
    }
}

pub fn movement(
    segments: ResMut<c::SnakeSegments>,
    mut last_tail_position: ResMut<c::LastTailPosition>,
    mut game_over_writer: EventWriter<c::GameOverEvent>,
    mut heads: Query<(Entity, &mut c::SnakeHead)>,
    mut positions: Query<&mut c::Position>,
) {
    if let Some((head_entity, mut head)) = heads.iter_mut().next() {
        let segment_positions = segments
            .segments
            .iter()
            .map(|e| *positions.get_mut(*e).unwrap())
            .collect::<Vec<c::Position>>();
        let mut head_pos = positions.get_mut(head_entity).unwrap();
        head.direction = head.intention;
        match &head.direction {
            c::Direction::Left => head_pos.x -= 1,
            c::Direction::Right => head_pos.x += 1,
            c::Direction::Up => head_pos.y += 1,
            c::Direction::Down => head_pos.y -= 1,
        }
        if head_pos.x < 0
            || head_pos.y < 0
            || head_pos.x as u32 >= c::ARENA_WIDTH
            || head_pos.y as u32 >= c::ARENA_HEIGHT
            || segment_positions.contains(&head_pos)
        {
            game_over_writer.send(c::GameOverEvent);
        }

        segment_positions
            .iter()
            .zip(segments.segments.iter().skip(1))
            .for_each(|(pos, segment)| {
                *positions.get_mut(*segment).unwrap() = *pos;
            });
        last_tail_position.pos = Some(*segment_positions.last().unwrap());
    }
}

pub fn eating(
    mut commands: Commands,
    mut growth_writer: EventWriter<c::GrowthEvent>,
    food_positions: Query<(Entity, &c::Position), With<c::Food>>,
    head_positions: Query<&c::Position, With<c::SnakeHead>>,
) {
    for head_pos in head_positions.iter() {
        for (ent, food_pos) in food_positions.iter() {
            if food_pos == head_pos {
                commands.entity(ent).despawn();
                growth_writer.send(c::GrowthEvent);
            }
        }
    }
}

pub fn growth(
    commands: Commands,
    last_tail_position: Res<c::LastTailPosition>,
    mut segments: ResMut<c::SnakeSegments>,
    mut growth_reader: EventReader<c::GrowthEvent>,
    materials: Res<c::Materials>,
) {
    if growth_reader.iter().next().is_some() {
        segments.segments.push(spawn_segment(
            commands,
            &materials.segment_material,
            last_tail_position.pos.unwrap(),
        ));
    }
}

pub fn die(
    mut commands: Commands,
    mut reader: EventReader<c::GameOverEvent>,
    materials: Res<c::Materials>,
    segments_res: ResMut<c::SnakeSegments>,
    food: Query<Entity, With<c::Food>>,
    segments: Query<Entity, With<c::SnakeSegment>>,
) {
    if reader.iter().next().is_some() {
        for ent in food.iter().chain(segments.iter()) {
            commands.entity(ent).despawn();
        }
        spawn(commands, materials, segments_res);
    }
}

pub fn food(
    mut commands: Commands,
    materials: Res<c::Materials>,
    segments: Query<&c::Position, With<c::SnakeSegment>>,
) {
    let snake_positions = segments.iter().collect::<Vec<&c::Position>>();
    // try three times, giving up if we still haven't found a free spot
    let mut x;
    let mut y;
    let mut attempts = 3;
    loop {
        x = (random::<f32>() * c::ARENA_WIDTH as f32) as i32;
        y = (random::<f32>() * c::ARENA_HEIGHT as f32) as i32;
        if !snake_positions.contains(&&c::Position { x, y }) {
            break;
        }
        attempts -= 1;
        if attempts == 0 {
            return;
        }
    }
    commands
        .spawn_bundle(SpriteBundle {
            material: materials.food_material.clone(),
            ..Default::default()
        })
        .insert(c::Food)
        .insert(c::Position { x, y })
        .insert(c::GridSize::square(0.8));
}