use bevy::ecs::{
    event::Event,
    system::{Commands, EntityCommands},
    world::{EntityWorldMut, World},
};
use lightyear::prelude::Disconnected;

pub trait CommandExt {
    fn trigger_server(&mut self, event: impl Event) -> &mut Self;
}

impl<'w, 's> CommandExt for Commands<'w, 's> {
    fn trigger_server(&mut self, event: impl Event) -> &mut Self {
        self.queue(|world: &mut World| {
            let mut query_state = world.query::<&Disconnected>();
            if cfg!(feature = "server") || !query_state.query(world).is_empty() {
                world.trigger(event);
            }
        });
        self
    }
}

pub trait EntityCommandExt {
    fn trigger_server(&mut self, event: impl Event) -> &mut Self;
}

impl<'w> EntityCommandExt for EntityCommands<'w> {
    fn trigger_server(&mut self, event: impl Event) -> &mut Self {
        self.queue(|mut entity: EntityWorldMut| {
            let mut query_state = entity.world_scope(|world| world.query::<&Disconnected>());
            if cfg!(feature = "server") || !query_state.query(entity.world()).is_empty() {
                entity.trigger(event);
            }
        })
    }
}
