Taking "Flight" with Bevy

// src/inputs.rs
#[derive(Bundle)]
pub struct ArrowKeysControllerBundle {
    input_manager: InputManagerBundle<Action>,
}

impl ArrowKeysControllerBundle {
    pub fn new() -> Self {
        ArrowKeysControllerBundle {
            input_manager: InputManagerBundle::<Action> {       
                action_state: ActionState::default(),
                input_map: InputMap::new([
                    (KeyCode::Up, Action::Up),
                    (KeyCode::Left, Action::Left),
                    (KeyCode::Down, Action::Down),
                    (KeyCode::Right, Action::Right),
                    (KeyCode::Escape, Action::Menu),
                ])
                .build(),
            },
        }
    }
}

// e2e/local-multi.rs
use bevy::{
    prelude::{Commands, IntoSystemAppConfig, OnEnter, Res},
    sprite::{SpriteSheetBundle, TextureAtlasSprite},
};

use common_e2e::Test;

use crabber::{
    components::{Crab, StepMotor, Transform},
    constants::PLAYER_Z,
    resources::SpriteSheetAssets,
    AppState, ArrowKeysControllerBundle, CoreGameLoopPlugin, GraphicsPlugin as CrabGraphicsPlugin,
    InputPlugin, LevelPlugin, TileRow, WASDControllerBundle,
};

fn spawn_players(mut commands: Commands, spritesheets: Res<SpriteSheetAssets>) {
    commands.spawn((
        Crab,
        SpriteSheetBundle {
            texture_atlas: spritesheets.crab.clone(),
            sprite: TextureAtlasSprite::new(0),
            transform: Transform::from_xyz(0., f32::from(TileRow(0)), PLAYER_Z),
            ..Default::default()
        },
        WASDControllerBundle::new(),
        StepMotor::new(),
    ));
    // !!! NEW CODE !!!
    commands.spawn((
        Crab,
        SpriteSheetBundle {
            texture_atlas: spritesheets.crab.clone(),
            sprite: TextureAtlasSprite::new(0),
            transform: Transform::from_xyz(0., f32::from(TileRow(0)), PLAYER_Z),
            ..Default::default()
        },
        ArrowKeysControllerBundle::new(),
        StepMotor::new(),
    ));
}

fn main() {
    Test {
        label: "Test local multiplayer".to_string(),
        setup: |app| {
            app.add_state::<AppState>()
                .add_plugin(InputPlugin)
                .add_plugin(CoreGameLoopPlugin)
                .add_plugin(LevelPlugin)
                .add_system(spawn_players.in_schedule(OnEnter(AppState::InGame)));
        },
        setup_graphics: |app| {
            app.add_plugin(CrabGraphicsPlugin);
        },
    }
    .run();
}

// crabber/protocol/src/lib.rs
pub mod components;
pub mod constants;

// crabber/protocol/src/components/mod.rs
use bevy_ecs::prelude::Component;
use bevy_math::prelude::Vec3;

use naia_bevy_shared::{Property, Replicate, Serde};

#[derive(Clone, Copy, PartialEq, Serde)]
//                               ^^^^^
pub enum Direction {
    Up,
    Right,
    Down,
    Left,
}

#[derive(Component, Replicate)]
pub struct ConstantMotor {
    pub speed: Property<f32>,
    //         ^^^^^^^^
    pub direction: Property<Direction>,
    //             ^^^^^^^^
}

impl ConstantMotor {
    pub fn drive_offscreen(&self, transform: &mut Transform) -> bool {
        // derefencing the `Property` type returns the underlying type
        transform.translation += *self.direction.to_vec() * *self.speed;
        //                       ^                          ^
        is_offscreen(transform, *self.direction)
        //                      ^
    }

    pub fn drive_and_loop(&self, transform: &mut Transform) {
        if self.drive_offscreen(transform) {
            transform.translation += get_offset_for_loop(*self.direction);
            //                                           ^
        }
    }
}

// impl ConstantMotor { ...
    pub fn new(speed: f32, direction: Direction) -> Self {
        Self::new_complete(speed, direction)
    }

// impl StepMotor { ...
    pub fn new() -> Self {
        Self::new_complete(None)
    }

#[derive(Component, Replicate)]
pub struct Position {
    // represents transform.translation.x
    pub x: Property<f32>,
    // represents transform.translation.y
    pub y: Property<f32>,
    // represents rotation around the z axis to control
    // the direction the entity is "facing"
    pub direction: Property<Direction>,
}

impl Position {
    pub fn new(x: f32, y: f32, direction: Direction) -> Self {
        Position::new_complete(x, y, direction)
    }
}

fn is_offscreen(position: &Position, direction: Direction) -> bool {
    let max_abs = match direction {
        Direction::Left | Direction::Right => MAX_X_F32,
        Direction::Up | Direction::Down => MAX_Y_F32,
    };
    let current_value = match direction {
        Direction::Left | Direction::Right => *position.x,
        Direction::Up | Direction::Down => *position.y,
    };
    let motion_vector = match direction {
        Direction::Left | Direction::Right => direction.to_vec().x,
        Direction::Up | Direction::Down => direction.to_vec().y,
    };
    current_value.abs() > max_abs
        && current_value.is_sign_positive() == motion_vector.is_sign_positive()
}

impl ConstantMotor {
    // ...

    pub fn drive_offscreen(&self, position: &mut Position) -> bool {
        let delta = self.direction.to_vec() * *self.speed;
        *position.x += delta.x;
        *position.y += delta.y;
        is_offscreen(position, *self.direction)
    }

    pub fn drive_and_loop(&self, position: &mut Position) {
        if self.drive_offscreen(position) {
            let delta = get_offset_for_loop(*self.direction);
            *position.x += delta.x;
            *position.y += delta.y;
        }
    }
}

impl Position {
    // ...
    pub fn move_direction(&mut self, delta: f32, direction: Direction) {
        match direction {
            Direction::Up => {
                *self.y += delta;
            }
            Direction::Down => {
                *self.y -= delta;
            }
            Direction::Right => {
                *self.x += delta;
            }
            Direction::Left => {
                *self.x -= delta;
            }
        }
    }

    pub fn move_forward(&mut self, delta: f32) {
        self.move_direction(delta, *self.direction);
    }
}

impl StepMotor {
    // ...
    pub fn start(&mut self, position: &mut Position, direction: Direction) {
        *position.direction = direction;
        *transform = transform.looking_to(Vec3::Z, direction.to_vec());
        *self.step = Some(0);
    }

    pub fn drive(&mut self, position: &mut Position) {
        if self.is_running() {
            position.move_forward(STEP_SPEED);
            *self.step = self
                .step
                .map(|step| step + 1)
                .filter(|&step| step < MOTION_STEPS);
        }
    }
}

#[derive(Component, Replicate)]
pub struct Score {
    pub value: Property<u16>,
}

// crabber/protocol/src/components/level.rs
#[derive(Component, Replicate)]
pub struct Level {
    pub rows: Property<Vec<LevelRow>>,
}

impl Level {
    // this builds a random vector of rows
    pub fn new_random() -> Self {
        let mut rows = Vec::new();
        // The level should start with grass
        let mut level_row_kind = LevelRow::Grass;
        rows.push(level_row_kind);
        // Then we should go up to the N-1 row from there
        for _ in 1..(LEVEL_HEIGHT_I16 - 1) {
            level_row_kind = level_row_kind.get_random_next();
            rows.push(level_row_kind);
        }
        // Finally, add a finish line.
        rows.push(LevelRow::Finish);
        Level::new_complete(rows)
    }

    // this spawns cars and rafts for a level
    pub fn spawn_level_entities(&self, commands: &mut Commands) {
        // Then we should go up to the N-1 row from there
        for (row_index, row_kind) in self.rows.iter().enumerate() {
            if LevelRow::Road == *row_kind {
                for (position, motor) in build_random_motors(row_index as i16).into_iter() {
                    commands.spawn((Car, position, motor));
                }
            }
            if LevelRow::River == *row_kind {
                for (position, motor) in build_random_motors(row_index as i16).into_iter() {
                    commands.spawn((Raft, position, motor));
                }
            }
        }
    }

    // a useful resource for determining the row kind
    pub fn is_row_of_kind(&self, row: TileRow, target: LevelRow) -> bool {
        self.rows
            .get(row.0 as usize)
            .and_then(|row_kind| if *row_kind == target { Some(()) } else { None })
            .is_some()
    }
}

// .../components/mod.rs
mod level;
pub use level::{Level, LevelRow};

// .../messages.rs
use naia_bevy_shared::{EntityProperty, Message, Serde};

#[derive(Clone, Copy, Debug, PartialEq, Serde)]
pub enum InputAction {
    Up,
    Down,
    Left,
    Right,
}

#[derive(Message)]
pub struct InputMessage {
    pub entity: EntityProperty,
    pub action: Option<InputAction>,
}

impl InputMessage {
    pub fn new(action: Option<InputAction>) -> Self {
        InputMessage {
            entity: EntityProperty::new_empty(),
            action,
        }
    }
}

// .../channels.rs
use naia_bevy_shared::{
    Channel, ChannelDirection, ChannelMode, Protocol, TickBufferSettings,
};

// We will use this channel to send inputs to the server
#[derive(Channel)]
pub struct PlayerInputChannel;
impl PlayerInputChannel {
    pub fn add_to_protocol(protocol: &mut Protocol) {
        protocol.add_channel::<PlayerInputChannel>(
            // the channel allows messages to be sent from client to server
            ChannelDirection::ClientToServer,
            // the channel is "tick-buffered",
            // meaning it is syncronized with NAIA's tick event loop
            ChannelMode::TickBuffered(TickBufferSettings::default()),
        );
    }
}

// .../lib.rs
use std::time::Duration;
use naia_bevy_shared::{LinkConditionerConfig, Protocol, ProtocolPlugin};

pub mod channels;
pub mod components;
pub mod constants;
pub mod messages;

struct CrabberProtocolPlugin;

impl ProtocolPlugin for CrabberProtocolPlugin {
    fn build(&self, protocol: &mut Protocol) {
        // register the channels
        channels::PlayerInputChannel::add_to_protocol(protocol);

        protocol
            // register the message types
            .add_message::<messages::InputMessage>()
            // register the component types
            .add_component::<components::Crab>()
            .add_component::<components::Car>()
            .add_component::<components::Raft>()
            .add_component::<components::Knockout>()
            .add_component::<components::ConstantMotor>()
            .add_component::<components::StepMotor>()
            .add_component::<components::Position>()
            .add_component::<components::Level>()
            .add_component::<components::Score>();
    }
}

pub fn protocol() -> Protocol {
    Protocol::builder()
        .tick_interval(Duration::from_millis(16))
        .link_condition(LinkConditionerConfig::good_condition())
        .add_plugin(CrabberProtocolPlugin)
        .build()
}

// crabber/server/src/main.rs
use std::time::Duration;

use bevy_app::{App, ScheduleRunnerPlugin, ScheduleRunnerSettings};
use bevy_core::{FrameCountPlugin, TaskPoolPlugin, TypeRegistrationPlugin};
use bevy_log::{info, LogPlugin};

use naia_bevy_server::{Plugin as ServerPlugin, ServerConfig};

use crabber_protocol::protocol;

fn main() {
    info!("Starting up Crabber server...");

    App::default()
        // Some base Bevy plugins
        .add_plugin(TaskPoolPlugin::default())
        .add_plugin(TypeRegistrationPlugin::default())
        .add_plugin(FrameCountPlugin::default())
        .insert_resource(ScheduleRunnerSettings::run_loop(Duration::from_millis(3)))
        .add_plugin(ScheduleRunnerPlugin::default())
        .add_plugin(LogPlugin::default())
        // NAIA's server plugin
        .add_plugin(ServerPlugin::new(
            // settings to monitor ping, jitter, manage connection details, etc.
            ServerConfig {
                // we don't have any form of auth yet
                require_auth: false,
                ..Default::default(),
            },
            // our protocol!
            protocol(),
        ))
        // Run App
        .run();
}

use bevy_log::info;

use naia_bevy_server::{transport::webrtc, Server};

pub fn init(mut server: Server) {
    info!("Starting Crabber server");

    let server_addresses = webrtc::ServerAddrs::new(
        "127.0.0.1:14191" // TODO: replace this with an env variable or something
            .parse()
            .expect("could not parse Signaling address/port"),
        // IP Address to listen on for UDP WebRTC data channels
        "127.0.0.1:14192" // TODO: replace this with an env variable or something
            .parse()
            .expect("could not parse WebRTC data address/port"),
        // The public WebRTC IP address to advertise
        "http://127.0.0.1:14192", // TODO: replace this with an env variable or something
    );
    let socket = webrtc::Socket::new(&server_addresses, server.socket_config());
    // Initialize the server
    server.listen(socket);
}

// .../connection.rs
pub fn connect_events(
    mut commands: Commands,
    mut server: Server,
    mut event_reader: EventReader<ConnectEvent>,
) {
    for ConnectEvent(user_key) in event_reader.iter() {
        // get the server's first key if there is one
        // TODO: check if the room is "full"
        let room_key = server
            .room_keys()
            .into_iter()
            // get the first RoomKey
            .next()
            // if there are no RoomKeys, spawn one
            .unwrap_or_else(|| server.make_room().key());

        // add user to the room
        let address = server
            .user_mut(user_key)
            .enter_room(&room_key)
            .address();

        // write a note to the log
        info!("Client connected from: {}", address);

        // spawn new entity for the player
        let entity = commands
            .spawn_empty()
            // call this to register the entity for replication with NAIA
            .enable_replication(&mut server)
            .insert(todo!("Add components here")) // N.B. the todo
            .id();

        server.room_mut(&room_key).add_entity(&entity);
    }
}

// .../connection.rs
pub fn disconnect_events(
    mut commands: Commands,
    mut server: Server,
    mut global: ResMut<Global>,
    mut event_reader: EventReader<DisconnectEvent>,
) {
    for DisconnectEvent(user_key, user) in event_reader.iter() {
        info!("Crabber Server disconnected from: {:?}", user.address);
        
        // TODO Remove entities that are associated with user_key
        let entity = todo!();
        server
            .room_mut(todo!()) // also need to figure out which remove they are in
            .remove_entity(&entity);
        commands.entity(entity).despawn();
        // later on, this code also might want to set fields that help identify reconnecting users, etc.
    }
}

pub fn error_events(mut event_reader: EventReader<ErrorEvent>) {
    for ErrorEvent(error) in event_reader.iter() {
        info!("Crabber Server Error: {:?}", error);
    }
}

// .../resources.rs
// this is similar to the naia-bevy-server-demo "Global" resource
#[derive(Resource, Default)]
pub struct UserEntities {
    user_to_entity_map: HashMap<UserKey, Entity>,
    entity_to_user_map: HashMap<Entity, UserKey>,
}

// Some methods that operate on the inner HashMaps for ease-of-use
impl UserEntities {
    fn get_by_user(&self, &UserKey) -> Option<&Entity> { /* ... */ }
    fn get_by_entity(&self, &Entity) -> Option<&UserKey> { /* ... */ }
    fn insert(&mut self, UserKey, Entity) { /* ... */ }
    fn remove(&mut self, &UserKey) -> Option<&Entity> { /* ... */ }
}

fn main() {
    App::default()
        // initialize the resource using its `Default` implementation
        .init_resource::<UserEntities>()
}

// .../connection.rs
pub fn connect_events(
    /* ... */
    mut user_entities: ResMut<UserEntities>,
) {
    for ConnectEvent(user_key) in event_reader.iter() {
        /* ... */
        user_entities.insert(*user_key, entity);
    }
}

pub fn disconnect_events(
    /* ... */
    mut user_entities: ResMut<UserEntities>,
) {
    for DisconnectEvent(user_key, user) in event_reader.iter() {
        if let Some(entity) = user_entities.remove(user_key) {
            server
                .room_mut(todo!()) // still TODO
                .remove_entity(&entity);
            commands.entity(entity).despawn();
        }
    }
}

// .../connection.rs
pub fn disconnect_events(/* ... */) {
    for DisconnectEvent(user_key, user) in event_reader.iter() {
        if let Some(entity) = user_entities.remove(user_key) {
            let room_keys = {
                let user_ref = server.user(user_key);
                // collecting the iterator so that we don't maintain a reference to server
                user_ref.room_keys().map(|key| *key).collect::<Vec<_>>()
            };
            for room_key in room_keys.into_iter() {
                server.room_mut(&room_key).remove_entity(&entity);
            }
            commands.entity(entity).despawn();
        }
    }
}

// crabber/protocol/src/bundles.rs
use bevy_ecs::prelude::Bundle;

use crate::components::{Crab, Direction, Position, StepMotor, Score};

#[derive(Bundle)]
pub struct CrabBundle {
    crab: Crab,
    motor: StepMotor,
    position: Position,
    score: Score,
}

impl CrabBundle {
    pub fn new() -> Self {
        CrabBundle {
            crab: Crab,
            motor: StepMotor::new(),
            position: Position::new(0., f32::from(TileRow(0)), Direction::Up),
            score: Score::new()
        }
    }
}

// crabber/server/src/connection.rs
pub fn connect_events(
    mut commands: Commands,
    mut server: Server,
    mut user_entities: ResMut<UserEntities>,
    mut event_reader: EventReader<ConnectEvent>,
) {
    for ConnectEvent(user_key) in event_reader.iter() {
        // ...
        let entity = commands
            .spawn_empty()
            .enable_replication(&mut server)
            .insert(CrabBundle::new())
            // new: ^^^^^^^^^^^^^^^^^
            .id();
        // ...
    }
}

// .../main.rs
use bevy_ecs::schedule::IntoSystemConfigs;
use naia_bevy_server::ReceiveEvents;

mod connection;
mod init;

fn main() {
    App::default()
        // ... (skipping the base Bevy plugins and ServerPlugin)
        .init_resource::<UserEntities>()
        .add_startup_system(init::init)
        .add_systems(
            (
                connection::connect_events,
                connection::disconnect_events,
                connection::error_events,
            )
                // Add them to the SystemSet that NAIA uses for network-related systems
                .in_set(ReceiveEvents),
        )
        .run();
}

pub fn connect_events(/* ... */) {
    for ConnectEvent(user_key) in event_reader.iter() {
        // ...

        let num_players = players_query.into_iter().count();

        // spawn a level if we are about to spawn the second player
        if num_players == 1 {
            let level = Level::new_random();
            // level.spawn_level_entities(&mut commands);
            // TODO! ^^^ how to `enable_replication`?
            let entity = commands
                .spawn_empty()
                .enable_replication(&mut server)
                .insert(level)
                .id(); // spawn our level!
            server.room_mut(&room_key).add_entity(&entity);
        }

        // only spawn player entities for the first two players
        if num_players < 2 {
            let entity = commands
                .spawn_empty()
                .enable_replication(&mut server)
                .insert(CrabBundle::new())
                .id();
            server.room_mut(&room_key).add_entity(&entity);
            // ...
        }
    }
}

// crabber/protocol/src/components/level.rs
// impl Level
    pub fn create_level_bundles(
        &self,
    ) -> (
        Vec<(Car, Position, ConstantMotor)>,
        Vec<(Raft, Position, ConstantMotor)>,
    ) {
        let mut car_bundles = Vec::new();
        let mut raft_bundles = Vec::new();
        // Then we should go up to the N-1 row from there
        for (row_index, row_kind) in self.rows.iter().enumerate() {
            if LevelRow::Road == *row_kind {
                for (position, motor) in build_random_motors(row_index as i16).into_iter() {
                    car_bundles.push((Car, position, motor));
                }
            }
            if LevelRow::River == *row_kind {
                for (position, motor) in build_random_motors(row_index as i16).into_iter() {
                    raft_bundles.push((Raft, position, motor));
                }
            }
        }
        (car_bundles, raft_bundles)
    }

// crabber/server/src/connection.rs
// pub fn connect_events(...) { ...
//  let room_key = ...;
    if num_players == 1 {
        let level = Level::new_random();
        let (car_bundles, raft_bundles) = level.create_level_bundles();
        // for every entity we want to create:
        // 1. spawn the bundle
        // 2. enable replication
        // 3. the entity to room_key's room
        for bundle in car_bundles.into_iter() {
            let entity = commands.spawn(bundle).enable_replication(&mut server).id();
            server.room_mut(&room_key).add_entity(&entity);
        }
        for bundle in raft_bundles.into_iter() {
            let entity = commands.spawn(bundle).enable_replication(&mut server).id();
            server.room_mut(&room_key).add_entity(&entity);
        }
        let entity = commands
            .spawn_empty()
            .enable_replication(&mut server)
            .insert(level)
            .id();
        server.room_mut(&room_key).add_entity(&entity);
    }

// naia-bevy-demo-server/src/events.rs
pub fn tick_events(
    mut server: Server,
    mut position_query: Query<&mut Position>,
    mut tick_reader: EventReader<TickEvent>,
) {
    let mut has_ticked = false;

    for TickEvent(server_tick) in tick_reader.iter() {
        has_ticked = true;

        // All game logic happens here!
        // Meaning all the logic in our "CoreGameLoopPlugin" systems should be moved into here
        // (Or into other systems that read from the EventReader<TickEvent>)

        // This lets us read from the InputMessageChannel for the current tick
        let mut messages = server.receive_tick_buffer_messages(server_tick);
        for (_user_key, command) in messages.read::<PlayerInputChannel, InputMessage>() {
            let Some(entity) = &command.entity.get(&server) else {
                continue;
            };
            let Ok(mut position) = position_query.get_mut(*entity) else {
                continue;
            };
            shared_behavior::process_command(&command, &mut position);
            // ^^^ we don't have this! but we can define it
        }
    }

    // this block is determines which users should be sent information about which entities
    if has_ticked {
        for (_, user_key, entity) in server.scope_checks() {
            server.user_scope(&user_key).include(&entity);
        }
    }
}

// crabber/server/src/main.rs

mod tick;

fn main() {
    App::default()
        // ...
        .add_systems(
            (
                connection::connect_events,
                connection::disconnect_events,
                connection::error_events,
                tick::tick_events, // <-- new!
            )
                .in_set(ReceiveEvents),
        )
        .run();
}

fn register_inputs(
    mut player_query: Query<(&mut Position, &mut StepMotor, &PlayerActionState), Without<Knockout>>,
) {
    for (mut position, mut motor, action) in player_query.iter_mut() {
        let direction = if action.pressed(Action::Up) {
            Some(Direction::Up)
        } else if action.pressed(Action::Down) {
            Some(Direction::Down)
        } else if action.pressed(Action::Left) {
            Some(Direction::Left)
        } else if action.pressed(Action::Right) {
            Some(Direction::Right)
        } else {
            None
        };
        if let Some(direction) = direction && !motor.is_running() {
            motor.start(&mut position, direction);
        }
    }
}

// crabber/protocol/src/inputs.rs
pub fn process_input_messages(
    // Each player entity and the associated input action for this tick
    inputs: Vec<(Entity, Option<InputAction>)>,
    // we will need a mutable reference to this query since it will be called in a loop
    player_query: &mut Query<(&mut Position, &mut StepMotor), Without<Knockout>>,
) {
    for (entity, action) in inputs.into_iter() {
        let direction = action.map(|action| match action {
            InputAction::Up => Direction::Up,
            InputAction::Down => Direction::Down,
            InputAction::Left => Direction::Left,
            InputAction::Right => Direction::Right,
        });
        if let Some(direction) = direction {
            if let Ok((mut position, mut motor)) = player_query.get_mut(entity) {
                if !motor.is_running() {
                    motor.start(&mut position, direction);
                }
            }
        }
    }
}

// crabber/server/src/tick.rs
use crabber_protocol::inputs;

pub fn tick_events(
    mut server: Server,
    // new query type
    mut player_query: Query<(&mut Position, &mut StepMotor), Without<Knockout>>,
    mut tick_reader: EventReader<TickEvent>,
) {
    for TickEvent(server_tick) in tick_reader.iter() {
        let mut messages = server.receive_tick_buffer_messages(server_tick);
        let mut player_actions = Vec::new();
        for (_user_key, input) in messages.read::<PlayerInputChannel, InputMessage>() {
            if let Some(entity) = command.entity.get(&server) {
                player_actions.push((entity, input.action))
            };
        }
        inputs::process_input_messages(player_actions, &mut player_query);
        // TODO: other tick systems
    }
}

fn tick_constant_motors(motor_query: &mut Query<(&mut Position, &ConstantMotor)>) {}

fn tick_step_motors(motor_query: &mut Query<(&mut Position, &mut StepMotor), Without<Knockout>>) {}

fn tick_score(player_query: &mut Query<(&mut Score, &Position), Without<Knockout>>) {}

pub fn tick_events(
    mut server: Server,
    mut player_query: Query<(&mut Position, &mut StepMotor), Without<Knockout>>,
    mut score_query: Query<(&mut Score, &Position), Without<Knockout>>,
    mut tick_reader: EventReader<TickEvent>,
) { /* ... */ }

pub fn tick_events(
    mut server: Server,
    mut query_set: ParamSet<(
        Query<(&mut Position, &mut StepMotor), Without<Knockout>>,
        Query<(&mut Position, &ConstantMotor)>,
        Query<(&mut Score, &Position), Without<Knockout>>,
        Query<(Entity, &mut Position, &StepMotor), (With<Crab>, Without<Knockout>)>,
    )>,
    mut tick_reader: EventReader<TickEvent>,
) {
    for TickEvent(server_tick) in tick_reader.iter() {
        let mut messages = server.receive_tick_buffer_messages(server_tick);
        let mut player_actions = Vec::new();
        for (_user_key, input) in messages.read::<PlayerInputChannel, InputMessage>() {
            if let Some(entity) = input.entity.get(&server) {
                player_actions.push((entity, input.action))
            };
        }
        inputs::process_inputs(player_actions, &mut query_set.p0());
        tick::tick_step_motors(&mut query_set.p0());
        tick::tick_constant_motors(&mut query_set.p1());
        tick::tick_score(&mut query_set.p2());
    }
    // ...
}

pub fn tick_events(
    mut commands: Commands,
    mut server: Server,
    mut player_query_set: ParamSet<(
        Query<(&mut Position, &mut StepMotor), (With<Crab>, Without<Knockout>)>,
        Query<(&mut Score, &Position), (With<Crab>, Without<Knockout>)>,
        Query<(Entity, &mut Position, &StepMotor), (With<Crab>, Without<Knockout>)>,
    )>,
    mut objects_query_set: ParamSet<(
        Query<(&mut Position, &ConstantMotor), Without<Crab>>,
        Query<&Position, (With<Car>, Without<Raft>, Without<Crab>)>,
        Query<(&Position, &ConstantMotor), (With<Raft>, Without<Car>, Without<Crab>)>,
    )>,
    level_query: Query<&Level>,
    mut tick_reader: EventReader<TickEvent>,
) {
    for TickEvent(server_tick) in tick_reader.iter() {
        // ...
        inputs::process_inputs(player_actions, &mut player_query_set.p0());
        tick::tick_step_motors(&mut player_query_set.p0());
        tick::tick_constant_motors(&mut objects_query_set.p0());
        tick::tick_score(&mut player_query_set.p1());
        tick::tick_road_collisions(
            &mut commands,
            &level_query,
            &player_query_set.p2().to_readonly(),
            &objects_query_set.p1(),
        );
        tick::tick_river_collisions(
            &mut commands,
            &level_query,
            &mut player_query_set.p2(),
            &objects_query_set.p2(),
        );
    }
    // ...
}

use naia_bevy_client::{
    transport::webrtc,
    Client,
};
use crate::AppState;

pub fn inititate_connection(mut client: Client) {
    // create a socket
    let socket = webrtc::Socket::new("http://127.0.0.1:14191", client.socket_config());
    // connect to it
    client.connect(socket);
}

use bevy::prelude::{info, EventReader, NextState, ResMut};
use naia_bevy_client::events::ConnectEvent;

pub fn connection_events(
    mut event_reader: EventReader<ConnectEvent>,
    client: Client,
    mut state: ResMut<NextState<AppState>>,
) {
    for _event in event_reader.iter() {
        info!("Client connected to: {:?}", client.server_address());
        // enter the "InGame" state
        state.set(AppState::InGame);
    }
}

use naia_bevy_client::events::{DisconnectEvent, RejectEvent};

pub fn disconnection_events(
    client: Client,
    mut state: ResMut<NextState<AppState>>,
    mut event_reader: EventReader<DisconnectEvent>,
) {
    for _event in event_reader.into_iter() {
        info!("Client disconnected from: {:?}", client.server_address());
        // reset to loading state for now
        // eventually, we could use this to show an alert supporting user actions
        // or handle some sort of reconnection
        state.set(AppState::Loading);
    }
}

pub fn rejection_events(mut event_reader: EventReader<RejectEvent>) {
    for _ in event_reader.iter() {
        info!("Client rejected from connecting to Server");
    }
}

// crabber/app/src/lib.rs
use bevy::prelude::{App, IntoSystemConfigs, States};
use naia_bevy_client::{ClientConfig, Plugin as ClientPlugin, ReceiveEvents};
use crabber_protocol::protocol;

mod connection;

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, States)]
pub enum AppState {
    #[default]
    Loading,
    InGame,
}

pub fn build(app: &mut App) -> &mut App {
    app.add_state::<AppState>()
        // add the client plugin
        .add_plugin(ClientPlugin::new(ClientConfig::default(), protocol()))
        // our init system
        .add_startup_system(connection::inititate_connection)
        // and the event systems
        .add_systems(
            (
                connection::connection_events,
                connection::disconnection_events,
                connection::rejection_events,
            )
                .in_set(ReceiveEvents),
        )
}

// crabber/app/e2e/full-client.rs
use common_e2e::Test;

use crabber_app::build;

fn main() {
    Test {
        label: "Test full client".to_string(),
        setup: |app| {
            build(app);
        },
        setup_graphics: |app| {},
    }
    .run();
}

// crabber/app/src/graphics.rs
fn direction_to_angle(direction: Direction) -> f32 {
    match direction {
        Direction::Up => 0.,
        Direction::Right => -std::f32::consts::FRAC_PI_2,
        Direction::Down => std::f32::consts::PI,
        Direction::Left => std::f32::consts::FRAC_PI_2,
    }
}

fn position_to_transform(position: &Position, z: f32) -> Transform {
    Transform::from_xyz(*self.x, *self.y, z)
        .with_rotation(Quat::from_rotation_z(self.direction.to_angle()))
}

// .../graphics.rs
fn sync_transforms(mut position_query: Query<(&Position, &mut Transform)>) {
    for (position, mut transform) in position_query.iter_mut() {
        // don't change any z's
        *transform = position_to_transform(position, transform.translation.z);
    }
}

pub fn insert_component_events(
    mut commands: Commands,
    mut event_reader: EventReader<InsertComponentEvents>,
    spritesheets: Res<SpriteSheetAssets>,
) {
    for events in event_reader.iter() {
        for entity in events.read::<Raft>() {
            info!("raft spawned");
            commands.entity(entity)
                .insert(SpriteSheetBundle { /* ... */ });
        }
        /* ... more types */
    }
}

// .../graphics.rs
fn handle_knockout(mut ko_query: Query<&mut TextureAtlasSprite, Added<Knockout>>) {
    for mut sprite in ko_query.iter_mut() {
        sprite.color = Color::rgba(1., 1., 1., 0.5);
        sprite.flip_y = true;
    }
}

// .../graphics.rs
fn setup_crab_sprites(
    mut commands: Commands,
    added_crabs_query: Query<(Entity, &Position), Added<Crab>>,
    spritesheets: Res<SpriteSheetAssets>,
) {
    for (entity, position) in added_crabs_query.iter() {
        commands.entity(entity).insert(SpriteSheetBundle {
            texture_atlas: spritesheets.crab.clone(),
            sprite: TextureAtlasSprite::new(0),
            // Convert the `Position` to a `Transform`
            transform: position_to_transform(position, PLAYER_Z),
            ..Default::default()
        });
    }
}

// .../graphics.rs
fn setup_car_sprites(
    mut commands: Commands,
    added_cars_query: Query<(Entity, &Position), Added<Car>>,
    spritesheets: Res<SpriteSheetAssets>,
) {
    let mut rng = rand::thread_rng();
    for (entity, position) in added_cars_query.iter() {
        let random_color_offset = rng.gen_range(0..=2);
        let start_sprite_index = if Direction::Left == *position.direction {
            3
        } else {
            0
        };
        commands.entity(entity).insert(SpriteSheetBundle {
            texture_atlas: spritesheets.car.clone(),
            sprite: TextureAtlasSprite::new(start_sprite_index + random_color_offset),
            transform: position_to_transform(position, LEVEL_Z),
            ..Default::default()
        });
    }
}

// .../graphics.rs
fn position_to_transform(position: &Position, z: f32, rotated: bool) -> Transform {
    if rotated {
        Transform::from_xyz(*position.x, *position.y, z).with_rotation(Quat::from_rotation_z(
            direction_to_angle(*position.direction),
        ))
    } else {
        Transform::from_xyz(*position.x, *position.y, z)
    }
}

fn sync_transforms(mut position_query: Query<(&Position, &mut Transform, Option<&Crab>)>) {
    for (position, mut transform, crab /*: Option<&Crab> */) in position_query.iter_mut() {
        *transform = position_to_transform(position, transform.translation.z, crab.is_some());
    }
}

// .../lib.rs
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, States)]
pub enum AppState {
    #[default]
    Loading,      // loading assets
    Connecting,   // connecting to game
    InGame,       // in game actively
    Disconnected, // disconnected
}

// .../graphics.rs
// these traits enable `.run_if` and `.after` below
use bevy::prelude::{IntoSystemConfig, IntoSystemConfigs};

use bevy_ecs_tilemap::prelude::TilemapPlugin;

pub struct GraphicsPlugin;

impl Plugin for GraphicsPlugin {
    fn build(&self, app: &mut App) {
        app
        .add_plugin(TilemapPlugin) // make sure to keep bevy_ecs_tilemap's plugin
        // add the asset loading code
        .add_loading_state(
            LoadingState::new(AppState::Loading).continue_to_state(AppState::Connecting),
        )
        .add_collection_to_loading_state::<_, SpriteSheetAssets>(AppState::Loading)
        // now our systems
        .add_startup_system(camera)
        .add_systems(
            (
                // this code adds the run condition to each system, but could use distributed_run_if instead
                handle_knockout.run_if(in_state(AppState::InGame)),
                setup_crab_sprites.run_if(in_state(AppState::InGame)),
                setup_car_sprites.run_if(in_state(AppState::InGame)),
                setup_raft_sprites.run_if(in_state(AppState::InGame)),
                setup_level_tilemap.run_if(in_state(AppState::InGame)),
                animate_sprites.run_if(in_state(AppState::InGame)),
                sync_transforms.run_if(in_state(AppState::InGame)),
            )
                .after(ReceiveEvents),
        );
    }
}

// .../controller.rs
pub enum Action { /* ... */ }
pub type PlayerActionState = // ...

#[derive(Bundle)]
pub struct WASDControllerBundle { /* ... */ }
impl WASDControllerBundle { /* ... */ }

#[derive(Bundle)]
pub struct ArrowKeysControllerBundle { /* ... */ }
impl ArrowKeysControllerBundle { /* ... */ }

fn get_action(action: &PlayerActionState) -> Option<InputAction> {
    if action.pressed(Action::Up) {
        Some(InputAction::Up)
    } else if action.pressed(Action::Down) {
        Some(InputAction::Down)
    } else if action.pressed(Action::Left) {
        Some(InputAction::Left)
    } else if action.pressed(Action::Right) {
        Some(InputAction::Right)
    } else {
        None
    }
}

#[derive(Default, Resource)]
pub struct QueuedInputs(pub HashMap<Entity, InputAction>);

fn register_inputs(
    mut player_query: Query<(Entity, &PlayerActionState), Without<Knockout>>,
    mut player_inputs: ResMut<QueuedInputs>,
) {
    for (entity, action_state) in player_query.iter_mut() {
        if let Some(input_action) = get_action(action_state) {
            player_inputs.0.insert(entity, input_action);
        }
    }
}

impl Plugin for ControllerPlugin {
    fn build(&self, app: &mut App) {
        app.add_plugin(InputManagerPlugin::<Action>::default())
            .init_resource::<QueuedInputs>()
            .add_system(
                register_inputs
                    .in_base_set(CoreSet::PreUpdate)
                    .after(InputSystem),
            );
    }
}

// .../tick.rs
pub fn tick(
    mut client: Client,
    mut tick_reader: EventReader<ClientTickEvent>,
    mut player_query: Query<
        (Entity, &mut Position, &mut StepMotor),
        (With<Crab>, Without<Knockout>),
    >,
    mut player_inputs: ResMut<QueuedInputs>,
) {
    // each of these represents a client tick
    // the variable `client_tick` is a number representing the current tick count
    // we iterate through each predicted tick and execute game logic on predicted entities
    let did_tick = false;
    for ClientTickEvent(client_tick) in tick_reader.iter() {
        did_tick = true;
        // iterate through player inputs
        // if multiple ticks are processed in one frame, send the same inputs for both
        for (entity, action) in player_inputs.0.iter() {
            if let Ok((entity, mut position, mut motor)) = player_query.get_mut(*entity) {
                let mut input_message = InputMessage::new(Some(*action));
                input_message.entity.set(&client, &entity);
                // Send command to server
                client.send_tick_buffer_message::<PlayerInputChannel, InputMessage>(
                    client_tick,
                    &input_message,
                );
                // also execute them
                process_input(input_message.action, &mut position, &mut motor); // ***
            }
        }

        tick_players(...); // ***
        tick_level(...); // ***
    }
    // if we ticked, clear the map
    if did_tick {
        player_inputs.0.clear();
    }
}

// .../components.rs

// e.g. getting the server entity from a predicted entity
// fn system(query: Query<(Entity, &PredictionOf)>) {
//     for (prediction_entity, PredictionOf(source_entity)) in query.iter() { ... }
// }
#[derive(Component)]
pub struct PredictionOf(pub Entity);
// e.g. getting the predicted entity from a server entity
// fn system(query: Query<(Entity, &SourceOf)>) {
//     for (source_entity, SourceOf(prediction_entity)) in query.iter() { ... }
// }
#[derive(Component)]
pub struct SourceOf(pub Entity);

// .../events.rs
use naia_bevy_client::CommandsExt;

fn receive_player_assignment_message(
    mut commands: Commands,
) {
    for /* TODO */ {
        let entity = /* TODO: get_entity(event); */
        // create the prediction entity
        let prediction_entity =
            commands.entity(entity)
                .duplicate() // this requires `CommandsExt`
                .insert(PredictionOf(entity))
                .id();
        // attach controls to the source entity
        commands.entity(entity).insert((
            SourceOf(prediction_entity),
            WASDControllerBundle::new(),
        ));
    }
}

// crabber/protocol/src/messages.rs
use naia_bevy_shared::{EntityProperty, Message};

#[derive(Message)]
pub struct PlayerAssignmentMessage {
    pub entity: EntityProperty,
}

impl PlayerAssignmentMessage {
    pub fn new() -> Self {
        PlayerAssignmentMessage {
            entity: EntityProperty::new_empty(),
        }
    }
}

// crabber/protocol/src/channels.rs
use naia_bevy_shared::{
    Channel, ChannelDirection, ChannelMode, Protocol, ReliableSettings,
};

// This one will be used to tell clients which player entity is theirs
#[derive(Channel)]
pub struct PlayerAssignmentChannel;
impl PlayerAssignmentChannel {
    pub fn add_to_protocol(protocol: &mut Protocol) {
        protocol.add_channel::<PlayerAssignmentChannel>(
            // the channel sends server -> client
            ChannelDirection::ServerToClient,
            // the channel sends an unordered reliable message
            // we need to ensure that the player receives the message,
            // but the order and tick do not matter, and it is OK
            // to spend additional time to ensure reliability
            ChannelMode::UnorderedReliable(ReliableSettings::default()),
        );
    }
}

// crabber/server/src/connection.rs
pub fn connect_events(/* ... */) {
    // ...
    let entity = commands
        .spawn_empty()
        .enable_replication(&mut server)
        .insert(CrabBundle::new())
        .id();

    server.room_mut(&room_key).add_entity(&entity);
    user_entities.insert(*user_key, entity);

    // new: assign the player when we spawn its entity
    let mut assignment_message = PlayerAssignmentMessage::new();
    assignment_message.entity.set(&server, &entity);
    server.send_message::<PlayerAssignmentChannel, PlayerAssignmentMessage>(
        user_key,
        &assignment_message,
    );
}

// crabber/app/src/events.rs
use naia_bevy_client::CommandsExt;

pub fn receive_entity_assignment_message(
    mut event_reader: EventReader<MessageEvents>,
    mut commands: Commands,
    client: Client,
) {
    for event in event_reader.iter() {
        for assignment in event.read::<PlayerAssignmentChannel, PlayerAssignmentMessage>() {
            let entity = assignment.entity.get(&client).unwrap();
            let prediction_entity = commands
                .entity(entity)
                .duplicate() // create a new entity and copy all `Replicate`
                .insert(PredictionOf(entity)) // this is the prediction entity
                .id();

            commands
                .entity(entity)
                .insert((
                    SourceOf(prediction_entity),
                    WASDControllerBundle::new(),
                ));
        }
    }
}

// .../graphics.rs
fn setup_crab_sprites(
    mut commands: Commands,
    added_crabs_query: Query<(Entity, &Position, Option<&SourceOf>), Added<Crab>>,
    //                                           ^^^^^^^^^^^^^^^^^
    spritesheets: Res<SpriteSheetAssets>,
) {
    for (entity, position, is_source) in added_crabs_query.iter() {
        let mut sprite = TextureAtlasSprite::new(0);
        // add a tint for non-player crabs
        if is_source.is_some() {
            sprite.color = Color::rgba(1., 1., 1., 0.9);
        }
        commands.entity(entity).insert(SpriteSheetBundle {
            texture_atlas: spritesheets.crab.clone(),
            sprite,
            transform: position_to_transform(position, PLAYER_Z, true),
            ..Default::default()
        });
    }
}

// .../tick.rs
pub fn tick(
    mut client: Client,
    mut tick_reader: EventReader<ClientTickEvent>,
    source_query: Query<&SourceOf>,
    mut prediction_query: Query<
        (&mut Position, &mut StepMotor),
        (With<PredictionOf>, With<Crab>, Without<Knockout>),
    >,
    mut player_inputs: ResMut<QueuedInputs>,
) {
    // ...
    for ClientTickEvent(client_tick) in tick_reader.iter() {
        for (entity, action) in player_inputs.0.iter() {
            let Ok(SourceOf(prediction)) = source_query.get(*entity) else { continue };
            if let Ok((mut position, mut motor)) = player_query.get_mut(*prediction) {
                // send inputs to source entity, process inputs with prediction components
            }
            // process the rest of the tick with prediction components
        }
    }
    // ...
}

// .../events.rs
use naia_bevy_client::events::InsertComponentEvents;

pub fn receive_insert_component_events(
    mut commands: Commands,
    mut event_reader: EventReader<InsertComponentEvents>,
) {
    for event in event_reader.iter() {
        for entity in event.read::<Raft>() {
            let prediction_entity = commands
                .entity(entity)
                .duplicate()
                .insert(PredictionOf(entity))
                .id();

            commands.entity(entity)
                .insert(SourceOf(prediction_entity));
        }
        // ... same thing but for `Car`
    }
}

// .../graphics.rs
fn setup_car_sprites(
    // ...
    added_cars_query: Query<(Entity, &Position), (Added<Car>, With<PredictionOf>)>,
    //                                                        ^^^^^^^^^^^^^^^^^^
) {
    // ...
}
fn setup_raft_sprites(
    // ...
    added_rafts_query: Query<(Entity, &Position), (Added<Raft>, With<PredictionOf>)>,
    //                                                          ^^^^^^^^^^^^^^^^^^
) {
    // ...
}

// .../tick.rs
pub fn tick(
    mut client: Client,
    mut commands: Commands,
    mut tick_reader: EventReader<ClientTickEvent>,
    level_query: &Query<&Level>,
    source_query: Query<&SourceOf>,
    mut player_query_set: ParamSet<(
        Query<
            (Entity, &mut Position, &mut StepMotor),
            (With<PredictionOf>, With<Crab>, Without<Knockout>),
        >,
        Query<(&mut Position, &mut StepMotor), (With<PredictionOf>, With<Crab>, Without<Knockout>)>,
        Query<
            (Entity, &mut Position, &StepMotor),
            (With<PredictionOf>, With<Crab>, Without<Knockout>),
        >,
    )>,
    mut objects_query_set: ParamSet<(
        Query<(&mut Position, &ConstantMotor), (With<PredictionOf>, Without<Crab>)>,
        Query<&Position, (With<PredictionOf>, With<Car>, Without<Raft>, Without<Crab>)>,
        Query<
            (&Position, &ConstantMotor),
            (With<PredictionOf>, With<Raft>, Without<Car>, Without<Crab>),
        >,
    )>,
    mut player_inputs: ResMut<QueuedInputs>,
) {
    let mut did_tick = false;
    for ClientTickEvent(client_tick) in tick_reader.iter() {
        did_tick = true;
        for (entity, action) in player_inputs.0.iter() {
            let Ok(SourceOf(prediction)) = source_query.get(*entity) else { continue };

            if let Ok((_, mut position, mut motor)) = player_query_set.p0().get_mut(*prediction) {
                let mut input_message = InputMessage::new(Some(*action));
                input_message.entity.set(&client, &entity);
                // Send command to server
                client.send_tick_buffer_message::<PlayerInputChannel, InputMessage>(
                    client_tick,
                    &input_message,
                );
                process_input(input_message.action, &mut position, &mut motor);
            }
        }
        tick::tick_step_motors(&mut player_query_set.p1());
        tick::tick_constant_motors(&mut objects_query_set.p0());
        tick::tick_road_collisions(
            &mut commands,
            &level_query,
            &player_query_set.p2().to_readonly(),
            &objects_query_set.p1(),
        );
        tick::tick_river_collisions(
            &mut commands,
            &level_query,
            &mut player_query_set.p2(),
            &objects_query_set.p2(),
        );
    }
    if did_tick {
        player_inputs.0.clear();
    }
}