From f6fc8f95a9a953b3d1f558de505eb5ed6eb18db9 Mon Sep 17 00:00:00 2001
From: hasareej <hasareej@amazon.com>
Date: Fri, 16 Apr 2021 15:28:05 +0100
Subject: [PATCH] Fixed random button highlights & toolbar order.

---
 .../EditorTransformComponentSelection.cpp     |  4 +--
 .../ViewportUi/ViewportUiSwitcher.cpp         | 30 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportSelection/EditorTransformComponentSelection.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportSelection/EditorTransformComponentSelection.cpp
index 3d4ed6ee0b..1f4a1d21be 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportSelection/EditorTransformComponentSelection.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportSelection/EditorTransformComponentSelection.cpp
@@ -1043,12 +1043,12 @@ namespace AzToolsFramework
         EditorEntityLockComponentNotificationBus::Router::BusRouterConnect();
         EditorManipulatorCommandUndoRedoRequestBus::Handler::BusConnect(entityContextId);
 
+        CreateTestSwitcher();
+
         CreateTransformModeSelectionCluster();
         RegisterActions();
         SetupBoxSelect();
         RefreshSelectedEntityIdsAndRegenerateManipulators();
-
-        CreateTestSwitcher();
     }
 
     EditorTransformComponentSelection::~EditorTransformComponentSelection()
diff --git a/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportUi/ViewportUiSwitcher.cpp b/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportUi/ViewportUiSwitcher.cpp
index 139d6cbb9f..5b2be40ec0 100644
--- a/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportUi/ViewportUiSwitcher.cpp
+++ b/Code/Framework/AzToolsFramework/AzToolsFramework/ViewportUi/ViewportUiSwitcher.cpp
@@ -49,6 +49,9 @@ namespace AzToolsFramework::ViewportUi::Internal
             QIcon buttonIcon = QString((button->m_icon).c_str());
 
             m_activeButton = new QToolButton();
+
+            // No hover effect for the main button as it's not clickable
+            m_activeButton->setProperty("IconHasHoverEffect", false);
             m_activeButton->setToolButtonStyle(Qt::ToolButtonTextBesideIcon);
             m_activeButton->setIcon(buttonIcon);
             m_activeButton->setText(buttonName);
@@ -85,11 +88,35 @@ namespace AzToolsFramework::ViewportUi::Internal
    
     void ViewportUiSwitcher::RemoveButton(ButtonId buttonId)
     {
-        ButtonId deleteLater = buttonId;
+        if (auto actionEntry = m_buttonActionMap.find(buttonId); actionEntry != m_buttonActionMap.end())
+        {
+            QAction* action = actionEntry->second;
+
+            // remove the action from the toolbar
+            removeAction(action);
+
+            // deregister from the widget manager
+            m_widgetCallbacks.RemoveWidget(action);
+
+            // resize to fit new area with minimum extra space
+            resize(minimumSizeHint());
+
+            m_buttonActionMap.erase(buttonId);
+
+            // reset current active mode if its the button being removed
+            if (buttonId == m_currentMode)
+            {
+                if (auto nextEntry = m_buttonActionMap.find(ButtonId(buttonId + 1)); nextEntry != m_buttonActionMap.end())
+                {
+                    SetActiveMode(nextEntry->first);
+                }
+            }
+        }
     }
     
     void ViewportUiSwitcher::Update()
     {
+        m_widgetCallbacks.Update();
     }
 
     void ViewportUiSwitcher::SetActiveMode(ButtonId buttonId)
@@ -117,6 +144,7 @@ namespace AzToolsFramework::ViewportUi::Internal
         {
             itr = m_buttonActionMap.find(m_currentMode);
             action = itr->second;
+
             addAction(action);
         }