From 1b860e1c3ec7088c358d09bc28a2ae6dad93f26b Mon Sep 17 00:00:00 2001
From: k1ngsterr1 <swerss564@gmail.com>
Date: Fri, 9 Jan 2026 23:32:57 +0500
Subject: [PATCH] feat: Improve logout handling and back navigation in AuthFlow
 for better user experience

---
 .../com/rosetta/messenger/MainActivity.kt     |  6 +++--
 .../com/rosetta/messenger/ui/auth/AuthFlow.kt | 25 +++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/app/src/main/java/com/rosetta/messenger/MainActivity.kt b/app/src/main/java/com/rosetta/messenger/MainActivity.kt
index 9d4f43c..d390b08 100644
--- a/app/src/main/java/com/rosetta/messenger/MainActivity.kt
+++ b/app/src/main/java/com/rosetta/messenger/MainActivity.kt
@@ -164,10 +164,11 @@ class MainActivity : ComponentActivity() {
                                         }
                                     },
                                     onLogout = {
+                                        // Set currentAccount to null immediately to prevent UI lag
+                                        currentAccount = null
                                         scope.launch {
                                             com.rosetta.messenger.network.ProtocolManager.disconnect()
                                             accountManager.logout()
-                                            currentAccount = null
                                         }
                                     }
                                 )
@@ -182,10 +183,11 @@ class MainActivity : ComponentActivity() {
                                         }
                                     },
                                     onLogout = {
+                                        // Set currentAccount to null immediately to prevent UI lag
+                                        currentAccount = null
                                         scope.launch {
                                             com.rosetta.messenger.network.ProtocolManager.disconnect()
                                             accountManager.logout()
-                                            currentAccount = null
                                         }
                                     }
                                 )
diff --git a/app/src/main/java/com/rosetta/messenger/ui/auth/AuthFlow.kt b/app/src/main/java/com/rosetta/messenger/ui/auth/AuthFlow.kt
index 0f17ccd..f17112d 100644
--- a/app/src/main/java/com/rosetta/messenger/ui/auth/AuthFlow.kt
+++ b/app/src/main/java/com/rosetta/messenger/ui/auth/AuthFlow.kt
@@ -1,5 +1,6 @@
 package com.rosetta.messenger.ui.auth
 
+import androidx.activity.compose.BackHandler
 import androidx.compose.animation.*
 import androidx.compose.animation.core.tween
 import androidx.compose.runtime.*
@@ -35,6 +36,19 @@ fun AuthFlow(
     var selectedAccountId by remember { mutableStateOf<String?>(accounts.firstOrNull()?.id) }
     var showCreateModal by remember { mutableStateOf(false) }
     
+    // Handle system back button
+    BackHandler(enabled = currentScreen != AuthScreen.UNLOCK && currentScreen != AuthScreen.WELCOME) {
+        when (currentScreen) {
+            AuthScreen.SEED_PHRASE -> currentScreen = AuthScreen.WELCOME
+            AuthScreen.CONFIRM_SEED -> currentScreen = AuthScreen.SEED_PHRASE
+            AuthScreen.SET_PASSWORD -> {
+                currentScreen = if (hasExistingAccount) AuthScreen.UNLOCK else AuthScreen.CONFIRM_SEED
+            }
+            AuthScreen.IMPORT_SEED -> currentScreen = AuthScreen.WELCOME
+            else -> {} // UNLOCK and WELCOME don't have back navigation
+        }
+    }
+    
     AnimatedContent(
         targetState = currentScreen,
         transitionSpec = {
@@ -99,7 +113,14 @@ fun AuthFlow(
                 SetPasswordScreen(
                     seedPhrase = seedPhrase,
                     isDarkTheme = isDarkTheme,
-                    onBack = { currentScreen = AuthScreen.CONFIRM_SEED },
+                    onBack = { 
+                        // If there are existing accounts, allow going back to unlock screen
+                        if (hasExistingAccount) {
+                            currentScreen = AuthScreen.UNLOCK
+                        } else {
+                            currentScreen = AuthScreen.CONFIRM_SEED
+                        }
+                    },
                     onAccountCreated = { onAuthComplete(null) }
                 )
             }
@@ -122,8 +143,8 @@ fun AuthFlow(
                     onUnlocked = { account -> onAuthComplete(account) },
                     onSwitchAccount = {
                         // Navigate to create new account screen
+                        // Don't reset selectedAccountId so last logged account is remembered
                         currentScreen = AuthScreen.WELCOME
-                        selectedAccountId = null
                     }
                 )
             }