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

pub trait CommandExt {
    fn trigger_server<'a, E: Event<Trigger<'a>: Default>>(&mut self, event: E) -> &mut Self;
}

impl<'w, 's> CommandExt for Commands<'w, 's> {
    fn trigger_server<'a, E: Event<Trigger<'a>: Default>>(&mut self, event: E) -> &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<'a, E: EntityEvent<Trigger<'a>: Default>>(
        &mut self,
        event: impl FnOnce(Entity) -> E + Send + Sync + 'static,
    ) -> &mut Self;
}

impl<'w> EntityCommandExt for EntityCommands<'w> {
    fn trigger_server<'a, E: EntityEvent<Trigger<'a>: Default>>(
        &mut self,
        event: impl FnOnce(Entity) -> E + Send + Sync + 'static,
    ) -> &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);
            }
        })
    }
}