From 4d8af832129e6b14b2d0e4d54ab75e675780507c Mon Sep 17 00:00:00 2001
From: Martin Zikmund <martin@zikmund.dev>
Date: Fri, 3 Sep 2021 08:58:20 +0200
Subject: [PATCH] fix: Flyout should return focus to previously focused

---
 .../UI/Xaml/Controls/Popup/PopupBase.cs       | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/Uno.UI/UI/Xaml/Controls/Popup/PopupBase.cs b/src/Uno.UI/UI/Xaml/Controls/Popup/PopupBase.cs
index d6f91353829b..720f24dd4767 100644
--- a/src/Uno.UI/UI/Xaml/Controls/Popup/PopupBase.cs
+++ b/src/Uno.UI/UI/Xaml/Controls/Popup/PopupBase.cs
@@ -1,5 +1,6 @@
 using System;
 using Uno.UI;
+using Uno.UI.Xaml.Core;
 using Windows.Foundation;
 using Windows.UI.Xaml.Input;
 using Windows.UI.Xaml.Media;
@@ -15,6 +16,9 @@ namespace Windows.UI.Xaml.Controls
 {
 	public partial class PopupBase : FrameworkElement, IPopup
 	{
+		private WeakReference<UIElement> _lastFocusedElement = null;
+		private FocusState _lastFocusState = FocusState.Unfocused;
+
 		private IDisposable _openPopupRegistration;
 		private bool _childHasOwnDataContext;
 
@@ -62,11 +66,32 @@ partial void OnIsOpenChangedPartial(bool oldIsOpen, bool newIsOpen)
 			if (newIsOpen)
 			{
 				_openPopupRegistration = VisualTreeHelper.RegisterOpenPopup(this);
+
+				if (IsLightDismissEnabled)
+				{
+					// Store last focused element
+					var focusManager = VisualTree.GetFocusManagerForElement(this);
+					var focusedElement = focusManager.FocusedElement as UIElement;
+					var focusState = focusManager.GetRealFocusStateForFocusedElement();
+					if (focusedElement != null && focusState != FocusState.Unfocused)
+					{
+						_lastFocusedElement = new WeakReference<UIElement>(focusedElement);
+						_lastFocusState = focusState;
+					}
+				}
+
 				Opened?.Invoke(this, newIsOpen);
 			}
 			else
 			{
 				_openPopupRegistration?.Dispose();
+				
+				if (_lastFocusedElement != null && _lastFocusedElement.TryGetTarget(out var target))
+				{
+					target.Focus(_lastFocusState);
+					_lastFocusedElement = null;
+				}
+
 				Closed?.Invoke(this, newIsOpen);
 			}
 		}