lightyear_inputs 0.28.0

//! Module to handle client inputs
//!
//! Client inputs are generated by the user and sent to the server.
//! They have to be handled separately from other messages, for several reasons:
//! - the history of inputs might need to be saved on the client to perform rollback and client-prediction
//! - we not only send the input for tick T, but we also include the inputs for the last N ticks before T. This redundancy helps ensure
//!   that the server isn't missing any client inputs even if a packet gets lost
//! - we must provide [`SystemSet`]s so that the user can order their systems before and after the input handling
//!
//! ### Adding a new input type
//!
//! An input type is an enum that implements the `UserAction` trait.
//! This trait is a marker trait that is used to tell Lightyear that this type can be used as an input.
//! In particular inputs must be `Serialize`, `Deserialize`, `Clone` and `PartialEq`.
//!
//! You can then add the input type by adding the `InputPlugin<InputType>` to your app.
//!
//! ```rust
//! use bevy_ecs::entity::{EntityMapper, MapEntities};
//! use serde::{Serialize, Deserialize};
//!
//! #[derive(Serialize, Deserialize, Clone, PartialEq, Debug)]
//! pub enum MyInput {
//!     Move { x: f32, y: f32 },
//!     Jump,
//!     // we need a variant for "no input", to differentiate between "no input" and "missing input packet"
//!     None,
//! }
//!
//! // every input must implement MapEntities
//! impl MapEntities for MyInput {
//!     fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
//!     }
//! }
//! ```
//!
//! ### Sending inputs
//!
//! There are several steps to use the `InputPlugin`:
//! - (optional) read the inputs from an external signal (mouse click or keyboard press, for instance)
//! - to buffer inputs for each tick. This is done by updating the `ActionState` component in a system.
//!   That system must run in the [`InputSystems::BufferClientInputs`] system set, in the `FixedPreUpdate` stage.
//! - handle inputs in your game logic in systems that run in the `FixedUpdate` schedule. These systems
//!   will read the inputs using the [`InputBuffer`] component.

use crate::config::{InputConfig, SharedInputConfig};
use crate::input_buffer::InputBuffer;
use crate::input_message::{
    ActionStateQueryData, ActionStateSequence, InputMessage, InputSnapshot, InputTarget,
    PerTargetData, StateMut, StateRef,
};
use crate::plugin::InputPlugin;
use crate::{HISTORY_DEPTH, InputChannel};
#[cfg(feature = "metrics")]
use alloc::format;
use alloc::{vec, vec::Vec};
use bevy_app::{
    App, FixedPostUpdate, FixedPreUpdate, Plugin, PostUpdate, PreUpdate, RunFixedMainLoopSystems,
};
use bevy_ecs::entity::MapEntities;
use bevy_ecs::prelude::*;
use bevy_time::{Real, Time, Timer, TimerMode};
use bevy_utils::prelude::DebugName;
use lightyear_connection::host::HostClient;
use lightyear_core::prelude::*;
use lightyear_core::tick::TickDuration;
#[cfg(feature = "interpolation")]
use lightyear_core::time::TickDelta;

use lightyear_connection::client::Client;
#[cfg(feature = "interpolation")]
use lightyear_interpolation::prelude::*;
use lightyear_messages::plugin::MessageSystems;
#[cfg(feature = "prediction")]
use lightyear_messages::prelude::MessageReceiver;
use lightyear_messages::prelude::MessageSender;
use lightyear_prediction::prelude::*;
use lightyear_replication::prelude::{ControlledBy, PreSpawned};
use lightyear_sync::plugin::SyncSystems;
use lightyear_sync::prelude::client::IsSynced;
use lightyear_sync::prelude::{InputTimeline, InputTimelineConfig};
use lightyear_transport::prelude::ChannelRegistry;
#[allow(unused_imports)]
use tracing::{debug, error, info, trace, warn};

#[deprecated(note = "Use InputSystems instead")]
pub type InputSet = InputSystems;
#[derive(SystemSet, Debug, Hash, PartialEq, Eq, Clone, Copy)]
pub enum InputSystems {
    // RUN-FIXED-MAIN-LOOP UPDATE
    /// Receive the InputMessage from other clients
    ReceiveInputMessages,

    // FIXED PRE UPDATE
    /// System Set where the user should emit InputEvents, they will be buffered in the InputBuffers in the BufferClientInputs set.
    ///
    /// (For Leafwing, there is nothing to do because the ActionState is updated by leafwing)
    /// (For BEI, there is nothing to do because the `Actions<C>` is updated by the BEI plugin)
    WriteClientInputs,
    /// System Set where we update the ActionState and the InputBuffers
    /// - no rollback: we write the ActionState to the InputBuffers
    /// - rollback: we fetch the ActionState value from the InputBuffers
    BufferClientInputs,

    // FIXED POST UPDATE
    /// Prepare a message for the server with the current tick's inputs.
    /// (we do this in the FixedUpdate schedule because if the simulation is slow (e.g. 10Hz)
    /// we don't want to send an InputMessage every frame)
    PrepareInputMessage,
    // TODO: could this run in RunFixedMainLoop::AfterFixedMainLoop?
    /// Restore the ActionState for the correct tick (without InputDelay) from the buffer
    RestoreInputs,

    // POST UPDATE
    /// Update the metadata where we store information about remote inputs that we received, for instance
    /// the [`LastConfirmedInput`] across all remote clients
    UpdateRemoteInputTicks,
    /// System Set to prepare the input message
    SendInputMessage,
    /// Clean up old values to prevent the buffers from growing indefinitely
    CleanUp,
}

/// Client-side plugin that buffers local action state each tick and sends
/// compressed input messages to the server.
///
/// The plugin also handles:
/// - **Input delay**: buffering inputs ahead of time so the server receives
///   them before it needs them.
/// - **Redundancy**: each message includes the last N ticks of input to
///   recover from packet loss.
/// - **Remote player prediction**: if `rebroadcast_inputs` is enabled,
///   processes input messages received from the server for other players.
pub struct ClientInputPlugin<S: ActionStateSequence> {
    config: InputConfig<S::Action>,
}

impl<S: ActionStateSequence> ClientInputPlugin<S> {
    pub fn new(config: InputConfig<S::Action>) -> Self {
        Self { config }
    }
}

impl<S: ActionStateSequence> ClientInputPlugin<S> {
    fn default() -> Self {
        Self::new(InputConfig::default())
    }
}

impl<S: ActionStateSequence + MapEntities> Plugin for ClientInputPlugin<S> {
    fn build(&self, app: &mut App) {
        if !app.is_plugin_added::<InputPlugin<S>>() {
            app.add_plugins(InputPlugin::<S>::default());
        }
        app.init_resource::<SharedInputConfig>();
        app.insert_resource(self.config);
        app.init_resource::<MessageBuffer<S>>();

        // SETS

        // NOTE: this is subtle! We receive remote players messages after
        //  RunFixedMainLoopSystems::BeforeFixedMainLoop to ensure that the local leafwing `states` have
        //  been switched to the `fixed_update` state (see https://github.com/Leafwing-Studios/leafwing-input-manager/blob/v0.16/src/plugin.rs#L170)
        //  We can move this system back in PreUpdate if we drop leafwing support.
        //  Conveniently, this also ensures that we run this after MessageSet::Receive.
        app.configure_sets(
            PreUpdate,
            InputSystems::ReceiveInputMessages
                .before(RunFixedMainLoopSystems::FixedMainLoop)
                .after(RunFixedMainLoopSystems::BeforeFixedMainLoop),
        );

        app.configure_sets(
            FixedPreUpdate,
            (
                InputSystems::WriteClientInputs,
                InputSystems::BufferClientInputs,
            )
                .chain(),
        );
        app.configure_sets(FixedPostUpdate, InputSystems::RestoreInputs);
        app.configure_sets(
            PostUpdate,
            ((
                InputSystems::PrepareInputMessage,
                SyncSystems::Sync,
                // run after SyncSet to make sure that the TickEvents are handled
                // and that the interpolation_delay injected in the message are correct
                InputSystems::SendInputMessage,
                InputSystems::CleanUp,
                MessageSystems::Send,
            )
                .chain(),),
        );

        // SYSTEMS
        #[cfg(feature = "prediction")]
        if self.config.rebroadcast_inputs {
            // NOTE: we do NOT need to run this after RunFixedMainLoopSystems::BeforeFixedMainLoop to ensure that the
            //  local leafwing `states` have been switched to the `fixed_update` state (see
            //  https://github.com/Leafwing-Studios/leafwing-input-manager/blob/v0.16/src/plugin.rs#L170)
            //  because when applying the diffs we manually update the fixed_update_state .

            //  because when we apply the diffs, we start by taking the initial `start_state` snapshot from the message.
            //  That snapshot has both `state` = `fixed_update_state` since it was buffered by the client
            //  during FixedUpdate.
            //  Then, the diffs will update `state`, and that value will be stored in the buffer.
            //
            //  The ActionState will only be modified during FixedPreUpdate in `get_action_state`, so we will have
            //  `state` = `fixed_update_state` in the ActionState component.
            app.configure_sets(
                PreUpdate,
                InputSystems::ReceiveInputMessages
                    .after(MessageSystems::Receive)
                    // NOTE: there is no point in running this after ReplicationSet::Receive, because even if we spawned the entity
                    //  before the corresponding input entity, entity-mapping is applied in MessageSet::Receive and would fail
                    .before(RollbackSystems::Check),
            );
            app.add_systems(
                PreUpdate,
                receive_remote_player_input_messages::<S>
                    .in_set(InputSystems::ReceiveInputMessages),
            );
            app.configure_sets(
                PostUpdate,
                InputSystems::UpdateRemoteInputTicks.after(SyncSystems::Sync),
            );
            app.add_systems(
                PostUpdate,
                update_last_confirmed_input::<S>.in_set(InputSystems::UpdateRemoteInputTicks),
            );
        }
        app.add_systems(
            FixedPreUpdate,
            (buffer_action_state::<S>, get_action_state::<S>)
                .chain()
                .in_set(InputSystems::BufferClientInputs),
        );
        app.add_systems(
            // This should be fine because the user should use the ActionState value in FixedUpdate
            FixedPostUpdate,
            (
                // we want:
                // - to write diffs for the delayed tick (in the next FixedUpdate run), so re-fetch the delayed action-state
                //   this is required in case the FixedUpdate schedule runs multiple times in a frame,
                // - next frame's input-map (in PreUpdate) to act on the delayed tick, so re-fetch the delayed action-state
                get_delayed_action_state::<S>.in_set(InputSystems::RestoreInputs),
            ),
        );
        app.add_systems(
            PostUpdate,
            (
                // TODO: instead, store directly in MessageSender, SyncSet::Sync before MessageSet::Send and register an observer to update the ticks from MessageSender directly!
                prepare_input_message::<S>.in_set(InputSystems::PrepareInputMessage),
                clean_buffers::<S>.in_set(InputSystems::CleanUp),
                send_input_messages::<S>.in_set(InputSystems::SendInputMessage),
            ),
        );
        // if the client tick is updated because of a desync, update the ticks in the input buffers
        app.add_observer(receive_tick_events::<S>);
    }
}

// equivalent to &ActionState<S::Action>

/// Write the value of the ActionState in the InputBuffer.
/// (so that we can pull it for rollback or for delayed inputs)
///
/// If we have input-delay, we will store the current ActionState in the buffer at the delayed-tick,
/// and we will pull ActionStates from the buffer instead of just using the ActionState component directly.
///
/// We do not need to buffer inputs during rollback, as they have already been buffered!
///
/// We only buffer inputs after the timelines have synced, otherwise it would be possible to buffer
/// inputs that would not arrive on time in the server (i.e. the server would simulate tick T
/// before receiving the input for tick T).
fn buffer_action_state<S: ActionStateSequence>(
    local_timeline: Res<LocalTimeline>,
    // we buffer inputs even for the Host-Server so that
    // 1. the HostServer client can broadcast inputs to other clients
    // 2. the HostServer client can have input delay
    input_timeline: Single<
        (Entity, &InputTimeline),
        (
            With<Client>,
            With<IsSynced<InputTimeline>>,
            Without<Rollback>,
        ),
    >,
    mut action_state_query: Query<
        (
            Entity,
            StateRef<S>,
            &mut InputBuffer<S::Snapshot, S::Action>,
            Option<&ControlledBy>,
        ),
        (With<S::Marker>, Allow<PredictionDisable>),
    >,
) {
    let (client_entity, input_timeline) = input_timeline.into_inner();
    let current_tick = local_timeline.tick();
    let tick = current_tick + input_timeline.input_delay() as i32;
    for (entity, action_state, mut input_buffer, controlled_by) in action_state_query.iter_mut() {
        if controlled_by.is_some_and(|controlled_by| controlled_by.owner != client_entity) {
            continue;
        }
        input_buffer.set(tick, S::to_snapshot(action_state));
        trace!(
            ?entity,
            // ?action_state,
            ?current_tick,
            delayed_tick = ?tick,
            input_buffer = %input_buffer.as_ref(),
            "set action state in input buffer",
        );
        trace!(
            target: "lightyear_debug::input",
            kind = "buffer_action_state",
            schedule = "FixedPreUpdate",
            sample_point = "FixedPreUpdate",
            entity = ?entity,
            action = ?DebugName::type_name::<S::Action>(),
            local_tick = current_tick.0,
            input_tick = tick.0,
            buffer_len = input_buffer.len(),
            input_buffer = %input_buffer.as_ref(),
            "buffered local action state"
        );
        #[cfg(feature = "metrics")]
        {
            metrics::gauge!(format!(
                "inputs::{}::{}::buffer_size",
                core::any::type_name::<S::Action>(),
                entity
            ))
            .set(input_buffer.len() as f64);
        }
    }
}

/// Retrieve the ActionState for the current tick.
fn get_action_state<S: ActionStateSequence>(
    tick_duration: Res<TickDuration>,
    config: Res<InputConfig<S::Action>>,
    local_timeline: Res<LocalTimeline>,
    // NOTE: we skip this for host-client because a similar system does the same thing
    //  in the server, and also clears the buffers
    sender: Query<
        (&InputTimeline, &InputTimelineConfig, Has<Rollback>),
        (With<Client>, Without<HostClient>),
    >,

    // NOTE: we want to apply the Inputs for BOTH the local player and remote player
    // - local player: we need to get the input from the InputBuffer because of input delay
    // - remote player: during rollbacks, we need to fetch the ActionState from the InputBuffer
    // (for the remote players, we update the ActionState as soon as we receive the RemoteMessage.)
    //  TODO: We could maybe have some decay logic where the input decays to the middle)
    mut action_state_query: Query<
        (
            Entity,
            StateMut<S>,
            &mut InputBuffer<S::Snapshot, S::Action>,
            Has<S::Marker>,
        ),
        Allow<PredictionDisable>,
    >,
) {
    let Ok((input_timeline, input_config, is_rollback)) = sender.single() else {
        return;
    };
    let input_delay = input_timeline.input_delay() as i32;
    let tick = local_timeline.tick();
    if is_rollback && config.ignore_rollbacks {
        return;
    }
    // TODO!: if config.rebroadcast = False, we don't need to handle remote players, try to encode that statically!
    for (entity, action_state, mut input_buffer, is_local) in action_state_query.iter_mut() {
        if !is_rollback && is_local && input_delay == 0 {
            // for local clients, if there is no rollback and no input_delay:
            // we just buffered the input for the current tick so the action state is already up to date
            continue;
        }

        // NOTE: for remote players:
        // - if we receive a remote input from an older tick (because of prediction), then upon receipt we update the ActionState
        //   immediately so that we can trigger a rollback (in `receive_input_message`) and the ActionState is correctly set
        //   to the past state for the rollback
        // - we cannot just rely on that, because in some cases (lockstep), we receive the remote input in the future, so we need
        //   to read from the buffer to restore the ActionState to the correct tick.
        // - we cannot use `is_lockstep` to check if we receive inputs in the future, because there are cases in non-lockstep where
        //   we could receive some inputs from the future, if we had a high enough input delay.

        // NOTE: we use `get` and not `get_predict ` here.
        // This means that we try to get the value for this exact tick.
        // - For local inputs with input_delay: our current state is in the future, so we need to fetch the exact value from the buffer
        // - For remote inputs with lockstep: the last input is in the future, so we need to fetch the exact value from the buffer
        // - For remote inputs without lockstep: we might receive a message that updates our input buffer, and the last input is
        //   in the past. We already updated the ActionState in `receive_input_message` for that past tick, which means we are
        //   predicting that the action hasn't changed since. We just need to decay it (for rollback or without rollback)
        if let Some(snapshot) = input_buffer.get(tick) {
            // TODO: should we decay_tick the snapshot?
            S::from_snapshot(S::State::into_inner(action_state), snapshot);
            trace!(
                ?entity,
                ?tick,
                ?is_local,
                ?snapshot,
                // ?action_state,
                "fetched action state from input buffer: {:?}",
                // action_state.get_pressed(),
                input_buffer
            );
            trace!(
                target: "lightyear_debug::input",
                kind = "get_action_state",
                schedule = "FixedPreUpdate",
                sample_point = "FixedPreUpdate",
                entity = ?entity,
                action = ?DebugName::type_name::<S::Action>(),
                local_tick = tick.0,
                input_tick = tick.0,
                is_local,
                is_rollback,
                snapshot = ?snapshot,
                buffer_len = input_buffer.len(),
                input_buffer = %*input_buffer,
                "restored action state from input buffer"
            );
        } else if !is_local && config.rebroadcast_inputs {
            if input_config.is_lockstep() {
                error!("We are in lockstep mode but didn't receive an input for tick {tick:?}!");
            }
            // we are here if:
            // - we are in rollback and we reach a tick further than the last tick we received from the remote
            // - we are not in rollback, in which case we want to decay the ActionState
            let mut snapshot = S::to_snapshot(S::State::as_read_only(&action_state));
            snapshot.decay_tick(tick_duration.0);
            trace!(
                ?entity,
                ?tick,
                "Action = {}, For remote input; no input for tick so we decay the ActionState to: {:?}",
                DebugName::type_name::<S::Action>(),
                snapshot
            );
            trace!(
                target: "lightyear_debug::input",
                kind = "decay_missing_remote_action_state",
                schedule = "FixedPreUpdate",
                sample_point = "FixedPreUpdate",
                entity = ?entity,
                action = ?DebugName::type_name::<S::Action>(),
                local_tick = tick.0,
                input_tick = tick.0,
                is_rollback,
                snapshot = ?snapshot,
                buffer_len = input_buffer.len(),
                "decayed missing remote action state"
            );
            // update the action state with decay
            S::from_snapshot(S::State::into_inner(action_state), &snapshot);
            // add the new snapshot in the buffer
            input_buffer.set(tick, snapshot);
        }
    }
}

/// At the start of the frame, restore the ActionState to the latest-action state in buffer
/// (e.g. the delayed action state) because all inputs (i.e. diffs) are applied to the delayed action-state.
fn get_delayed_action_state<S: ActionStateSequence>(
    timeline: Res<LocalTimeline>,
    sender: Query<
        (Entity, &InputTimeline, Has<Rollback>),
        (With<Client>, With<IsSynced<InputTimeline>>),
    >,
    mut action_state_query: Query<
        (
            Entity,
            StateMut<S>,
            &InputBuffer<S::Snapshot, S::Action>,
            Option<&ControlledBy>,
        ),
        // Filter so that this is only for directly controlled players, not remote players
        (With<S::Marker>, Allow<PredictionDisable>),
    >,
) {
    let Ok((client_entity, input_timeline, is_rollback)) = sender.single() else {
        return;
    };
    let input_delay_ticks = input_timeline.input_delay() as i32;
    if is_rollback || input_delay_ticks == 0 {
        return;
    }
    let tick = timeline.tick();
    let delayed_tick = tick + input_delay_ticks;
    for (entity, action_state, input_buffer, controlled_by) in action_state_query.iter_mut() {
        if controlled_by.is_some_and(|controlled_by| controlled_by.owner != client_entity) {
            continue;
        }
        // TODO: lots of clone + is complicated. Shouldn't we just have a DelayedActionState component + resource?
        //  the problem is that the Leafwing Plugin works on ActionState directly...
        if let Some(delayed_action_state) = input_buffer.get(delayed_tick) {
            S::from_snapshot(S::State::into_inner(action_state), delayed_action_state);
            trace!(
                ?entity,
                ?delayed_tick,
                // ?action_state,
                "fetched delayed action state from input buffer: {}",
                input_buffer
            );
            trace!(
                target: "lightyear_debug::input",
                kind = "get_delayed_action_state",
                schedule = "RunFixedMainLoop",
                sample_point = "RunFixedMainLoop",
                entity = ?entity,
                action = ?DebugName::type_name::<S::Action>(),
                local_tick = tick.0,
                input_tick = delayed_tick.0,
                snapshot = ?delayed_action_state,
                buffer_len = input_buffer.len(),
                input_buffer = %input_buffer,
                "restored delayed action state"
            );
        }
        // TODO: if we don't find an ActionState in the buffer, should we reset the delayed one to default?
    }
}

/// System that removes old entries from the InputBuffer
fn clean_buffers<S: ActionStateSequence>(
    timeline: Res<LocalTimeline>,
    // NOTE: we skip this for host-client because the get_action_state system on the server
    //  also clears the buffers
    sender: Query<(), (With<Client>, With<InputTimeline>, Without<HostClient>)>,
    prediction_manager: Option<Single<&PredictionManager, With<Client>>>,
    mut input_buffer_query: Query<
        &mut InputBuffer<S::Snapshot, S::Action>,
        Allow<PredictionDisable>,
    >,
) {
    if sender.single().is_err() {
        return;
    }
    let old_tick = timeline.tick() - input_history_depth(prediction_manager.as_deref().copied());

    // trace!(
    //     "popping all input buffers since old tick: {old_tick:?}",
    // );
    for mut input_buffer in input_buffer_query.iter_mut() {
        input_buffer.pop(old_tick);
        //  for now do NOT spawn Transform, instead directly use Position/Rotation!
    }
}

fn input_history_depth(prediction_manager: Option<&PredictionManager>) -> u32 {
    prediction_manager
        .map(|manager| u32::from(manager.rollback_policy.max_rollback_ticks) + 1)
        .unwrap_or(0)
        .max(HISTORY_DEPTH)
}

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

    #[test]
    fn input_history_depth_covers_prediction_rollback_window() {
        assert_eq!(input_history_depth(None), HISTORY_DEPTH);

        let mut manager = PredictionManager {
            rollback_policy: RollbackPolicy {
                max_rollback_ticks: 100,
                ..Default::default()
            },
            ..Default::default()
        };
        assert_eq!(input_history_depth(Some(&manager)), 101);

        manager.rollback_policy.max_rollback_ticks = 5;
        assert_eq!(input_history_depth(Some(&manager)), HISTORY_DEPTH);
    }
}

// TODO: is this actually necessary? The sync happens in PostUpdate,
//  so maybe it's ok if the InputMessages contain the pre-sync tick! (since those inputs happened
//  before the sync). If it's not needed, send the messages directly in FixedPostUpdate!
//  Actually maybe it is, because the send-tick on the server will be updated.
/// Buffer that will store the InputMessages we want to write this frame.
///
/// We need this because:
/// - we write the InputMessages during FixedPostUpdate
/// - we apply the TickUpdateEvents (from doing sync) during PostUpdate, which might affect the ticks from the InputMessages.
///   During this phase, we want to update the tick of the InputMessages that we wrote during FixedPostUpdate.
#[derive(Debug, Resource)]
pub(crate) struct MessageBuffer<S>(Vec<InputMessage<S>>);

impl<A> Default for MessageBuffer<A> {
    fn default() -> Self {
        Self(vec![])
    }
}

/// Take the input buffer, and prepare the input message to send to the server.
///
/// This runs once per frame in PostUpdate. It needs to run before SyncSet::Sync, because we buffer
/// the input in a MessageBuffer, and if a SyncEvent triggers, we want to adjust the ticks from the InputMessages.
fn prepare_input_message<S: ActionStateSequence>(
    mut message_buffer: ResMut<MessageBuffer<S>>,
    tick_duration: Res<TickDuration>,
    timeline: Res<LocalTimeline>,
    input_config: Res<InputConfig<S::Action>>,
    sender: Single<
        (Entity, &InputTimeline, Has<HostClient>),
        // the host-client doesn't need to send input messages since the ActionState is already on the entity
        // unless we want to rebroadcast the HostClient inputs to other clients (in which
        // case we prepare the input-message, which will be send_local to the server)
        (
            With<Client>,
            With<IsSynced<InputTimeline>>,
            Without<Rollback>,
        ),
    >,
    _channel_registry: Res<ChannelRegistry>,
    input_buffer_query: Query<
        (
            Entity,
            &InputBuffer<S::Snapshot, S::Action>,
            Option<&PreSpawned>,
            Option<&ControlledBy>,
        ),
        (With<S::Marker>, Allow<PredictionDisable>),
    >,
    real_time: Res<Time<Real>>,
    mut send_timer: Local<Option<Timer>>,
) {
    // Only prepare input every send-interval.
    if !input_config.send_interval.is_zero() {
        let timer = send_timer
            .get_or_insert_with(|| Timer::new(input_config.send_interval, TimerMode::Repeating));
        timer.tick(real_time.delta());
        if !timer.is_finished() {
            return;
        }
    }

    let (client_entity, input_timeline, is_host_client) = sender.into_inner();
    #[cfg(not(feature = "prediction"))]
    if is_host_client {
        // if there is not prediction, no need to rebroadcast inputs
        return;
    }

    #[cfg(feature = "prediction")]
    if is_host_client && !input_config.rebroadcast_inputs {
        // the host-client doesn't need to send input messages since the ActionState is already on the entity
        // unless we want to rebroadcast the HostClient inputs to other clients
        return;
    }

    // we send a message from the latest tick that we have available, which is the delayed tick
    let current_tick = timeline.tick();
    let tick = current_tick + input_timeline.input_delay() as i32;
    // TODO: the number of messages should be in SharedConfig
    trace!(delayed_tick = ?tick, ?current_tick, "prepare_input_message");
    trace!(
        target: "lightyear_debug::input",
        kind = "prepare_input_message_start",
        schedule = "PostUpdate",
        sample_point = "PostUpdate",
        action = ?DebugName::type_name::<S::Action>(),
        local_tick = current_tick.0,
        input_tick = tick.0,
        is_host_client,
        "preparing input message"
    );
    // TODO: instead of redundancy, send ticks up to the latest yet ACK-ed input tick
    //  this means we would also want to track packet->message acks for unreliable channels as well, so we can notify
    //  this system what the latest acked input tick is?

    // Send redundant inputs so that if a packet is lost, we can still recover.
    // The size of the input bundle scales with `send_interval`.
    let mut num_ticks: u32 = ((input_config.send_interval.as_nanos() / tick_duration.as_nanos())
        + 1)
    .try_into()
    .unwrap();
    num_ticks *= input_config.packet_redundancy as u32;
    let mut message = InputMessage::<S>::new(tick);
    for (entity, input_buffer, pre_spawned, controlled_by) in input_buffer_query.iter() {
        if controlled_by.is_some_and(|controlled_by| controlled_by.owner != client_entity) {
            continue;
        }
        trace!(
            ?tick,
            ?entity,
            "Preparing input message with buffer: {:?}",
            input_buffer
        );

        let target = if let Some(prespawned) = pre_spawned
            && let Some(hash) = prespawned.hash
        {
            debug!(?hash, ?entity, "Sending input for prespawned entity");
            InputTarget::PreSpawned(hash)
        } else {
            InputTarget::Entity(entity)
        };

        if let Some(state_sequence) = S::build_from_input_buffer(input_buffer, num_ticks, tick) {
            trace!(
                target: "lightyear_debug::input",
                kind = "prepare_input_message_target",
                schedule = "PostUpdate",
                sample_point = "PostUpdate",
                entity = ?entity,
                action = ?DebugName::type_name::<S::Action>(),
                local_tick = current_tick.0,
                input_tick = tick.0,
                num_ticks = num_ticks,
                buffer_len = input_buffer.len(),
                target = ?target,
                states = ?state_sequence,
                "added target data to input message"
            );
            message.inputs.push(PerTargetData {
                target,
                states: state_sequence,
            });
        }
    }

    // TODO: revisit this; maybe we should not send an empty message?
    // we send a message even when there are 0 inputs because that itself is information
    debug!(
        ?tick,
        ?num_ticks,
        ?is_host_client,
        "sending input message for {:?}: {}",
        DebugName::type_name::<S::Action>().shortname(),
        message
    );
    trace!(
        target: "lightyear_debug::input",
        kind = "prepare_input_message_finish",
        schedule = "PostUpdate",
        sample_point = "PostUpdate",
        action = ?DebugName::type_name::<S::Action>(),
        local_tick = current_tick.0,
        input_tick = tick.0,
        end_tick = tick.0,
        num_ticks = num_ticks,
        num_targets = message.inputs.len(),
        is_host_client,
        message = ?message,
        "prepared input message"
    );
    message_buffer.0.push(message);

    // NOTE: keep the older input values in the InputBuffer! because they might be needed when we rollback for client prediction
}

#[cfg(feature = "prediction")]
/// Read the InputMessages of other clients from the server to update their InputBuffer and ActionState.
/// This is useful if we want to do client-prediction for remote players.
///
/// If the InputBuffer/ActionState is missing, we will add it.
///
/// We will apply the diffs on the Predicted entity.
fn receive_remote_player_input_messages<S: ActionStateSequence>(
    mut commands: Commands,
    tick_duration: Res<TickDuration>,
    timeline: Res<LocalTimeline>,
    link: Single<
        (
            &mut MessageReceiver<InputMessage<S>>,
            Option<&LastConfirmedInput>,
            &PredictionManager,
        ),
        // the host-client won't receive input messages from the Server
        (
            With<Client>,
            With<IsSynced<InputTimeline>>,
            Without<HostClient>,
        ),
    >,
    mut predicted_query: Query<
        Option<&mut InputBuffer<S::Snapshot, S::Action>>,
        (Without<S::Marker>, Allow<PredictionDisable>),
    >,
    prespawned: Query<(Entity, &PreSpawned)>,
) {
    let (mut receiver, last_confirmed_input, prediction_manager) = link.into_inner();
    let tick = timeline.tick();
    let mut received_relevant_input = false;
    receiver.receive().for_each(|message| {
        trace!(?message.end_tick, ?message, "received remote input message for action: {:?}", DebugName::type_name::<S::Action>());
        trace!(
            target: "lightyear_debug::input",
            kind = "remote_input_message_recv",
            schedule = "PreUpdate",
            sample_point = "PreUpdate",
            action = ?DebugName::type_name::<S::Action>(),
            local_tick = tick.0,
            end_tick = message.end_tick.0,
            num_targets = message.inputs.len(),
            rebroadcast = message.rebroadcast,
            message = ?message,
            "received remote player input message"
        );
        for target_data in message.inputs {
            let Some(entity) = (match target_data.target {
                InputTarget::Entity(entity) => {
                    Some(entity)
                }
                InputTarget::PreSpawned(hash) => {
                    prespawned
                        .iter()
                        .find_map(|(e, p)| p.hash.is_some_and(|h| h == hash).then_some(e))
                }
            }) else {
                if message.rebroadcast {
                    debug!(
                        target = ?target_data.target,
                        end_tick = ?message.end_tick,
                        "ignored stale remote player input message for unmapped entity"
                    );
                } else {
                    warn!("Could not find entity in entity_map for remote player input message {:?}", target_data.target);
                }
                continue;
            };
            debug!(
                ?tick, ?message.end_tick,
                "received remote client input message for entity: {:?}. Applying to diff buffer.",
                entity
            );
            trace!(
                target: "lightyear_debug::input",
                kind = "remote_input_target",
                schedule = "PreUpdate",
                sample_point = "PreUpdate",
                entity = ?entity,
                action = ?DebugName::type_name::<S::Action>(),
                local_tick = tick.0,
                end_tick = message.end_tick.0,
                target = ?target_data.target,
                states = ?target_data.states,
                "applying remote input target data"
            );
            // TODO: careful of entity collisions!
            let Ok(input_buffer) = predicted_query.get_mut(entity) else {
                if message.rebroadcast {
                    debug!(
                        ?entity,
                        ?target_data.states,
                        end_tick = ?message.end_tick,
                        "ignored stale remote player input message for unrecognized entity"
                    );
                } else {
                    error!(?entity, ?target_data.states, end_tick = ?message.end_tick, "received input message for unrecognized entity");
                }
                continue
            };
            trace!(predicted=?entity, end_tick = ?message.end_tick, "update action diff buffer for remote player PREDICTED using input message");

            if let Some(mut input_buffer) = input_buffer {
                // do not parse the remote message our current Buffer end_tick is later than the message end_tick
                // this can happen if we receive multiple messages out of order.
                if input_buffer.last_remote_tick.is_some_and(|t| t >= message.end_tick) {
                    trace!("Ignoring input message because our current last_remote_tick {:?} is more recent than the remote_end_tick {:?}", input_buffer.last_remote_tick, message.end_tick);
                    trace!(
                        target: "lightyear_debug::input",
                        kind = "remote_input_ignored_stale",
                        schedule = "PreUpdate",
                        sample_point = "PreUpdate",
                        entity = ?entity,
                        action = ?DebugName::type_name::<S::Action>(),
                        local_tick = tick.0,
                        end_tick = message.end_tick.0,
                        last_remote_tick = ?input_buffer.last_remote_tick,
                        "ignored stale remote input message"
                    );
                    continue
                }
                received_relevant_input = true;
                update_buffer_from_remote_player_message::<S>(
                    target_data.states,
                    &mut input_buffer,
                    tick,
                    message.end_tick,
                    entity,
                    prediction_manager,
                    *tick_duration
                );
            } else {
                // add the ActionState and InputBuffer if they are missing
                let mut input_buffer = InputBuffer::<S::Snapshot, S::Action>::default();
                received_relevant_input = true;
                update_buffer_from_remote_player_message::<S>(
                    target_data.states,
                    &mut input_buffer,
                    tick,
                    message.end_tick,
                    entity,
                    prediction_manager,
                    *tick_duration
                );
                // Initialize ActionState from the latest buffered input rather
                // than from base_value(). When the remote player's end_tick is
                // in the past, get_action_state will call get(current_tick)
                // which returns None and leaves ActionState unchanged. If we
                // initialized to base_value() the player would simulate with
                // empty/released inputs until the buffer catches up.
                let mut action_state = S::State::base_value();
                if let Some(last) = input_buffer.get_last() {
                    S::from_snapshot(S::State::as_mut(&mut action_state), last);
                }
                commands.entity(entity).insert((
                    input_buffer,
                    action_state,
                ));
            };
        }
    });

    if let Some(last_confirmed_input) = last_confirmed_input
        && received_relevant_input
    {
        last_confirmed_input
            .received_any_messages
            .store(true, bevy_platform::sync::atomic::Ordering::Relaxed);
    }
}

#[cfg(feature = "prediction")]
/// Update the LastConfirmedInput tick by looking at latest tick buffered in each InputBuffer.
///
/// This runs in PostUpdate, and not in PreUpdate, because we want the rollback to use the previous LastConfirmedInput tick.
/// For example if the last confirmed input was at tick 10, and we received an InputMessage with end_tick 14, we want to rollback to tick 10
/// (and not to tick 14) because we need to potentially re-apply inputs for ticks 11, 12, 13, 14.
fn update_last_confirmed_input<S: ActionStateSequence>(
    timeline: Res<LocalTimeline>,
    last_confirmed_input: Single<
        (&mut LastConfirmedInput, &InputTimelineConfig),
        (With<Client>, With<IsSynced<InputTimeline>>),
    >,
    predicted_query: Query<
        &InputBuffer<S::Snapshot, S::Action>,
        (Without<S::Marker>, Allow<PredictionDisable>),
    >,
) {
    let (mut last_confirmed_input, input_config) = last_confirmed_input.into_inner();
    let tick = timeline.tick();
    // in lockstep mode, we don't need last confirmed input because we always have all inputs for a given tick.
    // we will just use the current tick as the last confirmed input tick
    if input_config.is_lockstep() {
        last_confirmed_input.tick.set_if_lower(tick);
        return;
    }
    // TODO: how to handle multiple actions S?

    // find the earliest last_confirmed_tick for each client
    // (we replaced LastConfirmedInput in  with a high tick to avoid any tick-wrapping issues)
    last_confirmed_input.received_for_all_clients = true;
    predicted_query.iter().for_each(|buffer| {
        // if we received any messages, we update the LastConfirmedInput
        // (this is used to determine the last confirmed tick for each client)
        if let Some(end_tick) = buffer.last_remote_tick {
            last_confirmed_input.tick.set_if_lower(end_tick);
        } else {
            last_confirmed_input.received_for_all_clients = false;
        }
    });
    trace!(
        target: "lightyear_debug::input",
        kind = "last_confirmed_input",
        schedule = "PostUpdate",
        sample_point = "PostUpdate",
        action = ?DebugName::type_name::<S::Action>(),
        local_tick = tick.0,
        confirmed_tick = last_confirmed_input.tick.get().0,
        "updated LastConfirmedInput"
    );
}

#[cfg(feature = "prediction")]
fn update_buffer_from_remote_player_message<S: ActionStateSequence>(
    sequence: S,
    input_buffer: &mut InputBuffer<S::Snapshot, S::Action>,
    tick: Tick,
    end_tick: Tick,
    entity: Entity,
    prediction_manager: &PredictionManager,
    tick_duration: TickDuration,
) {
    // we don't need to update the ActionState here because:
    // - if it's in the future, we will fetch the ActionState from the buffer in the `get_action_state` system when
    //   we reach that tick.
    // - if it's in the past, we will start a rollback and fetch the correct ActionState from the buffer in the
    //   `get_action_state` system
    if let Some(mismatch) = sequence.update_buffer(input_buffer, end_tick, tick_duration.0) {
        trace!(
            target: "lightyear_debug::input",
            kind = "remote_input_buffer_update",
            schedule = "PreUpdate",
            sample_point = "PreUpdate",
            entity = ?entity,
            action = ?DebugName::type_name::<S::Action>(),
            local_tick = tick.0,
            end_tick = end_tick.0,
            mismatch_tick = mismatch.0,
            buffer_len = input_buffer.len(),
            last_remote_tick = ?input_buffer.last_remote_tick,
            input_buffer = %input_buffer,
            "updated remote input buffer with mismatch"
        );
        // find the earliest mismatch tick across all clients
        // NOTE: if the mismatch tick is more recent than our current tick, then it's not a mismatch!
        //  it just means that we are receiving a remote tick in advance of simulating that tick.
        //  (for example in lockstep mode, we should have all player inputs for tick T before simulating tick T,
        //  so we will receive those inputs in advance)
        if let RollbackMode::Check = prediction_manager.rollback_policy.input
            && mismatch <= tick
        {
            debug!(
                ?entity,
                ?tick,
                ?end_tick,
                ?mismatch,
                "Mismatch detected for remote player input message!",
            );
            prediction_manager
                .earliest_mismatch_input
                .has_mismatches
                .store(true, bevy_platform::sync::atomic::Ordering::Relaxed);
            prediction_manager
                .earliest_mismatch_input
                .tick
                .set_if_lower(mismatch);
        }

        #[cfg(feature = "metrics")]
        {
            metrics::counter!(format!(
                "inputs::{}::remote_player::receive",
                DebugName::type_name::<S::Action>(),
            ))
            .increment(1);
            let margin = input_buffer.last_remote_tick.unwrap() - tick;
            metrics::gauge!(format!(
                "inputs::{}::remote_player::{}::buffer_margin",
                DebugName::type_name::<S::Action>(),
                entity
            ))
            .set(margin as f64);
            metrics::gauge!(format!(
                "inputs::{}::remote_player::{}::buffer_size",
                DebugName::type_name::<S::Action>(),
                entity
            ))
            .set(input_buffer.len() as f64);
        }
    };
    trace!(
        target: "lightyear_debug::input",
        kind = "remote_input_buffer_state",
        schedule = "PreUpdate",
        sample_point = "PreUpdate",
        entity = ?entity,
        action = ?DebugName::type_name::<S::Action>(),
        local_tick = tick.0,
        end_tick = end_tick.0,
        buffer_len = input_buffer.len(),
        last_remote_tick = ?input_buffer.last_remote_tick,
        input_buffer = %input_buffer,
        "remote input buffer state after message"
    );
}

/// Drain the messages from the buffer and send them to the server
fn send_input_messages<S: ActionStateSequence>(
    input_config: Res<InputConfig<S::Action>>,
    mut message_buffer: ResMut<MessageBuffer<S>>,
    sender: Single<
        (&mut MessageSender<InputMessage<S>>, Has<HostClient>),
        (With<Client>, With<IsSynced<InputTimeline>>),
    >,
    #[cfg(feature = "interpolation")] interpolation_query: Single<
        (&InputTimeline, &InterpolationTimeline),
        (
            With<Client>,
            With<IsSynced<InterpolationTimeline>>,
            With<IsSynced<InputTimeline>>,
        ),
    >,
) {
    let (mut sender, is_host_client) = sender.into_inner();

    #[cfg(not(feature = "prediction"))]
    if is_host_client {
        message_buffer.0.clear();
        return;
    }
    #[cfg(feature = "prediction")]
    if is_host_client && !input_config.rebroadcast_inputs {
        // the host-client doesn't need to send input messages since the ActionState is already on the entity
        // unless we want to rebroadcast the HostClient inputs to other clients
        message_buffer.0.clear();
        return;
    }
    trace!(
        "Number of input messages to send: {:?}",
        message_buffer.0.len()
    );
    trace!(
        target: "lightyear_debug::input",
        kind = "send_input_messages",
        schedule = "PostUpdate",
        sample_point = "PostUpdate",
        action = ?DebugName::type_name::<S::Action>(),
        num_messages = message_buffer.0.len(),
        is_host_client,
        "sending buffered input messages"
    );

    #[cfg(feature = "interpolation")]
    let interpolation_delay = {
        let (input_timeline, interpolation_timeline) = interpolation_query.into_inner();

        // NOTE: this can be negative because of input-delay!
        let mut delay = input_timeline.now() - interpolation_timeline.now();
        if delay.is_negative() {
            delay = TickDelta::from(Tick(0));
        }
        InterpolationDelay {
            delay: delay.into(),
        }
    };

    #[cfg_attr(
        not(feature = "interpolation"),
        expect(unused_mut, reason = "Used by expression behind feature flag")
    )]
    for mut message in message_buffer.0.drain(..) {
        // if lag compensation is enabled, we send the current delay to the server
        // (this runs here because the delay is only correct after the SyncSet has run)
        // TODO: or should we actually use the interpolation_delay BEFORE SyncSet
        //  because the user is reacting to stuff from the previous frame?
        #[cfg(feature = "interpolation")]
        if input_config.lag_compensation {
            message.interpolation_delay = Some(interpolation_delay);
        }
        sender.send::<InputChannel>(message);
    }
}

/// In case the client tick changes suddenly, we also update the InputBuffer accordingly
fn receive_tick_events<S: ActionStateSequence>(
    trigger: On<SyncEvent<InputTimelineConfig>>,
    mut message_buffer: ResMut<MessageBuffer<S>>,
    clients: Query<(), With<Client>>,
    mut input_buffer_query: Query<
        (
            &mut InputBuffer<S::Snapshot, S::Action>,
            Option<&ControlledBy>,
        ),
        Allow<PredictionDisable>,
    >,
) {
    if clients.get(trigger.entity).is_err() {
        return;
    }
    let delta = trigger.tick_delta;
    for (mut input_buffer, controlled_by) in input_buffer_query.iter_mut() {
        if controlled_by.is_some_and(|controlled_by| controlled_by.owner != trigger.entity) {
            continue;
        }
        if let Some(start_tick) = input_buffer.start_tick {
            input_buffer.start_tick = Some(start_tick + delta);
            debug!(
                "Receive tick snap event {:?}. Updating input buffer start_tick to {:?}!",
                trigger.event(),
                input_buffer.start_tick
            );
        }
        if let Some(last_remote_tick) = input_buffer.last_remote_tick {
            input_buffer.last_remote_tick = Some(last_remote_tick + delta);
        }
    }
    for message in message_buffer.0.iter_mut() {
        message.end_tick = message.end_tick + delta;
    }
}