Nerfed/Nerfed.Editor/Systems/EditorHierarchyWindow.cs

using ImGuiNET;
using MoonTools.ECS;
using Nerfed.Editor.Components;
using Nerfed.Runtime;
using Nerfed.Runtime.Components;
using Nerfed.Runtime.Util;

namespace Nerfed.Editor.Systems
{
    // Window that draws entities.
    internal class EditorHierarchyWindow : MoonTools.ECS.DebugSystem
    {
        private const ImGuiTreeNodeFlags baseFlags = ImGuiTreeNodeFlags.OpenOnArrow | ImGuiTreeNodeFlags.OpenOnDoubleClick | ImGuiTreeNodeFlags.SpanAvailWidth;

        //private readonly Filter rootEntitiesWithTransformFilter;
        //private readonly Filter rootEntitiesFilterBroken;
        private readonly Filter rootEntitiesFilter;

        private readonly EditorHierachySelectionSystem hierachySelectionSystem;
        private readonly EditorHierachyDragAndDropSystem hierachyDragAndDropSystem;

        public EditorHierarchyWindow(World world) : base(world)
        {
            //rootEntitiesWithTransformFilter = FilterBuilder.Include<LocalTransform>().Exclude<Child>().Build();

            // TODO: this doesn't work.
            //rootEntitiesFilterBroken = FilterBuilder.Exclude<Child>().Build();

            // Maybe the parent/child functions should add a root component when not being a child.
            rootEntitiesFilter = FilterBuilder.Include<Root>().Build();

            // Maybe instead of a root, if we need a component that is always on an entity and has some use we could create something like a VersionComponent which only hold an int.
            // The version would update each time something changes on the entity.
            // Or a EditorComponent, just a component that always gets added when in editor mode.

            hierachySelectionSystem = new EditorHierachySelectionSystem(world);
            hierachyDragAndDropSystem = new EditorHierachyDragAndDropSystem(world);
        }

        public override void Update(TimeSpan delta)
        {
            ImGui.Begin("Hierarchy");

            ImGuiTreeNodeFlags flags = baseFlags;
            flags |= ImGuiTreeNodeFlags.DefaultOpen;

            if (ImGui.TreeNodeEx("World", flags))
            {
                //if (ImGui.BeginDragDropTarget())
                //{
                //    unsafe
                //    {
                //        ImGuiPayloadPtr payload = ImGui.AcceptDragDropPayload($"{nameof(EditorHierarchyWindow)}");
                //        if (payload.NativePtr != null)
                //        {
                //            Entity* data = (Entity*)payload.Data;
                //            Entity child = data[0];

                //            Log.Info($"Dropped {child.ID}");

                //            Transform.RemoveParent(World, child);
                //        }
                //    }
                //    ImGui.EndDragDropTarget();
                //}

                //foreach (Entity entity in rootEntitiesWithTransformFilter.Entities)
                //{
                //    DrawEntityAndChildren(entity);
                //}

                foreach (Entity entity in rootEntitiesFilter.Entities)
                {
                    DrawEntityAndChildren(entity);
                }

                ImGui.TreePop();
            }

            ImGui.End();

            hierachySelectionSystem.Update(delta);
            hierachyDragAndDropSystem.Update(delta);
        }

        private void DrawEntityAndChildren(in Entity entity)
        {
            ImGuiTreeNodeFlags flags = baseFlags;

            if (!World.HasInRelation<ChildParentRelation>(entity))
            {
                flags |= ImGuiTreeNodeFlags.Leaf;
            }

            if (World.Has<SelectedInHierachy>(entity))
            {
                flags |= ImGuiTreeNodeFlags.Selected;
            }

            if (ImGui.TreeNodeEx($"{entity.ID} | {GetTag(entity)}", flags))
            {
                // Selection.
                if (ImGui.IsItemClicked() && !ImGui.IsItemToggledOpen())
                {
                    World.Set(entity, new ClickedInHierachy());
                }

                // Drag and drop.
                if (ImGui.BeginDragDropSource())
                {
                    ImGui.SetDragDropPayload($"{nameof(EditorHierarchyWindow)}", nint.Zero, 0);
                    Set(entity, new PayloadSourceInHierachy());
                    Log.Info($"SetSource {entity.ID}");

                    //unsafe
                    //{
                    //    fixed (Entity* payload = &entity)
                    //    {
                    //        ImGui.SetDragDropPayload($"{nameof(EditorHierarchyWindow)}", (IntPtr)payload, (uint)sizeof(Entity));
                    //    }
                    //}

                    ImGui.EndDragDropSource();
                }

                if (ImGui.BeginDragDropTarget())
                {
                    ImGuiPayloadPtr payload = ImGui.AcceptDragDropPayload($"{nameof(EditorHierarchyWindow)}");
                    unsafe
                    {
                        if (payload.NativePtr != null)
                        {
                            Log.Info($"SetTarget {entity.ID}");
                            Set(entity, new PayloadTargetInHierachy());
                        }
                    }
                    //unsafe
                    //{
                    //    ImGuiPayloadPtr payload = ImGui.AcceptDragDropPayload($"{nameof(EditorHierarchyWindow)}");
                    //    if (payload.NativePtr != null)
                    //    {
                    //        Entity ent = *(Entity*)payload.Data;

                    //        Log.Info($"Dropped {ent.ID}");

                    //        Transform.SetParent(World, ent, entity);
                    //    }
                    //}
                    ImGui.EndDragDropTarget();
                }

                ReverseSpanEnumerator<Entity> childEntities = World.InRelations<ChildParentRelation>(entity);
                foreach (Entity childEntity in childEntities)
                {
                    DrawEntityAndChildren(childEntity);
                }

                ImGui.TreePop();
            }
        }

        // System for handling the selected entities in the hierachy.
        private class EditorHierachySelectionSystem : MoonTools.ECS.System
        {
            private readonly Filter selectedEntities;
            private readonly Filter clickedEntities;

            public EditorHierachySelectionSystem(World world) : base(world)
            {
                selectedEntities = FilterBuilder.Include<SelectedInHierachy>().Build();
                clickedEntities = FilterBuilder.Include<ClickedInHierachy>().Build();
            }

            public override void Update(TimeSpan delta)
            {
                ImGuiIOPtr io = ImGui.GetIO();

                if (!clickedEntities.Empty && !io.KeyCtrl)
                {
                    foreach (Entity entity in selectedEntities.Entities)
                    {
                        Remove<SelectedInHierachy>(entity);
                    }
                }

                foreach (Entity entity in clickedEntities.Entities)
                {
                    // Unselect.
                    if (Has<SelectedInHierachy>(entity))
                    {
                        Remove<SelectedInHierachy>(entity);
                    }
                    // Select.
                    else
                    {
                        Set(entity, new SelectedInHierachy());
                    }

                    Remove<ClickedInHierachy>(entity);
                }
            }
        }

        private class EditorHierachyDragAndDropSystem : MoonTools.ECS.System
        {
            private readonly Filter sourceEntities;
            private readonly Filter targetEntities;

            public EditorHierachyDragAndDropSystem(World world) : base(world)
            {
                sourceEntities = FilterBuilder.Include<PayloadSourceInHierachy>().Build();
                targetEntities = FilterBuilder.Include<PayloadTargetInHierachy>().Build();
            }

            public override void Update(TimeSpan delta)
            {
                if (!targetEntities.Empty)
                {
                    Entity target = GetSingletonEntity<PayloadTargetInHierachy>();

                    foreach (Entity source in sourceEntities.Entities)
                    {
                        Transform.SetParent(World, source, target);
                    }
                }

                bool clear = false;
                unsafe
                {
                    ImGuiPayloadPtr payload = ImGui.GetDragDropPayload();
                    if (payload.NativePtr == null)
                    {
                        clear = true;
                    }
                }

                if (clear)
                {
                    foreach (Entity source in targetEntities.Entities)
                    {
                        Remove<PayloadTargetInHierachy>(source);
                    }

                    foreach (Entity source in sourceEntities.Entities)
                    {
                        Remove<PayloadSourceInHierachy>(source);
                    }
                }
            }
        }
    }
}