From db91bdbf51318e82d753a349cc1097c1d7da92b2 Mon Sep 17 00:00:00 2001
From: Mikhail Naumov <mnaumov@amazon.com>
Date: Tue, 17 Aug 2021 18:25:22 -0700
Subject: [PATCH] Improvements to Entity Outliner context menu

Signed-off-by: Mikhail Naumov <mnaumov@amazon.com>
---
 .../SandboxIntegration.cpp                    | 29 +++++++++++++------
 .../UI/Prefab/PrefabIntegrationManager.cpp    | 22 +++++++++++---
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/Code/Editor/Plugins/ComponentEntityEditorPlugin/SandboxIntegration.cpp b/Code/Editor/Plugins/ComponentEntityEditorPlugin/SandboxIntegration.cpp
index b10fe35513..6b60f96089 100644
--- a/Code/Editor/Plugins/ComponentEntityEditorPlugin/SandboxIntegration.cpp
+++ b/Code/Editor/Plugins/ComponentEntityEditorPlugin/SandboxIntegration.cpp
@@ -646,16 +646,27 @@ void SandboxIntegrationManager::PopulateEditorGlobalContextMenu(QMenu* menu, con
 
     QAction* action = nullptr;
 
-    action = menu->addAction(QObject::tr("Create entity"));
-    QObject::connect(action, &QAction::triggered, action, [this] { ContextMenu_NewEntity(); });
-
-    if (selected.size() == 1)
+    // when nothing is selected, entity is created at root level
+    if (selected.size() == 0)
     {
-        action = menu->addAction(QObject::tr("Create child entity"));
-        QObject::connect(action, &QAction::triggered, action, [selected]
-        {
-            EBUS_EVENT(AzToolsFramework::EditorRequests::Bus, CreateNewEntityAsChild, selected.front());
-        });
+        action = menu->addAction(QObject::tr("Create entity"));
+        QObject::connect(
+            action, &QAction::triggered, action,
+            [this]
+            {
+                ContextMenu_NewEntity();
+            });
+    }
+    // when a single entity is selected, entity is created as its child
+    else if (selected.size() == 1)
+    {
+        action = menu->addAction(QObject::tr("Create entity"));
+        QObject::connect(
+            action, &QAction::triggered, action,
+            [selected]
+            {
+                EBUS_EVENT(AzToolsFramework::EditorRequests::Bus, CreateNewEntityAsChild, selected.front());
+            });
     }
 
     bool prefabSystemEnabled = false;
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabIntegrationManager.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabIntegrationManager.cpp
index eb9cf2b65f..7f3fbbeb38 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabIntegrationManager.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/UI/Prefab/PrefabIntegrationManager.cpp
@@ -11,6 +11,7 @@
 #include <AzCore/IO/FileIO.h>
 #include <AzCore/IO/SystemFile.h>
 #include <AzCore/StringFunc/StringFunc.h>
+#include <AzCore/Component/TransformBus.h>
 
 #include <AzFramework/API/ApplicationAPI.h>
 #include <AzFramework/Asset/AssetSystemBus.h>
@@ -363,12 +364,25 @@ namespace AzToolsFramework
 
             if (hasUserSelectedValidSourceFile)
             {
-                // Get position (center of viewport). If no viewport is available, (0,0,0) will be used.
-                AZ::Vector3 viewportCenterPosition = AZ::Vector3::CreateZero();
-                EditorRequestBus::BroadcastResult(viewportCenterPosition, &EditorRequestBus::Events::GetWorldPositionAtViewportCenter);
+                AZ::EntityId parentId;
+                AZ::Vector3 position = AZ::Vector3::CreateZero();
+
+                // if one entity is selected, set it as parent
+                EntityIdList selectedEntities;
+                ToolsApplicationRequestBus::BroadcastResult(selectedEntities, &ToolsApplicationRequests::GetSelectedEntities);
+                if (selectedEntities.size() == 1)
+                {
+                    parentId = selectedEntities.front();
+                    AZ::TransformBus::EventResult(position, parentId, &AZ::TransformInterface::GetWorldTranslation);
+                }
+                else
+                {
+                    // Get position (center of viewport). If no viewport is available, (0,0,0) will be used.
+                    EditorRequestBus::BroadcastResult(position, &EditorRequestBus::Events::GetWorldPositionAtViewportCenter);
+                }
 
                 // Instantiating from context menu always puts the instance at the root level
-                auto createPrefabOutcome = s_prefabPublicInterface->InstantiatePrefab(prefabFilePath, AZ::EntityId(), viewportCenterPosition);
+                auto createPrefabOutcome = s_prefabPublicInterface->InstantiatePrefab(prefabFilePath, parentId, position);
 
                 if (!createPrefabOutcome.IsSuccess())
                 {