//! Extension to [`EntityCommands`] to modify `bevy_hierarchy` hierarchies
//! while preserving [`GlobalTransform`].

use crate::prelude::{GlobalTransform, Transform};
use bevy_ecs::{
    entity::Entity,
    system::EntityCommands,
    world::{EntityWorldMut, World},
};
use bevy_hierarchy::BuildChildren;

/// Collection of methods similar to [`BuildChildren`], but preserving each
/// entity's [`GlobalTransform`].
pub trait BuildChildrenTransformExt {
    /// Change this entity's parent while preserving this entity's [`GlobalTransform`]
    /// by updating its [`Transform`].
    ///
    /// See [`BuildChildren::set_parent`] for a method that doesn't update the [`Transform`].
    ///
    /// Note that both the hierarchy and transform updates will only execute
    /// the next time commands are applied
    /// (during [`ApplyDeferred`](bevy_ecs::schedule::ApplyDeferred)).
    fn set_parent_in_place(&mut self, parent: Entity) -> &mut Self;

    /// Make this entity parentless while preserving this entity's [`GlobalTransform`]
    /// by updating its [`Transform`] to be equal to its current [`GlobalTransform`].
    ///
    /// See [`BuildChildren::remove_parent`] for a method that doesn't update the [`Transform`].
    ///
    /// Note that both the hierarchy and transform updates will only execute
    /// the next time commands are applied
    /// (during [`ApplyDeferred`](bevy_ecs::schedule::ApplyDeferred)).
    fn remove_parent_in_place(&mut self) -> &mut Self;
}

impl BuildChildrenTransformExt for EntityCommands<'_> {
    fn set_parent_in_place(&mut self, parent: Entity) -> &mut Self {
        self.queue(move |entity: Entity, world: &mut World| {
            if let Ok(mut entity) = world.get_entity_mut(entity) {
                entity.set_parent_in_place(parent);
            }
        })
    }

    fn remove_parent_in_place(&mut self) -> &mut Self {
        self.queue(move |entity: Entity, world: &mut World| {
            if let Ok(mut entity) = world.get_entity_mut(entity) {
                entity.remove_parent_in_place();
            }
        })
    }
}

impl BuildChildrenTransformExt for EntityWorldMut<'_> {
    fn set_parent_in_place(&mut self, parent: Entity) -> &mut Self {
        let child = self.id();
        self.world_scope(|world| {
            world.entity_mut(parent).add_child(child);
            // FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
            let mut update_transform = || {
                let parent = *world.get_entity(parent).ok()?.get::<GlobalTransform>()?;
                let child_global = *world.get_entity(child).ok()?.get::<GlobalTransform>()?;
                let mut child_entity = world.get_entity_mut(child).ok()?;
                let mut child = child_entity.get_mut::<Transform>()?;
                *child = child_global.reparented_to(&parent);
                Some(())
            };
            update_transform();
        });
        self
    }

    fn remove_parent_in_place(&mut self) -> &mut Self {
        let child = self.id();
        self.world_scope(|world| {
            world.entity_mut(child).remove_parent();
            // FIXME: Replace this closure with a `try` block. See: https://github.com/rust-lang/rust/issues/31436.
            let mut update_transform = || {
                let child_global = *world.get_entity(child).ok()?.get::<GlobalTransform>()?;
                let mut child_entity = world.get_entity_mut(child).ok()?;
                let mut child = child_entity.get_mut::<Transform>()?;
                *child = child_global.compute_transform();
                Some(())
            };
            update_transform();
        });
        self
    }
}