From 8e9e0fa3469da66781c2aebde8a5ea7f18a01a91 Mon Sep 17 00:00:00 2001
From: Shadowfacts <me@shadowfacts.net>
Date: Sat, 25 Feb 2023 17:59:35 -0500
Subject: [PATCH] Persist state when switching accounts

---
 Tusker/Scenes/MainSceneDelegate.swift         |  4 ++--
 ...ountSwitchingContainerViewController.swift | 21 +++++++++++++++++--
 Tusker/Shortcuts/UserActivityManager.swift    |  9 --------
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/Tusker/Scenes/MainSceneDelegate.swift b/Tusker/Scenes/MainSceneDelegate.swift
index 395e107ccf..7ca55ae0a1 100644
--- a/Tusker/Scenes/MainSceneDelegate.swift
+++ b/Tusker/Scenes/MainSceneDelegate.swift
@@ -218,9 +218,9 @@ class MainSceneDelegate: UIResponder, UIWindowSceneDelegate, TuskerSceneDelegate
             } else {
                 direction = .none
             }
-            container.setRoot(newRoot, animating: direction)
+            container.setRoot(newRoot, for: account, animating: direction)
         } else {
-            window!.rootViewController = AccountSwitchingContainerViewController(root: newRoot)
+            window!.rootViewController = AccountSwitchingContainerViewController(root: newRoot, for: account)
         }
     }
     
diff --git a/Tusker/Screens/Main/AccountSwitchingContainerViewController.swift b/Tusker/Screens/Main/AccountSwitchingContainerViewController.swift
index d1fd769277..eb35c7ea51 100644
--- a/Tusker/Screens/Main/AccountSwitchingContainerViewController.swift
+++ b/Tusker/Screens/Main/AccountSwitchingContainerViewController.swift
@@ -11,9 +11,13 @@ import ScreenCorners
 
 class AccountSwitchingContainerViewController: UIViewController {
 
+    private var currentAccountID: String
     private(set) var root: TuskerRootViewController
     
-    init(root: TuskerRootViewController) {
+    private var userActivities: [String: NSUserActivity] = [:]
+    
+    init(root: TuskerRootViewController, for account: LocalData.UserAccountInfo) {
+        self.currentAccountID = account.id
         self.root = root
         
         super.init(nibName: nil, bundle: nil)
@@ -29,14 +33,27 @@ class AccountSwitchingContainerViewController: UIViewController {
         embedChild(root)
     }
     
-    func setRoot(_ newRoot: TuskerRootViewController, animating direction: AnimationDirection) {
+    func setRoot(_ newRoot: TuskerRootViewController, for account: LocalData.UserAccountInfo, animating direction: AnimationDirection) {
         let oldRoot = self.root
         if direction == .none {
             oldRoot.removeViewAndController()
         }
+        if let activity = oldRoot.stateRestorationActivity() {
+            stateRestorationLogger.debug("AccountSwitchingContainer: saving \(activity.activityType, privacy: .public) for \(self.currentAccountID, privacy: .public)")
+            userActivities[currentAccountID] = activity
+        }
+        
+        self.currentAccountID = account.id
         self.root = newRoot
         embedChild(newRoot)
         
+        if let activity = userActivities.removeValue(forKey: account.id) {
+            stateRestorationLogger.debug("AccountSwitchingContainer: restoring \(activity.activityType, privacy: .public) for \(account.id, privacy: .public)")
+            let context = StateRestorationUserActivityHandlingContext(root: newRoot)
+            _ = activity.handleResume(manager: UserActivityManager(scene: view.window!.windowScene!, context: context))
+            context.finalize(activity: activity)
+        }
+        
         if direction != .none {
             if UIAccessibility.prefersCrossFadeTransitions {
                 newRoot.view.alpha = 0
diff --git a/Tusker/Shortcuts/UserActivityManager.swift b/Tusker/Shortcuts/UserActivityManager.swift
index 3e3d97fa41..48fd88f13c 100644
--- a/Tusker/Shortcuts/UserActivityManager.swift
+++ b/Tusker/Shortcuts/UserActivityManager.swift
@@ -32,15 +32,6 @@ class UserActivityManager {
         scene.session.mastodonController!
     }
     
-    private func getMainViewController() -> TuskerRootViewController {
-        let window = scene.windows.first { $0.isKeyWindow } ?? scene.windows.first!
-        return window.rootViewController as! TuskerRootViewController
-    }
-    
-    private func present(_ vc: UIViewController, animated: Bool = true) {
-        getMainViewController().present(vc, animated: animated)
-    }
-    
     static func getAccount(from activity: NSUserActivity) -> LocalData.UserAccountInfo? {
         guard let id = activity.userInfo?["accountID"] as? String else {
             return nil