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.System
    {
        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;

        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);
        }

        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);
        }

        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())
                {
                    unsafe
                    {
                        fixed (Entity* payload = &entity)
                        {
                            ImGui.SetDragDropPayload($"{nameof(EditorHierarchyWindow)}", (IntPtr)payload, (uint)sizeof(Entity));
                        }
                    }

                    ImGui.EndDragDropSource();
                }

                if (ImGui.BeginDragDropTarget())
                {
                    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();
                }

                // Draw children.
                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);
                }
            }
        }
    }
}